Skip to content

v1.16.0

Latest

Choose a tag to compare

@github-actions github-actions released this 03 Aug 02:32
v1.16.0
86ec8b7

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 (#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 (#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 (#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:

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 (#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 (#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 (#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 (#1315).
  • wp.HashGridQueryH and wp.HashGridQueryD are scheduled for removal in Warp 1.18. Use wp.HashGridQuery in public type annotations (#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 (#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() (#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 (#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 (#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() with wp.float64 warp.fem.Grid2D and warp.fem.Grid3D geometries (#1660) and updating the FEM examples for current Matplotlib releases.
  • @Nas01010101 for preserving Adam optimizer state when resetting wp.float16 parameters (#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.