NanoVDB: per-leaf VBM build kernel; cache leafCount for async rebuilds

[swahtz]

Two optimizations to the CUDA VoxelBlockManager build (nanovdb/tools/cuda/VoxelBlockManager.cuh):

1. Per-leaf build kernel. BuildVoxelBlockManagerFunctor now launches one thread per leaf node, indexing the dense leaf array directly (valid since the build requires isSequential()), instead of one thread per lower-node child slot gated on the childMask. This removes the 2D (lowerCount × slices) launch, the getChild indirection, and all idle threads on sparse lower nodes, and matches the structure of the host-side builder. nBlocks is now uint64_t end-to-end.

2. Stream-asynchronous rebuilds. VoxelBlockManagerHandle caches the source grid's leaf count (defaulted trailing constructor parameter; existing callers unaffected). The in-place device rebuild overload uses the cached value instead of reading TreeData back from device memory on every call, removing a synchronous D2H copy from the rebuild path; if the handle lacks a cached count it is read once and cached. The allocating device overload now performs a single D2H TreeData read supplying both the voxel count and the leaf count (previously up to two).

The cached count is a launch-sizing hint only: the build kernel reads the authoritative leaf count from the device-resident tree and iterates with a grid-stride loop, so an undersized launch strides over the tail and an oversized one retires excess threads at the device-side bound. A stale cache (e.g. after regenerating the grid in place under a live handle) therefore degrades launch geometry but cannot cause out-of-bounds access or wrong output — verified byte-identical with the cache deliberately poisoned to ~1/7× and ~3× of the true count. The jumpMap is also zeroed before the empty-grid early return, so a rebuild against a grid reporting zero leaves leaves zeroed metadata rather than stale bits.

Performance

Measured on a 317M-active-voxel OnIndex grid (782k leaves, 1266 lower nodes, ~15% lower-node child occupancy) on an RTX PRO 6000 (Blackwell), CUDA 13.2. Per-rebuild averages over 100 iterations of the in-place rebuild overload:

Log2BlockWidth	rebuild, sync'd each call	rebuild, pipelined (single sync at end)
7	49.7 µs → 31.3 µs (1.6×)	40.2 µs → 25.1 µs (1.6×)
9	44.4 µs → 25.0 µs (1.8×)	36.9 µs → 18.6 µs (2.0×)
12	41.7 µs → 24.5 µs (1.7×)	37.0 µs → 18.1 µs (2.0×)

Notes on attribution of the gain:

• A bare memset of the 39.6 MB jumpMap (identical work in both versions, the floor for any rebuild) takes 9.1 µs, so the build kernel itself went from ~28–31 µs to ~9–16 µs (~2.5–3×). This matches the launch-width change: the old scheme launched 5.19M threads against 782k useful ones on this grid.
• The old rebuild barely benefited from back-to-back launching (40 µs pipelined vs 50 µs sync'd) because the per-call getTreeData D2H copy serialized every call; the new rebuild pipelines down to ~18 µs and is genuinely stream-ordered.
• The one-time allocating build at Log2BlockWidth=12 improved from 78 µs to 50 µs (one D2H read instead of two).

Verification

Device-built firstLeafID/jumpMap arrays compared byte-for-byte against the unchanged CPU reference builder (nanovdb/tools/VoxelBlockManager.h) on a dense 21³ box and a sparse spherical shell at Log2BlockWidth 6, 7, and 12, including repeated rebuilds through the cached-leafCount path — all identical. The existing TestNanoVDBCUDA.VoxelBlockManager_* unit tests cover the decode path end-to-end on top of this metadata.

🤖 Generated with Claude Code