v1.16.0

[shi-eric]

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.0

Non-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.py

For 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 small float32 rotations, the existing unsigned wp.quat_twist_angle() is now more accurate (Stabilize wp.quat_twist_angle() near zero and add a signed variant #1631).
• CUDA blocking streams implicitly synchronize with the legacy default stream, while non-blocking streams do not. wp.Stream.is_blocking reports 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 (Add documentation for working with non-blocking CUDA streams #1618).
• Array type annotations now produce evaluable, subscript-style repr() output such as wp.array4d[wp.uint32]. Generated API documentation uses the same form that appears in annotations (Array annotation repr() uses constructor form instead of subscript syntax #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 (@wp.kernel(module="unique") Python-side decoration cost makes factory-pattern kernels expensive #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 (Make the copy "Incompatible array data types" error show data types #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.py uses NCCL, NVIDIA's topology-aware inter-GPU communication library, through nccl4py, with MPI retained for rank setup (Add NCCL Jacobi example #1576).
• example_jacobi_nvshmem.py uses NVSHMEM, NVIDIA's symmetric-memory library for GPU clusters, to allocate Warp arrays on the symmetric heap and exchange halo rows through host-initiated nvshmem4py transfers (CUDA 12 package, CUDA 13 package), with MPI retained for bootstrap (Add NVSHMEM Jacobi example #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. pip does 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(...) or wp.mat22(...) when launching kernels or assigning struct fields (Deprecate Implicit Promotion of Numbers to Composite Types #1022).
• wp.Texture.copy_from_array() and wp.Texture.copy_to_array() are scheduled for removal in Warp 1.17. Use wp.Texture.copy_from() and wp.Texture.copy_to() instead ([REQ] Support zero-copy interop with external textures #1238).
• The warp.jax_experimental namespace will remain available through Warp 1.17. Its removal, including the legacy custom-call jax_kernel() and graph-cache getter and setter APIs, is now scheduled for Warp 1.18. Migrate to the top-level wp.jax_kernel() and wp.jax_callable() APIs (Promote warp.jax_experimental API to warp.jax #1370).
• warp.config.verbose and warp.config.quiet are scheduled for removal in Warp 1.18. Use warp.config.log_level = warp.LOG_DEBUG and warp.config.log_level = warp.LOG_WARNING, respectively. warp.config.verbose_warnings is unaffected ([BUG]: warp.utils.warn() overrides user warning filters, making deprecation warnings unsuppressible #1315).
• wp.HashGridQueryH and wp.HashGridQueryD are scheduled for removal in Warp 1.18. Use wp.HashGridQuery in public type annotations (Consolidate documented HashGrid query types #1452).
• The legacy max_verts, max_tris, and device arguments and compatibility attributes on wp.MarchingCubes are scheduled for removal in Warp 1.19. This includes the max_verts and max_tris arguments to wp.MarchingCubes.resize() and the id and runtime attributes. Output arrays are dynamically sized, and extraction uses the input field's device (Deprecate legacy MarchingCubes arguments and state #1594).
• The masked=True form of warp.sparse topology-changing operations is scheduled for removal in Warp 1.19. Use topology="masked" with bsr_set_from_triplets(), bsr_assign(), bsr_set_transpose(), bsr_axpy(), and bsr_mm() (Add 4-array BSR storage with row capacity #1537).
• The per-environment sequence form of warp.fem.Nanogrid.from_environment_voxels() and warp.fem.AdaptiveNanogrid.from_environment_voxels() is deprecated. It is scheduled for removal in Warp 1.19. Pass flat points, cell_levels where applicable, point_envs, and env_count instead (Graph-capturable NanoVDB volume allocation and rebuild #1606).
• warp.sparse.BsrMatrix.copy_nnz_async() is deprecated. It is scheduled for removal in Warp 1.20. Use warp.sparse.BsrMatrix.notify_nnz_changed() after updating the nonzero count ([REQ] Allocation-free path for masked sparse matrix operations #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 (Expose cudaProfilerStart and cudaProfilerStop for targeted captures #1596, Stabilize wp.quat_twist_angle() near zero and add a signed variant #1631).
• @pei-tian for fixing square-element node coordinates (node_coords_in_element() in SquareNedelecFirstKindShapeFunctions and SquareRaviartThomasShapeFunctions do not return coords #1685).
• @bpachev-nvidia for fixing warp.fem.lookup() with wp.float64 warp.fem.Grid2D and warp.fem.Grid3D geometries (fem.lookup fails to compile with FP64 geometries #1660) and updating the FEM examples for current Matplotlib releases.
• @Nas01010101 for preserving Adam optimizer state when resetting wp.float16 parameters (Adam.set_params re-zeros fp16 optimizer moment state on every call #1593).
• @snico432 for investigating loop-index reuse and raising the scoping question addressed by the final fix (Reusing for-loop variable name in while loop gives wrong result #1534).
• @loney7 for early implementation work toward JAX FFI CPU support (Add JAX FFI Host support #1446), which informed the feature shipped in [REQ] Support JAX FFI calls on CPU #1661.

---

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

---

This discussion was created from the release v1.16.0.