NVIDIA/warp

NVIDIA Warp – 用于 GPU 加速模拟和可微分物理的 Python

它是什么 – Warp 是一个 Python 库，可让您编写普通的 Python 函数，并将其即时编译（JIT）为快速的 GPU（或 CPU）内核。它针对高性能模拟任务，例如粒子动力学、流体流动、有限元分析和几何处理，且内核是可微分的，因此可以集成到 PyTorch、JAX 或 Paddle 流程中。

关键概念

• 内核装饰器 – @wp.kernel 将函数标记为 GPU 内核。在内部，您使用 Warp 的数组和数学原语（wp.array、wp.vec3、wp.tid()、wp.length_sq 等）。
• 自动 JIT – 在启动时，Warp 将内核编译为 CUDA（或 CPU）代码，为您处理线程索引和内存布局。
• 可微分 – 梯度可流经 Warp 内核，实现物理模拟的端到端学习。
• 互操作性 – Warp 张量可以与 NumPy 和深度学习框架相互转换，使模拟易于嵌入训练循环中。

典型工作流程

1. 从 PyPI 安装 warp-lang 包（或用于 CUDA-12 的 nightly wheel）。这些 wheel 包含针对现代 NVIDIA GPU 预编译的 CUDA 内核。
2. 使用 @wp.kernel 装饰器编写内核，并使用 Warp 的类型化数组（wp.array[wp.vec3] 等）。
3. 使用 wp.launch(kernel, dim=…, inputs=[…]) 启动内核。
4. 将结果转换回 NumPy（.numpy()）或直接馈入 PyTorch 模型。

为何重要

• 速度 – 通过编译为原生 CUDA 内核，Warp 模拟数百万个粒子或求解 PDE 的速度比纯 Python 或 NumPy 快数个数量级。
• 可微分模拟 – 研究人员可以训练控制或从物理中学习的神经网络（例如机器人控制、流体感知），而无需编写自定义 CUDA 梯度代码。
• 广泛的领域覆盖 – 此库提供重力、SPH 流体、纳维-斯托克斯方程、有限元法、光线追踪、网格处理等原语，以及大量可直接运行的示例。

开始使用

pip install warp-lang          # 核心库
pip install warp-lang[examples]  # 可选的示例依赖项

然后尝试 README 中的 20 行重力演示，或探索 NVIDIA Accelerated Computing Hub 中链接的许多笔记本。

进一步了解

• 官方文档：https://nvidia.github.io/warp/stable/
• GTC 和 SIGGRAPH 演讲（README 中的链接），逐步介绍如何构建可微分模拟。
• warp/examples 目录包含从粒子系统到 FEM 形状优化的现成演示。

社区 – 在 Apache 2.0 下开源。欢迎通过 GitHub 仓库提交问题、功能请求和拉取请求，并提供专用电子邮件地址（warp-python@nvidia.com）供非 GitHub 查询使用。

---

以上所有详细信息均直接取自仓库的 README；未添加任何外部假设。