Python That Compiles to CUDA When Physics Gets Heavy

What it does

Warp is a Python framework that takes decorated Python functions and JIT compiles them to efficient kernel code that can run on the CPU or GPU. It comes with a rich set of primitives for physics simulation, robotics, and geometry processing, and the kernels are differentiable, so they can be used as part of machine-learning pipelines.

The interesting bit

You write Python with typed arrays and scalar math, and Warp rewrites it into CUDA or CPU vector code at runtime. The kernels are differentiable and can be used as part of machine-learning pipelines with frameworks such as PyTorch, JAX, and Paddle.

Key highlights

• JIT compiles Python to CPU or CUDA kernels
• Kernels are differentiable and can be used with PyTorch, JAX, and Paddle
• Built-in primitives for physics, geometry, and robotics (examples include SPH fluids, FEM, and ray marching)
• Runs on x86-64 and ARMv8 (including Apple Silicon), though GPU support requires a CUDA-capable NVIDIA card
• Can generate USD time-sampled animations viewable in Blender or UsdView

Caveats

• GPU support requires a CUDA-capable NVIDIA GPU (minimum GeForce GTX 9xx), and a handful of examples are GPU-only
• On Linux aarch64, the framework automatically swaps in a USD dependency replacement because standard wheels are unavailable for that platform

Verdict

A strong fit for researchers and engineers who need custom differentiable physics or geometry kernels in Python and would rather not maintain a C++ CUDA codebase; less essential if your work is already handled by standard deep-learning tensor operations.