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 Warp's high-level workflows"（见 warp/_src/utils.py）。它位于内核语言（warp.lang）与底层运行时（warp._src.context）之间：内部实现直接调用 native 运行时函数，对外则提供类型安全的 Python 接口。

从仓库的公共导出层 warp/utils.py 可以看到，该模块按四个分类组织并重新导出（re-export）实现：

分类	导出符号	底层实现位置
Array Operations	array_cast、array_inner、array_scan、array_sum	warp/_src/utils.py
Sorting	radix_sort_pairs、runlength_encode、segmented_sort_pairs	warp/_src/utils.py
Graph Coloring	GraphColoringAlgorithm、graph_coloring_assign、graph_coloring_balance、graph_coloring_get_groups	warp/_src/coloring.py
Allocators	AllocatorRmm	warp/_src/rmm_allocator.py
Misc	create_warp_function	warp/_src/utils.py

所有实现均通过from warp._src... import ... as ...形式重导出，因此wp.utils.array_scan与warp._src.utils.array_scan是同一个对象，文档生成的签名与运行时行为保持一致。以下各节按官方文档的分类顺序逐一展开。

数组操作（Array Operations）

前缀扫描array_scan

array_scan对数组执行扫描（前缀和）操作，将结果写入输出数组，支持包含式（inclusive）与排除式（exclusive）两种语义，向量类型按分量逐分量扫描。其核心签名如下（实现见 warp/_src/utils.py）：

def array_scan(in_array, out_array, inclusive=True) -> None

参数语义：

• in_array：输入数组，标量类型必须是int32、int64、float32、float64中的一种，可为标量类型或向量类型；
• out_array：输出数组，类型与大小必须与输入完全一致（types_equal校验）；
• inclusive：True为包含式扫描（当前元素计入和），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_host，CUDA 上对应*_device后缀；CUDA 路径还会检查返回状态码并上报错误字符串。扫描针对正步长（非连续）的 1D 数组同样可用（步长通过strides[0]传入 native 层）。

图捕获限制：在 CPU 图捕获或apic=True的 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], dtype=wp.int32) inc = wp.empty(4, dtype=wp.int32) exc = wp.empty(4, dtype=wp.int32) wp.utils.array_scan(values, inc, inclusive=True) # [1, 3, 6, 10] wp.utils.array_scan(values, exc, inclusive=False) # [0, 1, 3, 6]

数组求和array_sum

array_sum计算数组元素之和，支持全数组归约或沿指定轴归约，标量类型必须是float32或float64（实现见 warp/_src/utils.py）：

def array_sum(values, out=None, value_count=None, axis=None)

参数语义与返回值：

• values：输入数组，标量类型限定为float32/float64；
• out：可选输出数组。为None时自动创建；当axis is None且out is None时返回 Pythonfloat，否则返回out数组；
• value_count：参与计算的元素个数，None表示处理整个数组（axis is None时取values.size，否则取values.shape[axis]）；
• axis：归约轴，支持负数索引（从最后一个维度倒数），越界抛IndexError；None表示对全部元素求和。

行为细节：

• out若显式给出，其 device、dtype、shape 均需与计算预期一致，否则抛RuntimeError；
• value_count == 0时输出清零并直接返回（axis is None且 host 返回模式返回0.0）；
• CPU/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]], dtype=wp.float32) total = wp.utils.array_sum(values) # 10.0 col_sum = wp.utils.array_sum(values, axis=0) # array([4., 6.]) row_sum = wp.utils.array_sum(values, axis=-1) # array([3., 7.])

内积array_inner

array_inner计算两个同形状数组的内积（点积），可沿指定轴规约，要求两数组 shape、dtype 完全一致（实现见 warp/_src/utils.py）：

def array_inner(a, b, out=None, count=None, axis=None)

参数语义：

• a、b：两个输入数组，shape 不一致抛ValueError，device / dtype 不一致抛RuntimeError；
• count：参与计算的元素个数，None表示全部（axis is None时为a.size）；
• axis：归约轴，语义与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], dtype=wp.float32) b = wp.array([1.0, 2.0, 3.0], dtype=wp.float32) print(wp.utils.array_inner(a, b)) # 14.0

类型转换array_cast

array_cast将数组元素逐元素转换到另一种 dtype（实现见 warp/_src/utils.py）：

def array_cast(in_array, out_array, count=None)

行为细节：

• 输入输出数组必须在同一 device；
• 若两数组的维度数与 dtype 的"数据形状"（dtype._shape_）均匹配，则直接按元素转换；否则自动展平（flatten）并按标量级执行转换，即支持向量/矩阵 dtype 与标量 dtype 之间的互转；
• count指定处理的元素个数；对多维数组不支持部分转换（count < size时抛RuntimeError），1D 数组可按count截断；
• 若输入输出 dtype 相同，则退化为一次wp.copy，不做任何转换（源码 L989–L991）；
• 转换通过内部内核_array_cast_kernel（dest[i] = dest.dtype(src[i])）在目标设备上启动完成。

