NVIDIA/warp

NVIDIA Warp – Python for GPU‑accelerated simulation and differentiable physics

What it is – Warp is a Python library that lets you write ordinary Python functions and have them JIT‑compiled into fast GPU (or CPU) kernels. It targets high‑performance simulation tasks such as particle dynamics, fluid flow, finite‑element analysis, and geometry processing, and the kernels are differentiable so they can be dropped into PyTorch, JAX or Paddle pipelines.

Key ideas

• Kernel decorator – @wp.kernel marks a function as a GPU kernel. Inside you use Warp’s array and math primitives (wp.array, wp.vec3, wp.tid(), wp.length_sq, etc.).
• Automatic JIT – At launch time Warp compiles the kernel to CUDA (or CPU) code, handling thread indexing and memory layout for you.
• Differentiable – Gradients flow through Warp kernels, enabling end‑to‑end learning on physical simulations.
• Interoperability – Warp tensors can be converted to/from NumPy and to deep‑learning frameworks, making it easy to embed simulations in training loops.

Typical workflow

1. Install the warp-lang package from PyPI (or a nightly wheel for CUDA‑12). The wheels include pre‑compiled CUDA kernels for modern NVIDIA GPUs.
2. Write a kernel with the @wp.kernel decorator, using Warp’s typed arrays (wp.array[wp.vec3], etc.).
3. Launch the kernel with wp.launch(kernel, dim=…, inputs=[…]).
4. Convert results back to NumPy (.numpy()) or feed them directly into a PyTorch model.

Why it matters

• Speed – By compiling to native CUDA kernels, Warp can simulate millions of particles or solve PDEs orders of magnitude faster than pure Python or NumPy.
• Differentiable simulation – Researchers can train neural networks that control or learn from physics (e.g., robot control, fluid‑aware perception) without writing custom CUDA gradient code.
• Broad domain coverage – The library ships with primitives for gravity, SPH fluids, Navier‑Stokes, FEM, ray‑marching, mesh processing, and more, plus a large collection of ready‑to‑run examples.

Getting started

pip install warp-lang          # core library
pip install warp-lang[examples]  # optional example dependencies

Then try the 20‑line gravity demo from the README, or explore the many notebooks linked from the NVIDIA Accelerated Computing Hub.

Where to learn more

• Official docs: https://nvidia.github.io/warp/stable/
• GTC and SIGGRAPH talks (links in the README) that walk through building differentiable simulations.
• The warp/examples directory contains ready‑made demos ranging from particle systems to FEM shape optimization.

Community – Open‑source under Apache 2.0. Issues, feature requests, and pull requests are welcomed via the GitHub repo, and a dedicated email address (warp-python@nvidia.com) is provided for non‑GitHub queries.

---

All details above are taken directly from the repository’s README; no external assumptions have been added.