NVIDIA CUDA Python编程框架--Warp开发文档

NVIDIA CUDA Python编程框架–Warp开发文档

Warp 是一个用于编写高性能模拟和图形代码的 Python 框架。 Warp 采用常规 Python 函数，JIT 将它们编译为可以在 CPU 或 GPU 上运行的高效内核代码。

Warp 专为空间计算而设计，并附带一组丰富的原语，可以轻松编写物理模拟、感知、机器人和几何处理程序。 此外，Warp 内核是可微分的，可以用作 PyTorch 和 JAX 等框架的机器学习管道的一部分。

以下是使用 Warp 实现的模拟的一些示例：

快速开始

Warp 支持 Python 3.7 及以上版本。 它可以在 Windows、Linux 和 macOS 上的 x86-64 和 ARMv8 CPU 上运行。 GPU 支持需要支持 CUDA 的 NVIDIA GPU 和驱动程序（最低 GeForce GTX 9xx）。

安装 Warp 最简单的方法是通过 PyPI：

$ pip install warp-lang

发布页面上还提供了预构建的二进制包。 要在本地 Python 环境中安装，请提取存档并从根目录运行以下命令：

$ pip install .

基本示例

下面给出了计算随机 3D 向量长度的第一个程序示例：

import warp as wp
import numpy as np

wp.init()

num_points = 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 = wp.array(np.random.rand(num_points, 3), dtype=wp.vec3)
lengths = wp.zeros(num_points, dtype=float)

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

print(lengths)

其他示例

Github 存储库中的示例目录包含许多脚本，这些脚本展示了如何使用 Warp API 实现不同的模拟方法。 大多数示例将在与示例相同的目录中生成包含时间采样动画的 USD 文件。 在运行示例之前，用户应确保使用以下命令安装了 usd-core 软件包：

pip install usd-core

可以从命令行运行示例，如下所示：

python -m warp.examples.<example_subdir>.<example>

大多数示例可以在 CPU 或支持 CUDA 的设备上运行，但少数示例需要支持 CUDA 的设备。 这些标记在示例脚本的顶部。

USD 文件可以在 NVIDIA Omniverse、Pixar 的 UsdView 和 Blender 中查看或渲染。 请注意，不建议在 macOS 中使用预览，因为它对时间采样动画的支持有限。

内置单元测试可以从命令行运行，如下所示：

python -m warp.tests

注意: 下方示例点击图片可以直接跳转到示例代码

在以后章节中会逐步更新编程细节

examples/core#

dem	fluid	graph capture	marching cubes
mesh	nvdb	raycast	raymarch
sph	torch	wave

examples/fem#

apic fluid	convection diffusion	diffusion 3d	diffusion
mixed elasticity	navier stokes	stokes transfer	stokes

examples/optim#

bounce	cloth throw	diffray	drone
inverse kinematics	spring cage	trajectory	walker

examples/sim#

cartpole	cloth	granular	granular collision sdf
jacobian ik	quadruped	rigid chain	rigid contact
rigid force	rigid gyroscopic	rigid soft contact	soft body