import warp as wp src = wp.array([1, 2, 3], dtype=wp.int32) dst = wp.empty(3, dtype=wp.float32) wp.utils.array_cast(src, dst) print(dst.numpy()) # [1. 2. 3.]

排序（Sorting）

键值对基数排序radix_sort_pairs

radix_sort_pairs基于基数排序对键-值对排序，稳定、近似线性时间复杂度，并且保持键值对应关系（实现见 warp/_src/utils.py）：

def radix_sort_pairs(keys, values, count, begin_bit=0, end_bit=None) -> None

参数语义：

• keys：键数组，dtype 支持int32、uint32、float32、int64、uint64、float64；
• values：值数组，元素必须为 4 或 8 字节宽（type_size_in_bytes校验，源码 L178–L182）；
• count：要排序的元素个数；
• begin_bit/end_bit：键位范围，end_bit=None时按键全宽（32 位键为 32，64 位键为 64）；要求0 <= begin_bit <= end_bit <= key_bit_width，且必须为整数，否则抛RuntimeError；begin_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], dtype=wp.int32) # 容量需 >= 2 * count，此处为 6 keys = wp.array([3, 1, 2, 0, 0, 0], dtype=wp.int32) values = wp.array([30, 10, 20, 0, 0, 0], dtype=wp.int32) wp.utils.radix_sort_pairs(keys, values, count=3) print(keys.numpy()[:3]) # [1 2 3] print(values.numpy()[:3]) # [10 20 30]

分段排序segmented_sort_pairs

segmented_sort_pairs在每个分段（segment）内部按键升序就地排序键值对，分段内稳定（相同键保持原始相对顺序），分段范围之外的元素不被修改（实现见 warp/_src/utils.py）：

def segmented_sort_pairs(keys, values, count, segment_start_indices, segment_end_indices=None) -> None

参数语义：

• keys：dtype 必须是int32或float32；
• values：dtype 必须是int32；
• count：参与分段的元素个数，必须是整数且满足0 <= count <= 2**31 - 1，布尔值会被拒绝（TypeError）；
• segment_start_indices：各分段的起始索引；当segment_end_indices为None时，相邻条目定义分段，因此长度为 N 的数组定义 N-1 个分段（源码 L371–L378 通过segment_start_indices[1:]推导）；
• segment_end_indices：可选的分段结束索引，提供时长度必须与起始索引数组一致。

约束与限制：

• 所有数组必须连续、位于同一 device；分段索引数组必须是 1Dint32；
• keys/values至少容纳2 * count个元素，后半段为可能被覆盖的 scratch 存储；
• 分段是半开区间[start, end)，必须满足0 <= start <= end <= count，分段之间不得重叠，分段索引数组也不得与keys/values存储重叠（Warp 当前不检测重叠，违反会导致未定义行为）；
• 错误报告：CPU 直接执行时非法分段边界抛ValueError；CPU 图回放时capture_launch抛RuntimeError；CUDA 设备上目前不报告非法边界，调用方需自行校验。

官方 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], dtype=wp.int32) values = wp.array([30, 10, 40, 20, 0, 0, 0, 0], dtype=wp.int32) offsets = wp.array([0, 2, 4], dtype=wp.int32) wp.utils.segmented_sort_pairs(keys, values, 4, offsets)

游程编码runlength_encode

runlength_encode对数组执行游程编码：将连续相同值压缩为"唯一值 + 游程长度"，例如[1,1,1,2,2,3]变为values=[1,2,3]、lengths=[3,2,1]（实现见 warp/_src/utils.py）：

def runlength_encode(values, run_values, run_lengths, run_count=None, value_count=None)

参数语义：

• values：输入数组，dtype 必须为int32；
• run_values：输出数组，存储唯一值，容量至少为value_count，dtype 必须与输入一致；
• run_lengths：输出数组，存储游程长度，dtype 必须为int32，容量至少为value_count；
• run_count：可选输出数组（int32），存储游程数量；为None时以整数形式在 host 返回结果；
• value_count：处理的元素个数，None表示处理整个数组，负值抛RuntimeError。

返回值：run_count is None时返回游程数量的 Python 整数；否则返回run_count数组（value_count == 0时对run_count清零后返回）。

APIC 捕获限制：CPU 图捕获或apic=True的 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"。

算法枚举GraphColoringAlgorithm

GraphColoringAlgorithm是IntEnum（见 warp/_src/coloring.py），提供两种着色算法：

