Warp 工具函数库全解析：数组规约、排序、图着色与自定义分配器（warp.utils）

Warp 工具函数库全解析数组规约、排序、图着色与自定义分配器warp.utils【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp本文以 NVIDIA Warp 官方 API 参考页 warp_utils.rst 为骨架系统讲解warp.utils模块提供的全部高阶工具数组扫描/求和/内积/类型转换、基数排序与分段排序、游程编码、并行约束求解的图着色管线以及基于 RAPIDS RMM 的自定义 CUDA 分配器。读完本文你将掌握这些工具的参数语义、设备与数据类型约束、在图捕获APIC场景下的使用限制并能直接在自己的仿真或机器学习工作流中落地调用。模块定位与导出结构warp.utils是 Warp 面向高层工作流high-level workflows提供的工具集合模块 docstring 明确写道Utilities supporting Warps high-level workflows见 warp/_src/utils.py。它位于内核语言warp.lang与底层运行时warp._src.context之间内部实现直接调用 native 运行时函数对外则提供类型安全的 Python 接口。从仓库的公共导出层 warp/utils.py 可以看到该模块按四个分类组织并重新导出re-export实现分类导出符号底层实现位置Array Operationsarray_cast、array_inner、array_scan、array_sumwarp/_src/utils.pySortingradix_sort_pairs、runlength_encode、segmented_sort_pairswarp/_src/utils.pyGraph ColoringGraphColoringAlgorithm、graph_coloring_assign、graph_coloring_balance、graph_coloring_get_groupswarp/_src/coloring.pyAllocatorsAllocatorRmmwarp/_src/rmm_allocator.pyMisccreate_warp_functionwarp/_src/utils.py所有实现均通过from warp._src... import ... as ...形式重导出因此wp.utils.array_scan与warp._src.utils.array_scan是同一个对象文档生成的签名与运行时行为保持一致。以下各节按官方文档的分类顺序逐一展开。数组操作Array Operations前缀扫描array_scanarray_scan对数组执行扫描前缀和操作将结果写入输出数组支持包含式inclusive与排除式exclusive两种语义向量类型按分量逐分量扫描。其核心签名如下实现见 warp/_src/utils.pydef array_scan(in_array, out_array, inclusiveTrue) - None参数语义in_array输入数组标量类型必须是int32、int64、float32、float64中的一种可为标量类型或向量类型out_array输出数组类型与大小必须与输入完全一致types_equal校验inclusiveTrue为包含式扫描当前元素计入和False为排除式扫描当前元素不计入。前置校验源码 L76–L97 可见输入输出设备不匹配、大小不一致、dtype 不一致都会直接抛出RuntimeError空数组size 0直接返回输入输出若为非连续数组则必须是 1 维只接受标量或向量类型其余类型抛RuntimeError。底层调用链源码 L119–L140 根据设备CPU / CUDA与标量类型分派到 native 运行时符号例如 CPU 上为wp_array_scan_int_host/wp_array_scan_int64_host/wp_array_scan_float_host/wp_array_scan_double_hostCUDA 上对应*_device后缀CUDA 路径还会检查返回状态码并上报错误字符串。扫描针对正步长非连续的 1D 数组同样可用步长通过strides[0]传入 native 层。图捕获限制在 CPU 图捕获或apicTrue的 CUDA 图捕获期间int32、float32、int64、float64标量与向量扫描会被记录进 APIC 操作流非空、负步长的数组会抛出NotImplementedError此类扫描必须放到捕获区域之外执行。测试覆盖见 warp/tests/test_utils.py 中的test_array_scan、test_array_scan_vector、test_array_scan_strided_views等用例含空数组与各类错误分支。import warp as wp values wp.array([1, 2, 3, 4], dtypewp.int32) inc wp.empty(4, dtypewp.int32) exc wp.empty(4, dtypewp.int32) wp.utils.array_scan(values, inc, inclusiveTrue) # [1, 3, 6, 10] wp.utils.array_scan(values, exc, inclusiveFalse) # [0, 1, 3, 6]数组求和array_sumarray_sum计算数组元素之和支持全数组归约或沿指定轴归约标量类型必须是float32或float64实现见 warp/_src/utils.pydef array_sum(values, outNone, value_countNone, axisNone)参数语义与返回值values输入数组标量类型限定为float32/float64out可选输出数组。为None时自动创建当axis is None且out is None时返回 Pythonfloat否则返回out数组value_count参与计算的元素个数None表示处理整个数组axis is None时取values.size否则取values.shape[axis]axis归约轴支持负数索引从最后一个维度倒数越界抛IndexErrorNone表示对全部元素求和。行为细节out若显式给出其 device、dtype、shape 均需与计算预期一致否则抛RuntimeErrorvalue_count 0时输出清零并直接返回axis is None且 host 返回模式返回0.0CPU/CUDA 下分别调用wp_array_sum_float_host/wp_array_sum_double_host*_device等 native 符号源码 L718–L731。APIC 捕获限制非空调用在捕获期间必须显式提供out数组否则抛NotImplementedError负步长负 stride的输入或输出在捕获期间同样不支持计数与归约步长必须能放进有符号 32 位整数参与计算的地址与步长需按标量类型对齐。完整的布局校验逻辑见 warp/_src/utils.py 的_validate_apic_array_reduction_layout。import warp as wp values wp.array([[1.0, 2.0], [3.0, 4.0]], dtypewp.float32) total wp.utils.array_sum(values) # 10.0 col_sum wp.utils.array_sum(values, axis0) # array([4., 6.]) row_sum wp.utils.array_sum(values, axis-1) # array([3., 7.])内积array_innerarray_inner计算两个同形状数组的内积点积可沿指定轴规约要求两数组 shape、dtype 完全一致实现见 warp/_src/utils.pydef array_inner(a, b, outNone, countNone, axisNone)参数语义a、b两个输入数组shape 不一致抛ValueErrordevice / dtype 不一致抛RuntimeErrorcount参与计算的元素个数None表示全部axis is None时为a.sizeaxis归约轴语义与array_sum相同out可选输出数组。注意此处输出数组的 dtype 必须是标量类型scalar_type即与输入的标量成分类型一致见源码 L847而array_sum的输出 dtype 与输入保持一致axis is None且out is None时返回 Pythonfloat。底层在 CPU/CUDA 上分派wp_array_inner_float_host/wp_array_inner_double_host*_device按轴归约时逐输出位置循环调用 native 函数。APIC 捕获的约束与array_sum一致显式out、无负步长、计数与步长对齐。官方测试test_array_inner验证了a[1,2,3]、b[1,2,3]时结果为14.0见 warp/tests/test_utils.py。import warp as wp a wp.array([1.0, 2.0, 3.0], dtypewp.float32) b wp.array([1.0, 2.0, 3.0], dtypewp.float32) print(wp.utils.array_inner(a, b)) # 14.0类型转换array_castarray_cast将数组元素逐元素转换到另一种 dtype实现见 warp/_src/utils.pydef array_cast(in_array, out_array, countNone)行为细节输入输出数组必须在同一 device若两数组的维度数与 dtype 的数据形状dtype._shape_均匹配则直接按元素转换否则自动展平flatten并按标量级执行转换即支持向量/矩阵 dtype 与标量 dtype 之间的互转count指定处理的元素个数对多维数组不支持部分转换count size时抛RuntimeError1D 数组可按count截断若输入输出 dtype 相同则退化为一次wp.copy不做任何转换源码 L989–L991转换通过内部内核_array_cast_kerneldest[i] dest.dtype(src[i])在目标设备上启动完成。import warp as wp src wp.array([1, 2, 3], dtypewp.int32) dst wp.empty(3, dtypewp.float32) wp.utils.array_cast(src, dst) print(dst.numpy()) # [1. 2. 3.]排序Sorting键值对基数排序radix_sort_pairsradix_sort_pairs基于基数排序对键-值对排序稳定、近似线性时间复杂度并且保持键值对应关系实现见 warp/_src/utils.pydef radix_sort_pairs(keys, values, count, begin_bit0, end_bitNone) - None参数语义keys键数组dtype 支持int32、uint32、float32、int64、uint64、float64values值数组元素必须为 4 或 8 字节宽type_size_in_bytes校验源码 L178–L182count要排序的元素个数begin_bit/end_bit键位范围end_bitNone时按键全宽32 位键为 3264 位键为 64要求0 begin_bit end_bit key_bit_width且必须为整数否则抛RuntimeErrorbegin_bit end_bit时直接返回无事可做。存储约束keys与values容量必须至少容纳2 * count个元素排序需要临时工作区且两者都必须是连续contiguous数组否则抛RuntimeError。CUDA 端排序由*_device后缀的 native 内核执行CPU 端为*_host。APIC 捕获会将两个数组的 base region 记录进字节流。import warp as wp keys wp.array([3, 1, 2], dtypewp.int32) # 容量需 2 * count此处为 6 keys wp.array([3, 1, 2, 0, 0, 0], dtypewp.int32) values wp.array([30, 10, 20, 0, 0, 0], dtypewp.int32) wp.utils.radix_sort_pairs(keys, values, count3) print(keys.numpy()[:3]) # [1 2 3] print(values.numpy()[:3]) # [10 20 30]分段排序segmented_sort_pairssegmented_sort_pairs在每个分段segment内部按键升序就地排序键值对分段内稳定相同键保持原始相对顺序分段范围之外的元素不被修改实现见 warp/_src/utils.pydef segmented_sort_pairs(keys, values, count, segment_start_indices, segment_end_indicesNone) - None参数语义keysdtype 必须是int32或float32valuesdtype 必须是int32count参与分段的元素个数必须是整数且满足0 count 2**31 - 1布尔值会被拒绝TypeErrorsegment_start_indices各分段的起始索引当segment_end_indices为None时相邻条目定义分段因此长度为 N 的数组定义 N-1 个分段源码 L371–L378 通过segment_start_indices[1:]推导segment_end_indices可选的分段结束索引提供时长度必须与起始索引数组一致。约束与限制所有数组必须连续、位于同一 device分段索引数组必须是 1Dint32keys/values至少容纳2 * count个元素后半段为可能被覆盖的 scratch 存储分段是半开区间[start, end)必须满足0 start end count分段之间不得重叠分段索引数组也不得与keys/values存储重叠Warp 当前不检测重叠违反会导致未定义行为错误报告CPU 直接执行时非法分段边界抛ValueErrorCPU 图回放时capture_launch抛RuntimeErrorCUDA 设备上目前不报告非法边界调用方需自行校验。官方 docstring 内置示例源码 L318–L326对keys[3,1,4,2,0,0,0,0]、values[30,10,40,20,0,0,0,0]、offsets[0,2,4]排序后前 4 个键变为[1,3,2,4]对应值变为[10,30,20,40]。import warp as wp keys wp.array([3, 1, 4, 2, 0, 0, 0, 0], dtypewp.int32) values wp.array([30, 10, 40, 20, 0, 0, 0, 0], dtypewp.int32) offsets wp.array([0, 2, 4], dtypewp.int32) wp.utils.segmented_sort_pairs(keys, values, 4, offsets)游程编码runlength_encoderunlength_encode对数组执行游程编码将连续相同值压缩为唯一值 游程长度例如[1,1,1,2,2,3]变为values[1,2,3]、lengths[3,2,1]实现见 warp/_src/utils.pydef runlength_encode(values, run_values, run_lengths, run_countNone, value_countNone)参数语义values输入数组dtype 必须为int32run_values输出数组存储唯一值容量至少为value_countdtype 必须与输入一致run_lengths输出数组存储游程长度dtype 必须为int32容量至少为value_countrun_count可选输出数组int32存储游程数量为None时以整数形式在 host 返回结果value_count处理的元素个数None表示处理整个数组负值抛RuntimeError。返回值run_count is None时返回游程数量的 Python 整数否则返回run_count数组value_count 0时对run_count清零后返回。APIC 捕获限制CPU 图捕获或apicTrue的 CUDA 图捕获期间非空调用必须显式提供run_count数组因为 host 返回形式无法表示回放时的结果源码 L523–L527 抛NotImplementedError。底层在 CPU/CUDA 上调用wp_runlength_encode_int_host/wp_runlength_encode_int_device。专项测试见 warp/tests/test_runlength_encode.py。图着色Graph Coloring图着色工具面向并行约束求解场景给无向图的每个节点分配颜色保证相邻节点颜色不同从而让同色节点可以安全地并行处理。整个管线的三个步骤在 warp/_src/coloring.py 中实现模块 docstring 明确其用途为Graph coloring utilities for parallel constraint solving。算法枚举GraphColoringAlgorithmGraphColoringAlgorithm是IntEnum见 warp/_src/coloring.py提供两种着色算法枚举值数值说明GraphColoringAlgorithm.MCS0基于最大基数搜索Maximum Cardinality Search的着色算法通常能产生更少的颜色数GraphColoringAlgorithm.GREEDY1按度数排序的贪心着色算法degree-ordered greedy着色graph_coloring_assigndef graph_coloring_assign(edges, node_colors, algorithmGraphColoringAlgorithm.MCS) - intedges形状为(edge_count, 2)的 2D 数组每行[i, j]表示一条无向边必须是 CPU 上的int32数组node_colors形状为(node_count,)的 1Dint32CPU 数组将被填充颜色结果其长度即图节点数返回值使用的颜色总数。源码 L56–L85 依次校验 device、dtype、维度、形状然后调用 native 的wp_graph_coloring返回负数表示失败并抛RuntimeError。空图node_count 0直接抛错。平衡graph_coloring_balance贪心/MCS 着色产生的各颜色组大小可能严重不均导致并行处理的负载不均衡。graph_coloring_balance在保持着色合法性的前提下调整节点所属颜色组让各组规模更均衡实现见 warp/_src/coloring.pydef graph_coloring_balance(edges, node_colors, color_count, target_max_min_ratio) - floatcolor_count当前着色使用的颜色数即graph_coloring_assign的返回值target_max_min_ratio期望的最大组/最小组规模比算法在达到该比例或无法再改进时停止返回值平衡后实际达到的 max/min 比例图结构可能阻止进一步平衡实际值可能高于目标值node_colors被就地修改。获取颜色分组graph_coloring_get_groupsdef graph_coloring_get_groups(node_colors, color_count, return_wp_arrayTrue, devicecpu)把node_colors转成按颜色划分的元组每个元素是拥有该颜色的节点 ID 数组实现见 warp/_src/coloring.pyreturn_wp_arrayTrue默认时返回 Warp 数组元组可指定device否则返回 NumPy 数组元组color_count 0时返回空元组负值抛RuntimeError内部流程源码 L246–L268先用count_color_group_sizes内核统计各组大小再经group_offsets前缀和确定各组偏移最后用fill_color_groups内核填充展平后的分组数组并切片返回。这两个内部内核由于存在对计数器数组的写竞争必须用dim(1,)单线程启动docstring 中有明确警告。完整管线示例与 warp/tests/test_coloring.py 中的用法一致import warp as wp edges wp.array([[0, 1], [1, 2], [2, 3]], dtypewp.int32, devicecpu) colors wp.empty(4, dtypewp.int32, devicecpu) color_count wp.utils.graph_coloring_assign( edges, colors, wp.utils.GraphColoringAlgorithm.MCS ) ratio wp.utils.graph_coloring_balance(edges, colors, color_count, 1.1) groups wp.utils.graph_coloring_get_groups(colors, color_count, return_wp_arrayTrue) # 之后即可按颜色顺序同色节点互不相邻并行处理各组节点测试验证依据warp/tests/test_coloring.py 中test_coloring_corner_case验证两个相连节点必须获得不同颜色需 2 色test_coloring_trimesh在 Stanford bunny 网格上验证 GREEDY 与 MCS 着色后不存在任何相邻同色的非法边且平衡后 max/min 比例不劣化test_combine_coloring验证三角形环 3恰好需要 3 色、正方形环 4二分图恰好需要 2 色。这些测试同时给出了着色正确性校验的参考实现统计colors[v1] colors[v2]的边数为 0 即合法。分配器AllocatorsRMM 分配器AllocatorRmmAllocatorRmm让 Warp 的设备内存分配走 RAPIDS Memory ManagerRMM从而复用 RMM 的内存池与流序分配能力实现见 warp/_src/rmm_allocator.pyimport rmm import warp as wp rmm.reinitialize(pool_allocatorTrue, initial_pool_size2**30) wp.set_cuda_allocator(wp.utils.AllocatorRmm()) # 之后所有 wp.array 的分配都会经过 RMM 池关键机制每次分配委托给rmm.DeviceBuffer使用rmm.mr.set_current_device_resource()设置的当前DeviceMemoryResource切换 RMM 资源会影响后续分配依赖rmm包仅 Linux 支持安装方式为pip install rmm-cu12需与 CUDA 版本匹配未安装时构造器直接抛ImportError单个AllocatorRmm实例可安全地在多个 CUDA 设备间共享wp.array的分配路径会以device.context_guard包裹每次allocate()因此分配总发生在正确的设备上分配在 Warp 设备流的当前流上按流序stream-ordered执行与 CuPy 的 RMM 集成模式一致保证与流序内存资源如rmm.mr.CudaAsyncMemoryResource以及 CUDA 图捕获的兼容性该分配器不是线程安全的多线程并发调用需要外部同步deallocate对未识别指针会抛RuntimeError可帮助发现 double-free 或指针归属错误__repr__报告当前活跃 buffer 数量。接入方式全局设置使用wp.set_cuda_allocator(allocator)实现在 warp/_src/context.py也可用wp.set_device_allocator(device, allocator)按设备设置、wp.get_device_allocator(device)查询当前分配器_validate_allocatorwarp/_src/context.py会校验自定义分配器是否实现了allocate/deallocate接口。其他工具Misc运行时创建 Warp 函数create_warp_functioncreate_warp_function把普通 Python 函数转换为 Warp 函数对象实现见 warp/_src/utils.pydef create_warp_function(func) - (wp.Function, Module)返回(wp.Function, warp._src.context.Module)二元组函数命名普通函数用__qualname__清洗后作为 keylambda 会通过Adjoint.extract_lambda_source提取函数体源码并生成基于 SHA-256 的唯一名无名函数则从源码提取并哈希命名源码 L1008–L1027创建的函数通过get_module(fmap_{key})注册到独立模块参数注解初始为Any可重载。create_warp_function是wp.map的实现基石——wp.map会把 Python 函数 / lambda 经它转为wp.Function后再生成映射内核见 warp/_src/utils.py。wp.map本身支持多数组广播遵循 NumPy 广播规则、out输出复用、return_kernelTrue仅返回内核、以及block_dim线程块配置。import warp as wp f, module wp.utils.create_warp_function(lambda x, y: x y) a wp.array([1.0, 2.0], dtypewp.float32) b wp.array([3.0, 4.0], dtypewp.float32) result wp.map(f, a, b) # [4. 6.]使用建议与限制总结设备与 dtype 是硬约束上述所有工具都有明确的 dtype / device / 连续性校验失败一律抛RuntimeError或文档中注明的ValueError/TypeError/IndexError调用前建议用array.numpy()或.device/.dtype属性预检图捕获APIC场景提前规划array_scan、array_sum、array_inner、runlength_encode、radix_sort_pairs、segmented_sort_pairs在 CPU 图捕获或apicTrue的 CUDA 捕获下均要求显式提供输出数组out/run_count并禁止负步长需要回放的结果必须在捕获区内显式分配存储排序需要 2 倍容量radix_sort_pairs与segmented_sort_pairs的 keys/values 数组容量必须 ≥2 * count后半段是工作区图着色管线是 CPU 工具graph_coloring_assign/graph_coloring_balance/graph_coloring_get_groups均要求 CPU 上的int32数组着色结果颜色分组可用于后续 GPU 内核的并行调度平衡步骤建议配合graph_coloring_get_groups观察各组规模RMM 分配器按需引入仅在需要 RMM 内存池、与 RAPIDS 生态共享显存时才安装rmm-cu12并启用注意其非线程安全与分配时生效的语义。参考源码与测试导航API 参考页docs/api_reference/warp_utils.rst公共导出层warp/utils.py数组操作与排序实现warp/_src/utils.py图着色实现warp/_src/coloring.pyRMM 分配器实现warp/_src/rmm_allocator.py分配器注册接口warp/_src/context.py测试图着色 warp/tests/test_coloring.py、数组工具 warp/tests/test_utils.py、游程编码 warp/tests/test_runlength_encode.py、数组归约 warp/tests/test_array_reduce.py【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考