NVIDIA/warp v1.10.0

v1.10.0

Repository: NVIDIA/warp

Tag: v1.10.0

Published: 2025-11-02T14:11:50Z

Prerelease: no

Release notes:

Warp v1.10.0

Warp v1.10 expands JAX integration with automatic differentiation support and multi-device jax.pmap() compatibility. The tile programming model has been enhanced with axis-specific reductions, component-level indexing, and convenience functions for creating tiles.

Performance has been significantly improved in several areas: BVH operations now support in-place rebuilding for CUDA graphs and configurable leaf sizes, built-in function calls from Python are up to 70× faster, and additional sparse matrix and FEM operations can now be captured in CUDA graphs.

Additional usability improvements include negative indexing and slicing for arrays, atomic bitwise operations, and new built-in functions including error functions and type casting.

Important: This release removes the warp.sim module (deprecated since v1.8), which has been superseded by the Newton physics engine. See the Announcements section below for migration guidance and other upcoming changes.

For a complete list of changes, see the full changelog.

New features

JAX automatic differentiation (experimental)

Warp now supports experimental automatic differentiation with JAX, allowing kernels to participate in JAX automatic differentiation workflows. This feature is contributed by @mehdiataei and builds on earlier work by @jaro-sevcik. It enables computing gradients through Warp kernels using jax.grad() by passing enable_backward=True to jax_kernel().

Key capabilities include:

• Single and multiple output kernels: Compute gradients for kernels with one or more output arrays
• Static input auto-detection: Scalar inputs are automatically treated as static (non-differentiable) arguments
• Vector and matrix arrays: Arrays of composite types like wp.vec2 or wp.mat22 are fully supported
• Multi-device execution: Compatible with jax.pmap() for distributed forward and backward passes across multiple GPUs

import jax
from warp.jax_experimental import jax_kernel

@wp.kernel
def my_kernel(a: wp.array(dtype=float), out: wp.array(dtype=float)):
i = wp.tid()
out[i] = a[i] ** 2.0

# Enable automatic differentiation
jax_func = jax_kernel(my_kernel, num_outputs=1, enable_backward=True)

# Compute gradients through the kernel
grad_fn = jax.grad(lambda a: jax.numpy.sum(jax_func(a)[0]))
gradient = grad_fn(input_array) # gradient: [2*a[0], 2*a[1], ...]

This feature is experimental and has some current limitations. See the JAX Automatic Differentiation documentation for complete examples, usage details, and limitations.

Multi-device JAX support with jax.pmap()

Warp now properly supports jax.pmap() and jax.shard_map() for multi-device parallel execution, thanks to fixes contributed by @chaserileyroberts. Previously, device targeting issues prevented Warp callables from working correctly within these JAX primitives—JAX would invoke callbacks from multiple threads targeting different devices, but Warp would always execute on the default device. The fix ensures proper device coordination by extracting device ordinals from XLA FFI and adding thread synchronization for concurrent callbacks, enabling efficient data-parallel workflows across multiple GPUs.

In-place BVH rebuilding with CUDA graph support

A new wp.Bvh.rebuild() method enables rebuilding BVH hierarchies in-place without allocating new memory. This complements the existing refit() method and is particularly useful when primitive distributions change significantly.

CUDA graph capture: Unlike creating a new BVH, rebuild() reuses existing buffers, making it safe to capture in CUDA graphs. Previously captured graphs that include queries on the BVH remain valid after rebuilding, enabling high-performance repeated updates without graph re-capture overhead.

Construction algorithms: On CUDA devices, in-place rebuild supports `"lbvh" only. On CPU, "sah" and "median"` are supported. Defaults are chosen automatically based on the device.

Tile programming enhancements

The tile programming model has been enhanced with new capabilities to make tile-based computations more expressive and convenient:

Axis-specific reductions

The tile-reduction functions wp.tile_reduce() and wp.tile_sum() now support an optional axis parameter, enabling reductions along a specific dimension of a tile rather than reducing the entire tile to a single value. This enhancement brings NumPy-like axis semantics to tile operations.

@wp.kernel
def tile_reduce_axis(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
a = wp.tile_load(x, shape=(4, 8), storage="shared")
# Sum along axis 0, reducing shape from (4, 8) to (8,)
b = wp.tile_sum(a, axis=0)
wp.tile_store(y, b)

x = wp.array(np.arange(32).reshape(4, 8), dtype=float)
# x = [[ 0. 1. 2. 3. 4. 5. 6. 7.]
# [ 8. 9. 10. 11. 12. 13. 14. 15.]
# [16. 17. 18. 19. 20. 21. 22. 23.]
# [24. 25. 26. 27. 28. 29. 30. 31.]]
y = wp.zeros(8, dtype=float)

wp.launch_tiled(tile_reduce_axis, dim=(1,), inputs=[x], outputs=[y], block_dim=32)
# y = [48. 52. 56. 60. 64. 68. 72. 76.] (column sums)

Component-level indexing

Tiles of composite types (vectors, matrices, quaternions) now support component-level indexing and assignment. You can directly index into individual components using extended indexing syntax:

• Vector components: tile[i][1] extracts the second component of a vector at position i
• Matrix elements: tile[i][1, 1] accesses the element at row 1, column 1 of a matrix at position i

This provides more convenient and expressive syntax for working with structured data in tiles.

Creating tiles filled with a constant value

The new wp.tile_full() function provides a convenient way to create tiles initialized with a constant value, similar to NumPy's np.full():

# Create an 8x8 tile filled with 3.14
tile = wp.tile_full(shape=(8, 8), value=3.14, dtype=float)

New example

The new `example_tile_mcgp.py` example demonstrates tile-based Monte Carlo methods by implementing a walk-on-spheres algorithm for...