枚举值	数值	说明
GraphColoringAlgorithm.MCS	0	基于最大基数搜索（Maximum Cardinality Search）的着色算法，通常能产生更少的颜色数
GraphColoringAlgorithm.GREEDY	1	按度数排序的贪心着色算法（degree-ordered greedy）

着色graph_coloring_assign

def graph_coloring_assign(edges, node_colors, algorithm=GraphColoringAlgorithm.MCS) -> int

• edges：形状为(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.py）：

def graph_coloring_balance(edges, node_colors, color_count, target_max_min_ratio) -> float

• color_count：当前着色使用的颜色数（即graph_coloring_assign的返回值）；
• target_max_min_ratio：期望的最大组/最小组规模比，算法在达到该比例或无法再改进时停止；
• 返回值：平衡后实际达到的 max/min 比例（图结构可能阻止进一步平衡，实际值可能高于目标值）；node_colors被就地修改。

获取颜色分组graph_coloring_get_groups

def graph_coloring_get_groups(node_colors, color_count, return_wp_array=True, device="cpu")

把node_colors转成按颜色划分的元组，每个元素是"拥有该颜色的节点 ID 数组"（实现见 warp/_src/coloring.py）：

• return_wp_array=True（默认）时返回 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]], dtype=wp.int32, device="cpu") colors = wp.empty(4, dtype=wp.int32, device="cpu") 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_array=True) # 之后即可按颜色顺序（同色节点互不相邻）并行处理各组节点

测试验证依据：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 即合法。

分配器（Allocators）

RMM 分配器AllocatorRmm

AllocatorRmm让 Warp 的设备内存分配走 RAPIDS Memory Manager（RMM），从而复用 RMM 的内存池与流序分配能力（实现见 warp/_src/rmm_allocator.py）：

import rmm import warp as wp rmm.reinitialize(pool_allocator=True, initial_pool_size=2**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_allocator（warp/_src/context.py）会校验自定义分配器是否实现了allocate/deallocate接口。

其他工具（Misc）

运行时创建 Warp 函数create_warp_function

create_warp_function把普通 Python 函数转换为 Warp 函数对象（实现见 warp/_src/utils.py）：

def create_warp_function(func) -> (wp.Function, Module)

• 返回(wp.Function, warp._src.context.Module)二元组；
• 函数命名：普通函数用__qualname__清洗后作为 key；lambda 会通过Adjoint.extract_lambda_source提取函数体源码并生成基于 SHA-256 的唯一名；无名函数则从源码提取并哈希命名（源码 L1008–L1027）；
• 创建的函数通过get_module(f"map_{key}")注册到独立模块，参数注解初始为Any（可重载）。

create_warp_function是wp.map的实现基石——wp.map会把 Python 函数 / lambda 经它转为wp.Function后再生成映射内核（见 warp/_src/utils.py）。wp.map本身支持多数组广播（遵循 NumPy 广播规则）、out输出复用、return_kernel=True仅返回内核、以及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], dtype=wp.float32) b = wp.array([3.0, 4.0], dtype=wp.float32) result = wp.map(f, a, b) # [4. 6.]

使用建议与限制总结

1. 设备与 dtype 是硬约束：上述所有工具都有明确的 dtype / device / 连续性校验，失败一律抛RuntimeError（或文档中注明的ValueError/TypeError/IndexError），调用前建议用array.numpy()或.device/.dtype属性预检；
2. 图捕获（APIC）场景提前规划：array_scan、array_sum、array_inner、runlength_encode、radix_sort_pairs、segmented_sort_pairs在 CPU 图捕获或apic=True的 CUDA 捕获下均要求显式提供输出数组（out/run_count），并禁止负步长；需要回放的结果必须在捕获区内显式分配存储；
3. 排序需要 2 倍容量：radix_sort_pairs与segmented_sort_pairs的 keys/values 数组容量必须 ≥2 * count，后半段是工作区；
4. 图着色管线是 CPU 工具：graph_coloring_assign/graph_coloring_balance/graph_coloring_get_groups均要求 CPU 上的int32数组，着色结果（颜色分组）可用于后续 GPU 内核的并行调度；平衡步骤建议配合graph_coloring_get_groups观察各组规模；
5. RMM 分配器按需引入：仅在需要 RMM 内存池、与 RAPIDS 生态共享显存时才安装rmm-cu12并启用，注意其非线程安全与"分配时生效"的语义。

参考源码与测试导航

• API 参考页：docs/api_reference/warp_utils.rst
• 公共导出层：warp/utils.py
• 数组操作与排序实现：warp/_src/utils.py
• 图着色实现：warp/_src/coloring.py
• RMM 分配器实现：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