WARP: A HIGH-PERFORMANCE PYTHON FRAMEWORK FOR GPU SIMULATION AND GRAPHICS

MILES MACKLIN, PRINCIPAL ENGINEER, NVIDIA

目录

• Warp 背景
• Warp 目标
• Warp 非目标
• 为什么不使用张量框架
• 功能比较

• 编译模型

  • Warp 编译流水线
  • 代码生成

• 类型系统

• C/C++ 互操作

• 自动微分

  • 计算梯度 - 示例

• 模拟示例

• 模拟 (SIMULATION

  • 布料模拟 (CLOTH SIMULATION
  • 颗粒材料 (GRANULAR MATERIALS
  • 海洋模拟 (OCEAN SIMULATION
  • 数字人变形器 (DIGITAL HUMAN DEFORMER
  • OpenVDB 支持 (OPENVDB SUPPORT

• 基准测试 (BENCHMARKING

  • 基准测试 - 布料求解器 (BENCHMARK - CLOTH SOLVER
  • 初始化时间 (越低越好
  • 模拟时间 (越低越好
  • CUDA 图 (CUDA GRAPHS

• Omniverse

  • Omniverse 集成 (OMNIVERSE INTEGRATION
  • OmniGraph + Warp (OMNIGRAPH + WARP

• 迭代时间

• 总结
• 参与Omniverse开发者日
• 结束页

Warp 背景

Warp 最初是作为可微分物理模拟的编译器后端而编写的。
未来将有更广泛的关注点，包括：
* 几何处理：例如，OmniGraph 节点收缩包络。
* 物理模拟：例如，海洋形变器、粒子物理。
* 强化学习研究：例如，IsaacGym、奖励函数建模。
* 合成数据生成：例如，图像处理。
* 可微分模拟：例如，药物发现、机器人技术。

Warp 目标

1. 便于使用 Python 编写 GPU 图形和模拟代码：

import warp as wp

@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]

    f_ext = f[tid]

    inv_mass = w[tid]

    # Semi-implicit Euler step
    v1 = v0 + (f_ext * inv_mass + gravity) * dt
    x1 = x0 + v1 * dt

    x[tid] = x1
    v[tid] = v1

# Kernel launch
wp.launch(integrate, dim=1024, inputs=[x, v, f, w], device="cuda")

1. 降低门槛，加快迭代时间。
2. 丰富的内置数学和几何库。
3. 与现有 C++/CUDA 代码的简单互操作。
4. 完全可微分。

概括为“模拟的 Shadertoy”。

Warp 非目标

1. 不支持完整的 CPython 语言。
2. 不进行基于张量的编程（如 CuPy, PyTorch 等）。
3. 不与 HPC 库竞争（如 cuBLAS, cuSparse 等）。
4. 不引入新的语言。

为什么不使用张量框架？

考虑一个点到三角形最近点的函数，包含复杂的条件逻辑和多个提前退出。

@wp.func
def triangle_closest_point(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
    ab = b - a
    ac = c - a
    ap = p - a

    d1 = wp.dot(ab, ap)
    d2 = wp.dot(ac, ap)

    if (d1 <= 0.0 and d2 <= 0.0):
        return wp.vec3(1.0, 0.0, 0.0)

    bp = p - b
    d3 = wp.dot(ab, bp)
    d4 = wp.dot(ac, bp)

    if (d3 >= 0.0 and d4 <= d3):
        return wp.vec3(0.0, 1.0, 0.0)

    vc = d1 * d4 - d3 * d2
    v = d1 / (d1 - d3)
    if (vc <= 0.0 and d1 > 0.0 and d3 < 0.0):
        return wp.vec3(1.0 - v, v, 0.0)

    cp = p - c
    d5 = wp.dot(ab, cp)
    d6 = wp.dot(ac, cp)

    vc = d1 * d4 - d3 * d2
    v = d1 / (d1 - d3)
    if (vc <= 0.0 and d1 > 0.0 and d3 < 0.0):
        return wp.vec3(1.0 - v, v, 0.0)

    cp = p - c
    d5 = wp.dot(ab, cp)
    d6 = wp.dot(ac, cp)

    v = (d4 * d5 - d1 * d6) / ((d4 - d6) * (d5 - d1) - d5 * d6 + d3 * d1)
    if (d6 > 0.0 and d5 <= d6):
        return wp.vec3(0.0, 0.0, 1.0)

    vb = d5 * d2 - d1 * d6
    w = d2 / (d2 - d6)
    if (vb <= 0.0 and d2 > 0.0 and d6 < 0.0):
        return wp.vec3(1.0 - w, 0.0, w)

    va = d3 * d6 - d5 * d4
    w = (d3 - d5) / ((d4 - d6) * (d5 - d1) - d5 * d6 + d3 * d1)
    if (va <= 0.0 and d4 - d3 > 0.0 and d5 - d6 > 0.0):
        return wp.vec3(0.0, w, 1.0 - w)

    denom = 1.0 / (va + vb + vc)
    v = vc * denom
    w = vb * denom

    return wp.vec3(1.0 - v - w, v, w)

对于张量框架，处理这类问题存在以下困难：
* 需要手动将条件语句转换为条件选择/掩码。
* 内核融合不可靠且不透明。
* 对分散写入/原子累加的支持不佳。

功能比较

下表比较了 Warp 与其他几种框架的功能：

• Legate：性能大致等同于单个 GPU 上的 CuPy。
• Pythran：不支持 GPU。

编译模型

Warp 编译流水线

该流程包括：
1. Python 代码（带有 @wp.kernel 装饰器）通过 AST 遍历。
2. 生成 .cpp, .cu, .hlsl 文件。
3. 这些文件通过 MSVC/GCC、NVRTC、FXC 编译器编译。
4. 生成 .dll, .so, .ptx 文件。
5. 编译过程利用 Kernel Cache。
6. 最终由 Python 通过 wp.launch() 导入并执行。

代码生成

Python AST -> SSA -> CUDA。
状态存储在寄存器中。
内核融合是隐式的。
支持调试。

@wp.func
def simple(a : float, b : float):
    return wp.sin(a) + a*b

// generated code
float simple(
    float var_a,
    float var_b)
{
    // primal vars
    float var_0;
    float var_1;
    float var_2;

    // forward pass
    var_0 = wp::sin(var_a);
    var_1 = wp::mul(var_a, var_b);
    var_2 = wp::add(var_0, var_1);

    return var_2;
}

类型系统

• Warp 是强类型化的。
• 类型从数组元数据 + Python 类型提示中推断出来。
• 内置数学类型类似于 OpenCL/HLSL，包括：vec2, vec3, vec4, mat22, mat33, mat44, quat, transform。
• 内置类型表现为元组（不可变），不支持原地操作。
• 主机/设备内存通过 wp.array 类型管理，支持 __array_interface__, __cuda_array_interface__。
• 提供了将 PyTorch 转换为数组的辅助函数：wp.from_torch(), wp.to_torch()。
• 能够将梯度反向传播回 PyTorch。

C/C++ 互操作

• 原生调用现有 C++/CUDA 代码。
• 包含数学函数和伴随函数。
• 使用 wp.add_builtin() 注册原生函数。
• 将函数暴露给任何 Warp 内核。
• 传播梯度。

// forward mode
inline CUDA_CALLABLE vec3 cross(vec3 a, vec3 b)
{
    vec3 c;
    c.x = a.y*b.z - a.z*b.y;
    c.y = a.z*b.x - a.x*b.z;
    c.z = a.x*b.y - a.y*b.x;

    return c;
}

// reverse mode (optional)
inline CUDA_CALLABLE void adj_cross(
    vec3 a,
    vec3 b,
    vec3 adj_a,
    vec3 adj_b,
    vec3 adj_ret)
{
    adj_a = cross(b, adj_ret);
    adj_b = cross(a, adj_ret);
}

# register
add_builtin("cross", input_types={"x": vec3, "y": vec3}, value_type=vec3, ...)

自动微分

计算梯度 - 示例

• wp.Tape() 的使用示例。
• 使用 tape 记录内核启动。
• 运行 tape.backward()。
• 从输入到梯度的映射可在 tape.gradients{} 中获取。

tape = wp.Tape()
with tape:
    wp.launch(kernel=kernel_1, dim=dim, inputs=[a], outputs=[b], device=device)
    wp.launch(kernel=kernel_2, dim=dim, inputs=[b], outputs=[c], device=device)
    wp.launch(kernel=kernel_3, dim=dim, inputs=[c], outputs=[l], device=device)

# run backward
tape.backward(loss=l)

# gradients of loss w.r.t inputs
print(tape.gradients[a])

模拟示例

模拟 (SIMULATION)

布料模拟 (CLOTH SIMULATION)

布料模拟利用了以下技术：
* Bridson et al. 提出的弯曲模型。
* Neo-Hookean FEM 弹性模型。
* 内置的网格查询功能，包括 wp.mesh_query_point() 和 wp.mesh_query_ray()。
* 针对动态网格的 SDF 接触。
* 快速 GPU BVH 重建。
* 能够实现 60k 粒子对比 32k 三角形在 0.25ms 内完成模拟。

颗粒材料 (GRANULAR MATERIALS)

颗粒材料模拟的特点包括：
* 内置哈希网格原语，包括 wp.hash_grid_query_point() 和 wp.hash_grid_query_next()。
* 快速的最近邻查询。
* 使用 200 行 Python 代码实现了 DEM（离散元方法）。
* 可模拟 128k 粒子，每帧约 10ms。

进一步地，颗粒材料模拟还支持：
* 将 DEM 与形变网格碰撞结合。
* 来自 USD 的时间采样动画。
* 处理凝聚力（Cohesion）和身体粘附（adhesion to body）。
* 可模拟 512k 粒子，每帧约 40ms。

海洋模拟 (OCEAN SIMULATION)

海洋模拟由以下组件构成：
* 由 Stefan Jeschke 开发的 Ocean deformer (omni.ocean)。
* 基于水面小波的程序网格位移。
* 参考来源：[Jeschke 2018]。

数字人变形器 (DIGITAL HUMAN DEFORMER)

数字人变形器功能包括：
* 由 Edy Susanto Lim 开发的 Shrink-wrap deformer。
* 利用核射线投射（ray-cast）技术将眼睑投影到眼球上。
* 通过 OmniGraph 管道驱动。

OpenVDB 支持 (OPENVDB SUPPORT)

系统支持 OpenVDB，具体包括：
* 支持 NanoVDB 稀疏体。
* 非常适用于窄带 SDF 碰撞。
* 提供简单易用的 API，例如：
* wp.volume_sample_world(vol, xyz)
* wp.volume_sample_local(vol, uvw)
* wp.volume_lookup(vol, ijk)
* wp.volume_transform(vol, xyz)
* wp.volume_transform_inv(vol, xyz)

基准测试 (BENCHMARKING)

基准测试 - 布料求解器 (BENCHMARK - CLOTH SOLVER)

布料求解器的基准测试基于以下实现：

• 半隐式弹簧-质量布料模拟器。
• 包含分支/收集/散布 (branches/gather/scatter) 操作。
• 对图结构没有假设。
• 在 7 个框架中实现，共有 11 种配置，包括：
  • Warp (CPU/GPU)
  • Taichi (CPU/GPU)
  • PyTorch (CPU/GPU)
  • JAX (CPU/GPU)
  • CuPy
  • NumPy
  • Numba

初始化时间 (越低越好) (LOWER IS BETTER)

• 冷启动时间 (Cold time)：表示代码更改后首次运行并得到结果所需的时间。
• 后续的“热启动”时间可能因缓存而更快。
• 测试时间包括了 PTX JIT 编译时间（所有 GPU 方法通用）。

模拟时间 (越低越好) (LOWER IS BETTER)

• 张量框架的开销显著更大。
• 这主要归因于缺乏核融合 (kernel fusion)。
• 在此工作负载规模下，恒定的开销因素主导了 GPU 时间。

CUDA 图 (CUDA GRAPHS)

• 旨在尽可能避免 CPython 开销。
• 类型检查/类型转换在执行数千次核启动时会迅速增加开销。

• Warp 对 CUDA 图提供了原生支持，通过以下 API 实现：

  • wp.capture_begin()
  • wp.capture_end()
  • wp.capture_launch()

• 这带来了启动开销 > 10 倍的减少。

Omniverse

Omniverse 集成 (OMNIVERSE INTEGRATION)

Warp 与 Omniverse 的集成方式如下：
* 通过 omni.warp 扩展暴露功能。
* 将 Warp 列为扩展的依赖项。
* 通过 import warp 即可开始编写核函数 (kernels)。

OmniGraph + Warp (OMNIGRAPH + WARP)

将 OmniGraph 与 Warp 结合使用具有以下优势：
* 非常适合编写 Python OmniGraph 节点。
* 实现超快的迭代时间。
* 支持 Ctrl-S 热重载 (Hot Reload)，使得 JIT 编译在 < 1s 内完成。
* 扩展附带了示例，包括：
* 波动方程求解器
* 布料模拟
* 粒子模拟

迭代时间

在第31页，标题为“ITERATION TIMES”，展示了一个包含3D模拟软件界面和代码编辑器的截图。左侧的软件界面显示了一个波浪形表面上的球体，实时帧率为102 Hz。界面还显示了系统资源信息：NVIDIA GeForce RTX 3090 GPU使用了2 GB，可用20 GB；主机内存使用了17 GB，可用128 GB。右侧是Python代码，其中包含对cupy和numpy库的引用，并定义了一个CUDA内核，用于处理“wave_solver”逻辑。这页幻灯片强调了迭代过程的效率和速度，展示了在NVIDIA Omniverse或类似环境中进行实时物理模拟及其背后的Python/CUDA代码。

总结

第32页的标题为“SUMMARY”，总结了之前介绍的关键特性：

• Python中CUDA的快速迭代。
• 用于模拟的丰富数学和几何图元集合。
• 一次编写，可同时针对CPU/GPU。
• 作为独立、开源库发布。

参与Omniverse开发者日

第33页的标题为“JOIN OMNIVERSE DEVELOPER DAYS”，介绍了NVIDIA GTC大会上与Omniverse相关的特色演讲和特别活动：

特色演讲
* 为Omniverse制作连接器 [S41592]
* Lou Rohan, Connect Engineer, NVIDIA

• 深入研究Omniverse的微服务构建 [S41593]

  • Jozef van Eenbergen, Senior Software Engineering Manager, NVIDIA
  • Philippe Sawicki, Senior Software Engineer, NVIDIA

• 如何构建Omniverse套件的扩展和应用程序 [S41591]

  • Damien Fagnou, Senior Director, NVIDIA
  • Koen Vroegindstijn, Software Engineer, NVIDIA

• 使用Omniverse Replicator进行合成数据生成 [S41581]

  • Gavriel State Senior Director - Simulation and AI NVIDIA

• Warp：面向GPU模拟的高性能Python框架 [S41599]

  • Miles Macklin, Principal Engineer, NVIDIA

特别活动与交流

• 开发者专题：物理模拟新纪元 [SE2312]

  • 3月21日 - 6:00 AM PDT / 2:00 PM CET

• 分组会议：使用NVIDIA Kaolin进行3D深度学习研究 [SE42115]

  • 3月21日 - 7:00 AM PDT / 3:00 PM CET

• 开发者专题：在Omniverse上构建您自己的微服务 [SE2306]

  • 3月21日 - 2:00 PM PDT / 10:00 PM CET

• 开发者专题：构建Omniverse扩展和应用程序 [SE2305]

  • 3月21日 - 3:00 PM PDT / 12 AM CET

• 与Omniverse专家见面：材料与渲染技术 [SE2310]

  • 3月23日 - 8:00 AM PDT / 4:00 PM CET

会议日期：3月21-24日, 2022 | http://www.nvidia.com/GTC

结束页

第34页是NVIDIA的品牌宣传图片，展示了一个抽象的绿色螺旋形图案，背景为黑色，左下角是NVIDIA的标志。