如何使用Warp在Python环境中编写CUDA内核

通常，实时物理模拟代码是用低级 CUDA C ++编写的，以获得最佳性能。在这篇文章中，我们将介绍 NVIDIA Warp ，这是一个新的 Python 框架，可以轻松地用 Python 编写可微图形和模拟 GPU 代码。 Warp 提供了编写高性能仿真代码所需的构建块，但它的工作效率与 Python 等解释语言相当。

在这篇文章的最后，您将学习如何使用 Warp 在 Python 环境中编写 CUDA 内核，并利用一些内置的高级功能，从而轻松编写复杂的物理模拟，例如海洋模拟。

安装

Warp 以 来自 GitHub 的开源库 的形式提供。克隆存储库后，可以使用本地软件包管理器进行安装。对于 pip ，请使用以下命令：

pip install warp

初始化

导入后，必须显式初始化扭曲：

import warp as wp
wp.init()

推出内核

Warp 使用 Python 装饰器的概念来标记可以在 GPU 上执行的函数。例如，可以编写一个简单的半隐式粒子积分方案，如下所示：

@wp.kernel
def integrate(x: wp.array(dtype=wp.vec3), v: wp.array(dtype=wp.vec3), f: wp.array(dtype=wp.vec3), w: wp.array(dtype=float), gravity: wp.vec3, dt: float): # thread id tid = wp.tid() x0 = x[tid] v0 = v[tid] # Semi-implicit Euler step f_ext = f[tid] inv_mass = w[tid] v1 = v0 + (f_ext * inv_mass + gravity) * dt x1 = x0 + v1 * dt # store results x[tid] = x1 v[tid] = v1 

