NanoVDB Python: GPU/device bindings (nanovdb.cuda)

[swahtz]

Builds on feature/nanovdb_python (the CPU api-mirror epoch) to add the GPU/device half of the NanoVDB Python bindings: a real nanovdb.cuda namespace, zero-copy interop, streams, device-side managers, the tools.cuda algorithm surface, and cuda::* infrastructure.

(Note: This is a PR into the feature/nanovdb_python branch and #2219 is a PR of that branch into main. This is to make this PR easier to read so that it only has the deltas of the CUDA bindings. The idea is we review/approve both and then merge this PR into the feature branch and #2219 can then merge into main.)

Example (CuPy)

For functionality (beyond the functionality this PR provides) that needs to be compiled as a kernel for speed, you can now use CuPy to write custom kernels that get JIT compiled and can i/o buffers from NanoVDB Python which is kind of neato.

Build a grid on the device, view its buffer zero-copy as a CuPy array, then run your own CUDA kernel over it — compiled against the headers the wheel ships:

import cupy as cp
import nanovdb

# Build a float level-set sphere directly on the device, then upload it.
handle = nanovdb.tools.cuda.createLevelSetSphere(nanovdb.GridType.Float, 20)
handle.deviceUpload(0, True)            # raw int stream handle; 0 == default stream

# Zero-copy: view the whole device buffer as a CuPy array (no copy).
buf = cp.asarray(handle)                # via __cuda_array_interface__ / __dlpack__
assert int(buf.data.ptr) == handle.device_ptr()

# Run a custom kernel over the device grid. compile_options() supplies the
# -I for the NanoVDB headers bundled in the wheel.
device_grid = handle.deviceGrid(0)      # raw NanoGrid<float>* in GPU memory
kernel = cp.RawKernel(r'''
#include <nanovdb/NanoVDB.h>
extern "C" __global__
void inspect(const nanovdb::NanoGrid<float>* grid, float* out) {
    auto acc = grid->getAccessor();
    out[0] = acc.getValue(nanovdb::Coord(0, 0, 0));
    out[1] = float(grid->activeVoxelCount());
}
''', "inspect", options=nanovdb.cuda.compile_options("-std=c++17"), backend="nvrtc")

out = cp.zeros(2, dtype=cp.float32)
kernel((1,), (1,), (device_grid.data_ptr(), out.data.ptr))
cp.cuda.runtime.deviceSynchronize()
print(out)   # -> [-3.0, 30254.0]  (value@origin, activeVoxelCount)

nanovdb.io.deviceReadGrid("grid.nvdb") reads a grid straight onto the device the same way. See examples/ for the runnable versions (gpu_load_inspect.py, cupy_rawkernel.py) and the examples README for the full GPU walkthrough.

What's new

nanovdb.cuda (infrastructure)

• DeviceBuffer / DeviceGridHandle moved from the root module into nanovdb.cuda (breaking — see Migration), now with __cuda_array_interface__ v3, __dlpack__/__dlpack_device__, device_ptr/host_ptr, from_external (wrap caller-owned device/managed memory) and from_buffer.
• Stream-aware deviceUpload/deviceDownload (raw int stream handles).
• UnifiedBuffer/UnifiedGridHandle (managed memory), DeviceMesh, DeviceStreamMap, DeviceResource, TempDevicePool.
• Device createNodeManager and buildVoxelBlockManager (with device firstLeafID()/jumpMap() views).

nanovdb.tools.cuda (algorithms) — all take an optional stream

• Generalized pointsToGrid (world-space float/double → Point) alongside the voxel rasterizers (voxelsTo{OnIndex,Index,RGBA8}Grid, pointsToRGBA8Grid).
• Morphology/topology: dilateGrid/coarsenGrid/refineGrid/pruneGrid/mergeGrids.
• Index/QC: indexToGrid, addBlindData, updateGridStats, evalChecksum/validateChecksum/updateChecksum, isValid, signedFloodFill, sampleFromVoxels.
• DistributedPointsToGrid (multi-GPU; single-GPU path exercised here).
• tools.cuda.inject / injectPredicateToMask / injectGridMask — bind the sidecar-injection operators from util/cuda/Injection.cuh (the NanoVDB 2.0 paper's §3.4 injectData): copy per-voxel data (scalar or, via a 2-D overload, vector-valued) across a topology change over the src∩dst intersection, and build the per-leaf Mask<3> that pruneGrid consumes.
• tools.cuda.gatherBoxStencil / activeVoxelCoords — dense decode of the VBM stencil pipeline (decodeInverseMaps + computeBoxStencil are device-only): gatherBoxStencil writes each active voxel's 3×3×3 neighbour values to a (valueCount, 27) device array (centre 13; faces 4/10/12/14/16/22), activeVoxelCoords writes index-space positions to (valueCount, 3) — so dense array/tile frameworks (CuPy, cuTile) can run VDB stencils, and bake results back, without a hand-written decode kernel.

Custom-kernel interop — grid.data_ptr() exposes the device grid pointer and nanovdb.cuda.compile_options() returns the NVRTC include flags, e.g. a CuPy RawKernel that #include <nanovdb/NanoVDB.h> and reads a const NanoGrid<float>*.

Tests / examples / docs — TestGpuInterop.py, a 37-test CuPy/Torch-guarded suite (self-skips without CUDA/GPU) wired into ctest; four runnable examples (gpu_load_inspect, cupy_rawkernel, numba_cuda, triton_kernel); and a GPU section in the examples README.

Build/CI — nanobind bumped to 2.12.0 scoped to the NanoVDB CI job + the wheel (pyproject.toml pinned >=2.12.0,<3). The other install sites and FUTURE_MINIMUM_NANOBIND_VERSION stay at 2.5.0, and the binding code remains ≥2.5.0 source-compatible (the shared build.yml co-compiles it against 2.5.0).

examples/levelset_filter.py + three backends — one full GPU level-set filter (Laplacian deform + Godunov reinit + narrow-band retrack via the VBM and dilateGrid/inject/injectPredicateToMask/pruneGrid, style-preserving over FloatGrid / OnIndexGrid+blind-SDF .nvdb files). The driver levelset_filter.py owns everything that isn't backend-specific (read / write / retrack / outer loop) behind a small Backend protocol and a {rawkernel|cupy|cutile} selector; each backend supplies only the dense per-voxel stencils and also runs standalone as a stencil-only smoke test. The three differ only in that compute: levelset_filter_rawkernel.py (a fused cupy.RawModule CUDA kernel), levelset_filter_cupy.py (kernel-free CuPy array ops), and levelset_filter_cutile.py (NVIDIA cuTile @ct.kernels) — all producing identical results.

Also included

A core fix in tools/cuda/GridChecksum.cuh: both evalChecksum overloads copied a 256-entry-LUT-sized region back into the 4-byte checksum field (OOB read → runtime abort); corrected to sizeof(uint32_t). Independent of the Python work, but tools.cuda.evalChecksum aborts without it.

Validation

Validated locally on a Blackwell (sm_120) GPU with CuPy: CAI/DLPack zero-copy round-trips, streams, from_external, device NodeManager/VBM, the rasterizer/morphology/QC ops, and the managed-memory distributed path; nanobind 2.12.0 confirmed to compile both nanovdb_python and openvdb_python. CI does not exercise GPU paths — ctest excludes cuda/mgpu and the Python GPU tests self-skip via isCudaAvailable()/isGpuAvailable().