Warp v1.16.0
Warp v1.16 adds in-place rebuilding for fixed-capacity NanoVDB volumes, including during CUDA graph capture. This lets fluid simulations update sparse-grid topologies without allocating new volumes. The release also adds grouped HashGrid queries for multi-environment workloads, NumPy-style tile slicing, CPU support for JAX FFI wrappers, experimental support for replaying more operations in CPU graphs and saved .wrp graphs, and CUDA profiler range controls.
New features
Rebuild sparse NanoVDB volumes in place
Previously, a volume's sparse topology was fixed at allocation time. Simulations whose active grid changed from one step to the next, such as the affine particle-in-cell (APIC) fluid example, had to allocate a new volume. That allocation required host synchronization and could not be replayed inside a CUDA graph. Warp 1.16 adds fixed-capacity, in-place rebuilding so simulations can reuse the same volume and dependent FEM topology buffers across graph replays (#1606).
The excerpt below follows that example's allocation, rebuild, and replay flow. It omits particle initialization, capacity estimation, FEM space construction, and the APIC transfer and solve.
import warp as wp
import warp.fem as fem
# API-shape excerpt. Simulation setup and solver work are omitted.
# particle_q, voxel_size, grid_capacity, and frame_count come from that setup.
grid_status = wp.zeros(1, dtype=wp.uint32, device=particle_q.device)
volume = wp.Volume.allocate_by_voxels(
voxel_points=particle_q,
voxel_size=voxel_size,
device=particle_q.device,
rebuildable=True,
**grid_capacity,
status=grid_status,
)
grid = fem.Nanogrid(volume, rebuildable=True)
# ... build linear_basis_space and strain_space once from grid ...
def simulate():
# Rebuild from the particle positions updated by the previous step.
grid.rebuild(particle_q, status=grid_status)
linear_basis_space.topology.rebuild()
strain_space.topology.rebuild()
# ... transfer particles to the grid, solve, and advect particles ...
with wp.ScopedCapture(particle_q.device) as capture:
simulate()
for _ in range(frame_count):
wp.capture_launch(capture.graph)
# Host-side status and topology queries stay outside capture.
status = int(grid_status.numpy()[0])
stats = volume.get_active_stats()
print(f"rebuild status=0x{status:x}, active voxels={stats.voxel_count}")Capacity does not grow automatically. Pass a one-element wp.uint32 status array to the allocation method and to rebuild(), then check for wp.Volume.REBUILD_SUCCESS, a REBUILD_*_CAPACITY_EXCEEDED flag, or REBUILD_INVALID_INPUT. A tile-allocated volume must be rebuilt from tile positions, while a voxel-allocated volume must be rebuilt from voxel positions. The point_mask parameter can exclude input points, and Warp deduplicates repeated positions.
On CUDA, rebuildable allocation and rebuilding can run inside graph capture when memory-pool allocation is enabled. Exact, non-rebuildable allocation and host-side queries such as get_active_stats() must remain outside capture.
See example_apic_fluid.py for the complete capacity estimation, topology construction, and captured simulation loop.
API changes:
| Status | API | Change in Warp 1.16 |
|---|---|---|
| Existing class method | wp.Volume.allocate_by_tiles() |
Adds CPU support and the rebuildable, max_tiles, max_lower_nodes, max_upper_nodes, status, and point_mask parameters. |
| Existing class method | wp.Volume.allocate_by_voxels() |
Adds CPU support and the rebuildable, max_active_voxels, max_leaf_nodes, max_lower_nodes, max_upper_nodes, status, and point_mask parameters. |
| Existing constructor | warp.fem.Nanogrid() |
Adds the rebuildable parameter for retaining capacity-sized topology buffers. |
| New instance methods | wp.Volume.rebuild() and warp.fem.Nanogrid.rebuild() |
Rebuild the volume directly or rebuild it through a Nanogrid that also refreshes its geometry buffers. |
| New topology methods | topology.rebuild() on compatible Nanogrid function spaces |
Refresh topology buffers after rebuilding the geometry. |
| New inspection APIs | wp.Volume.is_rebuildable, wp.Volume.get_rebuild_info(), and wp.Volume.get_active_stats() |
Report whether a volume can be rebuilt, its reserved capacity, and its current active topology. |
| New status constants | wp.Volume.REBUILD_* |
Report success, exceeded capacities, or invalid input through the optional status array. |
| New kernel builtin | wp.volume_voxel_count() |
Returns the current active voxel or index count inside kernels and captured graphs. |
Multi-environment spatial queries
Grouped HashGrid construction and queries
Independent simulation environments can now share one spatial grid without visiting candidates from neighboring worlds. Assign each point an int32 group in HashGrid.build(), then pass the desired group to wp.hash_grid_query() (#1579).
import warp as wp
@wp.kernel
def count_neighbors(
grid: wp.uint64,
points: wp.array[wp.vec3],
groups: wp.array[wp.int32],
counts: wp.array[wp.int32],
):
i = wp.tid()
count = int(0)
# Restrict the query to candidates in the current point's environment.
for j in wp.hash_grid_query(grid, points[i], 0.5, groups[i]):
if wp.length(points[j] - points[i]) <= 0.5:
count += 1
counts[i] = count
points = wp.array([(0.0, 0.0, 0.0), (0.4, 0.0, 0.0)] * 2, dtype=wp.vec3)
# The repeated two-point scene represents environments 0 and 1.
groups = wp.array([0, 0, 1, 1], dtype=wp.int32)
counts = wp.zeros(4, dtype=wp.int32)
grid = wp.HashGrid(8, 8, 8)
grid.build(points, radius=0.5, groups=groups)
wp.launch(count_neighbors, dim=4, inputs=[wp.uint64(grid.id), points, groups], outputs=[counts])
print(f"same-group neighbor counts: {counts.numpy()}") # same-group neighbor counts: [2 2 2 2]Group IDs may be any int32 values and can change between rebuilds, including during graph replay. If the query omits the group argument, it visits all candidates. Before capturing a grouped CUDA rebuild without a warm-up build, call grid.reserve(num_points, with_groups=True). Hash-grid queries return cell candidates, so kernels should still test the actual distance.
Tile programming
NumPy-style tile slicing
Tile kernels now use NumPy-style indexing to select, reverse, stride, gather, and assign subregions. An integer index collapses a dimension and may be negative. Use wp.tile_slice_indexed() for one-axis gathers such as tile[indices, :] (#1176).
import numpy as np
import warp as wp
@wp.kernel
def slice_tile(src: wp.array2d[float], dst: wp.array2d[float]):
tile = wp.tile_load(src, shape=(4, 4))
# Reverse the rows and keep columns 0 and 2.
wp.tile_store(dst, tile[::-1, ::2])
src = wp.array(np.arange(16, dtype=np.float32).reshape(4, 4), device="cuda:0")
dst = wp.zeros((4, 2), dtype=float, device="cuda:0")
# Launch one cooperative 32-thread block for the single tile.
wp.launch_tiled(slice_tile, dim=[1], inputs=[src], outputs=[dst], block_dim=32, device="cuda:0")
print(f"sliced tile (reversed rows, even columns):\n{dst.numpy()}")Output:
sliced tile (reversed rows, even columns):
[[12. 14.]
[ 8. 10.]
[ 4. 6.]
[ 0. 2.]]
The source tile in a slice assignment must have a matching shape. Tile views do not support compound assignment, and slices do not support scalar broadcasting. Indexed gathers select one axis at a time and require full : slices on the other axes.
The out-of-place wp.tile_lower_solve() now supports gradients for vector and matrix right-hand sides. This makes lower-triangular solves available in wp.Tape backward passes (#1378).
JAX integration
Run Warp FFI callbacks on CPU
JAX can now run wp.jax_kernel() and wp.jax_callable() wrappers on CPU as well as CUDA. JAX selects the device, so a CPU jax.jit program invokes the Warp callback without routing buffers through a GPU (#1661).
import jax
import jax.numpy as jnp
import numpy as np
import warp as wp
@wp.kernel
def triple(x: wp.array[float], out: wp.array[float]):
i = wp.tid()
out[i] = 3.0 * x[i]
run = wp.jax_kernel(triple)
with jax.default_device(jax.devices("cpu")[0]):
# JAX places both the arrays and the Warp FFI callback on CPU.
(result,) = jax.jit(run)(jnp.arange(4, dtype=jnp.float32))
print(f"Warp FFI result on CPU: {np.asarray(result)}") # Warp FFI result on CPU: [0. 3. 6. 9.]JAX 0.5.0 or newer is required. JaxCallableGraphMode.NONE and JaxCallableGraphMode.JAX work on CPU. The Warp-managed graph modes remain CUDA-only. example_jax_kernel.py shows more call shapes.
Set CUDA block sizes for JAX tile kernels
JAX FFI wrappers created with wp.jax_kernel() now accept block_dim, letting a CUDA tile kernel choose how many threads cooperate on each tile. The value is fixed when the wrapper is constructed and is also used for generated adjoint launches (#1436).
import jax
import jax.numpy as jnp
import numpy as np
import warp as wp
ROW_COUNT = 4
TILE_SIZE = 256
TILE_THREADS = 64
@wp.kernel
def row_sum(values: wp.array2d[float], output: wp.array[float]):
row = wp.tid()
# Threads in one CUDA block cooperate to reduce each row.
tile = wp.tile_load(values[row], shape=TILE_SIZE)
wp.tile_store(output, wp.tile_sum(tile), offset=row)
# JAX FFI uses wp.launch(), so append the block width to the logical row count.
jax_row_sum = wp.jax_kernel(
row_sum,
launch_dims=(ROW_COUNT, TILE_THREADS),
output_dims=(ROW_COUNT,),
block_dim=TILE_THREADS,
)
with jax.default_device(jax.devices("gpu")[0]):
values = jnp.arange(ROW_COUNT * TILE_SIZE, dtype=jnp.float32).reshape(ROW_COUNT, TILE_SIZE)
(row_sums,) = jax.jit(jax_row_sum)(values)
print(f"CUDA tile row sums: {np.asarray(row_sums)}") # CUDA tile row sums: [ 32640. 98176. 163712. 229248.]TILE_SIZE is the logical data shape, while TILE_THREADS is a CUDA execution choice. Keep output_dims equal to the logical result shape rather than including the block width. See CUDA block dimensions and tile kernels for more detail.
Graph capture
Replay more operations in CPU and saved .wrp graphs
Important
This is an experimental feature. The API may change without a formal deprecation cycle.
API Capture can now record wp.utils.array_sum() and wp.utils.array_inner() on CPU and CUDA, allowing saved .wrp graphs to recompute those reductions from current inputs during replay. Live CPU graphs can also rebuild wp.HashGrid data and refit or rebuild wp.Bvh trees from their current arrays (#1663, #1664, #1665).
import warp as wp
values = wp.array([1.0, 2.0, 3.0], dtype=float, device="cpu")
total = wp.zeros(1, dtype=float, device="cpu")
dot = wp.zeros(1, dtype=float, device="cpu")
with wp.ScopedCapture(device="cpu") as capture:
wp.utils.array_sum(values, out=total)
wp.utils.array_inner(values, values, out=dot)
wp.capture_launch(capture.graph)
print(
f"replayed reductions: sum={total.numpy()[0]}, "
f"inner product={dot.numpy()[0]}"
) # replayed reductions: sum=6.0, inner product=14.0Non-empty reductions recorded through API Capture require an explicit output array. Negative strides are unsupported. Counts and layout strides must fit a signed 32-bit integer. Because resources use process-local handles, graphs that record HashGrid.build(), Bvh.refit(), or Bvh.rebuild() are replay-only. wp.capture_save() rejects these graphs instead of writing a non-portable file. See CPU Graphs and Saving and Loading Graphs for the full lists of operations and limitations.
Profiling
Target external profiler captures
External profilers can now skip initialization and warm-up work and capture only the steady-state region. wp.cuda_profiler_start(), wp.cuda_profiler_stop(), and wp.ScopedCudaProfiler expose CUDA profiler range controls from Python (#1596).
import warp as wp
@wp.kernel
def increment(values: wp.array[float]):
values[wp.tid()] += 1.0
values = wp.zeros(1024, dtype=float)
# Compile and warm up before starting profiler collection.
wp.launch(increment, dim=values.size, inputs=[values])
wp.synchronize_device()
with wp.ScopedCudaProfiler():
for _ in range(10):
wp.launch(increment, dim=values.size, inputs=[values])
# Finish asynchronous launches before profiler collection stops.
wp.synchronize_device()The warm-up launch stays outside the capture range. The synchronization inside the scope ensures that the asynchronous launches complete before wp.cuda_profiler_stop() marks the end of collection. CUDA does not guarantee that the stop call synchronizes the device.
Configure the external profiler to honor these calls. For example:
nsys profile --capture-range=cudaProfilerApi python my_app.pyFor Nsight Compute, use --profile-from-start off. Warp invokes the profiler-control calls with the selected device's CUDA context current; the profiler and its configuration determine which activity is collected during that interval.
Additional API improvements
wp.quat_twist_angle_signed()recovers the signed rotational coordinate represented by a quaternion. For smallfloat32rotations, the existing unsignedwp.quat_twist_angle()is now more accurate (#1631).- CUDA blocking streams implicitly synchronize with the legacy default stream, while non-blocking streams do not.
wp.Stream.is_blockingreports which behavior a stream uses. Interoperability code can check borrowed streams before releasing temporary Warp resources, then either synchronize a non-blocking stream on the host or make a blocking Warp stream wait on it. See Warp's non-blocking stream guidance for both patterns (#1618). - Array type annotations now produce evaluable, subscript-style
repr()output such aswp.array4d[wp.uint32]. Generated API documentation uses the same form that appears in annotations (#1628).
Performance and diagnostics
- Repeated
@wp.kernel(module="unique")declarations from factory functions are roughly 2x faster in microbenchmarks, whether their captured values produce identical or specialized kernels (#1486). - Errors from array copies, texture copies, reshapes, views, and DLPack source-device checks now include the relevant shapes, data types, channels, or device identifiers (#1644).
New examples
Two new distributed Jacobi solver examples expand on example_jacobi_mpi.py, which uses CUDA-aware mpi4py directly for nearest-neighbor halo exchange. The new variants keep MPI for process setup but move the halo transfers to GPU communication libraries:
example_jacobi_nccl.pyuses NCCL, NVIDIA's topology-aware inter-GPU communication library, through nccl4py, with MPI retained for rank setup (#1576).example_jacobi_nvshmem.pyuses NVSHMEM, NVIDIA's symmetric-memory library for GPU clusters, to allocate Warp arrays on the symmetric heap and exchange halo rows through host-initiatednvshmem4pytransfers (CUDA 12 package, CUDA 13 package), with MPI retained for bootstrap (#1582).
Announcements
Tentative CUDA 13 PyPI wheel builds
- Warp's PyPI wheels are still tentatively planned to move to CUDA 13.x with Warp 1.17. The switch depends on CUDA Toolkit 13.4 being available in time and may move to a later Warp release if CUDA 13.4 is delayed. If the switch proceeds, an R580-series or newer NVIDIA driver will be required to use the CUDA backend from those wheels.
pipdoes not check the installed driver and may install a wheel whose CUDA backend cannot run on an older driver. - CUDA 12 support will continue after the PyPI wheel transition. Users who need CUDA 12 will be able to build Warp from source or install a pre-built CUDA 12 wheel from GitHub Releases.
This repeats the tentative plan announced in Warp v1.15.0. See NVIDIA's CUDA minor-version compatibility table for driver requirements.
Upcoming removals
- Implicit conversion of scalar values to composite types is scheduled for removal in Warp 1.17. Use an explicit constructor such as
wp.vec3(...)orwp.mat22(...)when launching kernels or assigning struct fields (#1022). wp.Texture.copy_from_array()andwp.Texture.copy_to_array()are scheduled for removal in Warp 1.17. Usewp.Texture.copy_from()andwp.Texture.copy_to()instead (#1238).- The
warp.jax_experimentalnamespace will remain available through Warp 1.17. Its removal, including the legacy custom-calljax_kernel()and graph-cache getter and setter APIs, is now scheduled for Warp 1.18. Migrate to the top-levelwp.jax_kernel()andwp.jax_callable()APIs (#1370). warp.config.verboseandwarp.config.quietare scheduled for removal in Warp 1.18. Usewarp.config.log_level = warp.LOG_DEBUGandwarp.config.log_level = warp.LOG_WARNING, respectively.warp.config.verbose_warningsis unaffected (#1315).wp.HashGridQueryHandwp.HashGridQueryDare scheduled for removal in Warp 1.18. Usewp.HashGridQueryin public type annotations (#1452).- The legacy
max_verts,max_tris, anddevicearguments and compatibility attributes onwp.MarchingCubesare scheduled for removal in Warp 1.19. This includes themax_vertsandmax_trisarguments towp.MarchingCubes.resize()and theidandruntimeattributes. Output arrays are dynamically sized, and extraction uses the input field's device (#1594). - The
masked=Trueform ofwarp.sparsetopology-changing operations is scheduled for removal in Warp 1.19. Usetopology="masked"withbsr_set_from_triplets(),bsr_assign(),bsr_set_transpose(),bsr_axpy(), andbsr_mm()(#1537). - The per-environment sequence form of
warp.fem.Nanogrid.from_environment_voxels()andwarp.fem.AdaptiveNanogrid.from_environment_voxels()is deprecated. It is scheduled for removal in Warp 1.19. Pass flatpoints,cell_levelswhere applicable,point_envs, andenv_countinstead (#1606). warp.sparse.BsrMatrix.copy_nnz_async()is deprecated. It is scheduled for removal in Warp 1.20. Usewarp.sparse.BsrMatrix.notify_nnz_changed()after updating the nonzero count (#987).
Acknowledgments
We also thank the following contributors from outside the core Warp development team:
- @felixmey for exposing CUDA profiler controls and improving quaternion twist-angle accuracy (#1596, #1631).
- @pei-tian for fixing square-element node coordinates (#1685).
- @bpachev-nvidia for fixing
warp.fem.lookup()withwp.float64warp.fem.Grid2Dandwarp.fem.Grid3Dgeometries (#1660) and updating the FEM examples for current Matplotlib releases. - @Nas01010101 for preserving Adam optimizer state when resetting
wp.float16parameters (#1593). - @snico432 for investigating loop-index reuse and raising the scoping question addressed by the final fix (#1534).
- @loney7 for early implementation work toward JAX FFI CPU support (#1446), which informed the feature shipped in #1661.
For a complete list of changes, see the full changelog.