因为 Warp 是强类型的，所以应该为内核参数提供类型提示。要启动内核，请使用以下语法：

 wp.launch(kernel=simple_kernel, # kernel to launch dim=1024, # number of threads inputs=[a, b, c], # parameters device="cuda") # execution device

与 NumPy 等基于张量的框架不同， Warp 使用 kernel-based 编程模型。基于内核的编程与底层 GPU 执行模型更为匹配。对于需要细粒度条件逻辑和内存操作的模拟代码，这通常是一种更自然的表达方式。然而， Warp 以一种易于使用的方式公开了这种以线程为中心的编程模型，它不需要 GPU 体系结构的低级知识。

编译模型

启动内核会触发实时（ JIT ）编译管道，该管道会自动从 Python 函数定义生成 C ++/ CUDA 内核代码。

属于 Python 模块的所有内核都在运行时编译到动态库和 PTX 中。图 2 。显示了编译管道，其中包括遍历函数 AST 并将其转换为直线 CUDA 代码，然后编译并加载回 Python 进程。

图 2 。 Warp 内核的编译管道

这个 JIT 编译的结果被缓存。如果输入内核源代码不变，那么预编译的二进制文件将以低开销的方式加载。

记忆模型

Warp 中的内存分配通过warp.array类型公开。阵列封装了可能位于主机（ CPU ）或设备（ GPU ）内存中的底层内存分配。与张量框架不同， Warp 中的数组是强类型的，并存储内置结构的线性序列（vec3， matrix33， quat，等等）。

您可以从 Python 列表或 NumPy 数组中构造数组，或使用与 NumPy 和 PyTorch 类似的语法进行初始化：

# allocate an uninitizalized array of vec3s v = wp.empty(length=n, dtype=wp.vec3, device="cuda") # allocate a zero-initialized array of quaternions q = wp.zeros(length=n, dtype=wp.quat, device="cuda") # allocate and initialize an array from a numpy array # will be automatically transferred to the specified device v = wp.from_numpy(array, dtype=wp.vec3, device="cuda")

Warp 支持__array_interface__和__cuda_array_interface__协议，允许在基于张量的框架之间进行零拷贝数据视图。例如，要将数据转换为 NumPy ，请使用以下命令：

# automatically bring data from device back to host view = device_array.numpy()

特征

Warp 包含几个更高级别的数据结构，使实现模拟和几何处理算法更容易。

网格

三角形网格在仿真和计算机图形学中无处不在。 Warp 提供了一种内置类型，用于管理网格数据，该数据支持几何查询，例如最近点、光线投射和重叠检查。

下面的示例演示如何使用“扭曲”计算网格上距离输入位置数组最近的点。这种类型的计算是碰撞检测中许多算法的基础（图 3 ）。 Warp 的网格查询使实现此类方法变得简单。

@wp.kernel
def project(positions: wp.array(dtype=wp.vec3), mesh: wp.uint64, output_pos: wp.array(dtype=wp.vec3), output_face: wp.array(dtype=int)): tid = wp.tid() x = wp.load(positions, tid) face_index = int(0) face_u = float(0.0) face_v = float(0.0) sign = float(0.0) max_dist = 2.0 if (wp.mesh_query_point(mesh, x, max_dist, sign, face_index, face_u, face_v)): p = wp.mesh_eval_position(mesh, face_index, face_u, face_v) output_pos[tid] = p output_face[tid] = face_index

稀疏卷

稀疏体对于表示大型域上的网格数据非常有用，例如复杂对象的符号距离场（ SDF ）或大规模流体流动的速度。 Warp 支持使用 NanoVDB 标准定义的稀疏卷。使用标准 OpenVDB 工具（如 Blender 、 Houdini 或 Maya ）构造卷，然后在 Warp 内核内部采样。

您可以直接从磁盘或内存中的二进制网格文件创建卷，然后使用volumes API 对其进行采样：

wp.volume_sample_world(vol, xyz, mode) # world space sample using interpolation mode
wp.volume_sample_local(vol, uvw, mode) # volume space sample using interpolation mode
wp.volume_lookup(vol, ijk) # direct voxel lookup
wp.volume_transform(vol, xyz) # map point from voxel space to world space
wp.volume_transform_inv(vol, xyz) # map point from world space to volume space

使用卷查询，您可以以最小的内存开销高效地碰撞复杂对象。

散列网格

许多基于粒子的模拟方法，如离散元法（ DEM ）或平滑粒子流体动力学（ SPH ），都涉及到在空间邻域上迭代以计算力的相互作用。哈希网格是一种成熟的数据结构，用于加速这些最近邻查询，特别适合 GPU 。

哈希网格由点集构成，如下所示：

哈希网格由点集构成，如下所示：

grid = wp.HashGrid(dim_x=128, dim_y=128, dim_z=128, device="cuda")
grid.build(points=p, radius=r)

创建散列网格后，可以直接从用户内核代码中查询它们，如以下示例所示，该示例计算所有相邻粒子位置的总和：

@wp.kernel
def sum(grid : wp.uint64, points: wp.array(dtype=wp.vec3), output: wp.array(dtype=wp.vec3), radius: float): tid = wp.tid() # query point p = points[tid] # create grid query around point query = wp.hash_grid_query(grid, p, radius) index = int(0) sum = wp.vec3() while(wp.hash_grid_query_next(query, index)): neighbor = points[index] # compute distance to neighbor point dist = wp.length(p-neighbor) if (dist <= radius): sum += neighbor output[tid] = sum

图 5 显示了粘性材料的 DEM 颗粒材料模拟示例。使用内置的哈希网格数据结构，您可以在不到 200 行 Python 中编写这样的模拟，并以交互速率运行超过 100K 个粒子。

使用扭曲散列网格数据可以轻松评估相邻粒子之间的成对力相互作用。

可微性

基于张量的框架，如 PyTorch 和 JAX ，提供了张量计算的梯度，非常适合于 ML 训练等应用。

Warp 的一个独特功能是能够生成 kernel code 的正向和反向版本。这使得编写可微模拟变得很容易，可以将梯度作为更大训练管道的一部分进行传播。一个常见的场景是，对网络层使用传统的 ML 框架，并使用 Warp 实现允许端到端差异性的模拟层。

当需要渐变时，应使用requires_grad=True创建阵列。例如，warp.Tape类可以记录内核启动并回放它们，以计算标量损失函数相对于内核输入的梯度：

tape = wp.Tape() # forward pass
with tape: wp.launch(kernel=compute1, inputs=[a, b], device="cuda") wp.launch(kernel=compute2, inputs=[c, d], device="cuda") wp.launch(kernel=loss, inputs=[d, l], device="cuda") # reverse pass
tape.backward(loss=l)

完成后向传递后，可通过Tape对象中的映射获得与输入相关的梯度：

# gradient of loss with respect to input a
print(tape.gradients[a])

图 6 。一个轨迹优化的例子，其中球的初始速度被优化以击中黑色目标。每行显示 LBFGS 优化步骤的一次迭代的结果。

总结

在这篇文章中，我们介绍了 NVIDIA Warp ，这是一个 Python 框架，可以很容易地为 GPU 编写可微模拟代码。

关于作者

迈尔斯·麦克林（ Miles Macklin ）是NVIDIA 的首席工程师，致力于模拟技术。他从哥本哈根大学获得计算机科学博士学位，从事计算机图形学、基于物理学的动画和机器人学的研究。他在 ACM SIGGRAPH 期刊上发表了几篇论文，他的研究已经被整合到许多商业产品中，包括NVIDIA 的 PhysX 和 ISAAC 健身房模拟器。他最近的工作旨在为 GPU 上的可微编程开发健壮高效的框架。

Fred Oh 是 CUDA 、 CUDA on WSL 和 CUDA Python 的高级产品营销经理。弗雷德拥有加州大学戴维斯分校计算机科学和数学学士学位。他的职业生涯开始于一名 UNIX 软件工程师，负责将内核服务和设备驱动程序移植到 x86 体系结构。他喜欢《星球大战》、《星际迷航》和 NBA 勇士队。

审核编辑：郭婷