1. 背景

Warp-MACA 是适配于沐曦、曦云系列 GPU 的 Warp 框架，提供完整的 Warp 功能，通过镜像、安装包、源码的形式发布，并能作为依赖支持第三方库。许可证类型为 Apache 2.0。

Warp 是一个用于编写高性能仿真和图形代码的 Python 框架。Warp 可以将普通的 Python 函数通过 JIT 转换为高效的内核代码，能够在 CPU 或 GPU 上运行。Warp 让开发者能够轻松编写物理仿真、感知、机器人以及几何处理等领域的程序。此外，Warp 的内核是可微分的，可以无缝集成到 PyTorch、JAX 和 Paddle 等机器学习框架的训练流程中。API 和语言参考请见官方链接：https://nvidia.github.io/warp/index.html

Warp 提供了多种示例：

2. 环境搭建

2.1. 环境需求

Hardware: MetaX C-series

OS: Linux

Software:

Python >= 3.10
Docker support

2.2. 构建

Warp-MACA 目前仅支持 MetaX 开发者社区发布的 maca-pytorch:*-torch2.4-py310-*-amd64 镜像。

下载源码

git clone https://github.com/MetaX-MACA/hpc_warp.git
cd hpc_warp

设置环境变量

export MACA_PATH=/opt/maca
export CUCC_PATH=/opt/maca/tools/cu-bridge
export CUDA_PATH=/opt/maca/tools/cu-bridge
export PATH=${CUCC_PATH}/bin:${CUCC_PATH}/tools:${MACA_CLANG_PATH}:${MACA_PATH}/bin:${CUDA_PATH}/bin:$PATH
export LD_LIBRARY_PATH=${CUCC_PATH}/lib:${MACA_PATH}/lib:${MACA_PATH}/ompi/lib:${MACA_PATH}/ucx/lib:$LD_LIBRARY_PATH
export MACA_DIRECT_DISPATCH=1

构建

python build_lib.py --no_libmathdx

安装

python install -e .

3. Warp-MACA 示例

作为示例，安装 Warp-MACA 后可以运行以下 Python 代码。@wp.kernel 装饰符将 Python 函数通过 JIT 编译为 MACA 代码，所有的参数都必须标注类型。length 内核计算 points 中每个点到原点的距离。

import warp as wp
import numpy as np

N = 1024

@wp.kernel
def length(points: wp.array(dtype=wp.vec3), lengths: wp.array(dtype=float)):
    # thread index
    tid = wp.tid()

    # compute distance of each point from origin
    lengths[tid] = wp.length(points[tid])

# allocate an array of 3d points
points = [[np.sin(i), np.cos(i), np.tan(i)] for i in range(N)]
points = wp.array(points, dtype=wp.vec3, requires_grad=True)
lengths = wp.zeros(N, dtype=float, requires_grad=True)

# launch kernel
wp.launch(kernel=length, dim=N, inputs=[points, lengths])

print(lengths)

输出：

[1.        1.8508157 2.402998  ... 1.0001551 1.8016782 2.4995737]

@warp.kernel 装饰符不支持 JIT 编译共享内存等功能，但共享内存可以通过 @wp.func_native 装饰符在内核中插入 C++ 代码实现。

snippet = """
    __shared__ float sum[%d];

    sum[tid] = arr[tid];
    __syncthreads();
    for (int stride = %d; stride > 0; stride >>= 1) {
        if (tid < stride)
            sum[tid] += sum[tid + stride];
        __syncthreads();
    }
    if (tid == 0)
        out[0] = sum[0];
    """
snippet %= (N, N // 2)
adj_snippet = """
    adj_arr[tid] = 1;
    """

@wp.func_native(snippet, adj_snippet)
def reduce(arr: wp.array(dtype=float), out: wp.array(dtype=float), tid: int):
    return

@wp.kernel
def reduce_kernel(arr: wp.array(dtype=float), out: wp.array(dtype=float)):
    tid = wp.tid()
    reduce(arr, out, tid)

s = wp.zeros(1, dtype=float, requires_grad=True)

wp.launch(kernel=reduce_kernel, dim=N, inputs=[lengths, s])

print(s)

输出：

[1062.221]

@wp.kernel 标注的内核都实现了自动微分。wp.func_native 的自动微分通过 adj_snippet 参数实现。以下代码用 @wp.Tape 从 s 反向传播计算 lengths、points 的梯度。

tape = wp.Tape()
with tape:
    wp.launch(kernel=length, dim=N, inputs=[points, lengths])
    wp.launch(kernel=reduce_kernel, dim=N, inputs=[lengths, s])

# run backward
tape.backward(s)

print(lengths.grad)
print(points.grad)

# dl/dx = x/l
print(np.array_equal(points.grad.numpy(), (points / lengths).numpy()))

输出：

[1. 1. 1. ... 1. 1. 1.]
[[ 0.          1.          0.        ]
 [ 0.4546487   0.29192656  0.841471  ]
 [ 0.37840125 -0.1731782  -0.9092974 ]
 ...
 [ 0.01760877 -0.9996899  -0.01761151]
 [-0.46169454 -0.3080673   0.83182496]
 [-0.36665666  0.16005456 -0.91648537]]
True

4. 支持的第三方库

Warp-MACA 支持以下第三方库，不需要修改代码即可在曦云系列 GPU 上运行：

Rewarped：一个用于在并行可微分多物理模拟中进行强化学习的平台。
XLB (Accelerated Lattice Boltzmann)：一个带有 Warp 后端选项的格子玻尔兹曼求解器，由 Autodesk 维护。
warp-mpm: 一个 MPM（物质点法）模拟器，被用于诸如 Neural Stress Fields for Reduced-order Elastoplasticity and Fracture 和 PhysGaussian: Physics-Integrated 3D Gaussians for Generative Dynamics 等项目中。

5. Warp 原理

Warp 的基本原理为：@wp.kernel 修饰符将 Python 函数编译为 C++ 代码，再通过 JIT 编译为可执行模块。@wp.launch 模块导入可执行模块并启动内核。

6. 致谢

感谢 NVIDIA 开发并开源 Warp 框架。本项目基于 NVIDIA Warp 进行移植与适配，谨向原作者及 NVIDIA 团队致以诚挚谢意。