[swahtz]

Summary

Adopts NanoVDB's VoxelBlockManager (VBM) for per-active-voxel iteration on CUDA. The leaf scan launches one thread per 512-slot leaf position and discards inactive slots. The VBM decode costs O(1) per active voxel, in registers, so its cost tracks active voxels instead of leaf slots. A batch's VBMs are built together in one launch, and a batch is iterated in one launch.

First step of #817, which tracks the remaining consumers and follow-ups.

Depends on AcademySoftwareFoundation/openvdb#2291 for buildVoxelBlockManagerBatch. Until it merges, get_nanovdb.cmake pins NanoVDB to its head (swahtz/openvdb 29165acb, on top of develop/nanovdb). The pin moves back to AcademySoftwareFoundation/openvdb before this merges.

Changes

1. VbmCache: lazily built VBMs for every grid of a GridBatchData (~0.16 B/voxel), reached through GridBatchData::vbm() / vbmAt(bi).
   • The first use builds all of the batch's VBMs into one buffer with one memset and one kernel. Grid g's VBM starts at block floor(cumVoxels_g / 128) + g, which leaves room for every grid without a scan (at most one unused block per grid), so the build kernel derives each grid's entry from the batch's device-side voxel offsets and metadata. Sizes come from host metadata, and the build never synchronizes the host.
   • Each GridBatchData, views included, owns its own cache. Never invalidated (topology is immutable) and never serialized.
   • The buffer is one tensor from the CUDA caching allocator, read through typed uint32_t/uint64_t pointers. The build event is created only when a second stream first acquires the cache; that stream waits on it and record_stream()s the buffer once, and later calls from it skip both.
   • The destructor takes no device guard, so freeing a grid in a forked child cannot throw.
2. forEachActiveVoxelVbmCUDA: one launch per batch. Each 128-thread group takes one VBM block and finds its grid by binary search on the batch's compact voxel offsets (8 bytes per grid). Blocks between grids are skipped.
3. reinitialize_sdf: stencil kernels use the register decode instead of the shared-memory one, and reuse the cached VBM instead of rebuilding it every call.
4. BasePerActiveVoxelProcessor (.ijk, morton/hilbert): CUDA path iterates via the VBM. CPU and PrivateUse1 keep the leaf scan.
5. PadGrid: uses allocateProcessedRootTiles, since the newer develop/nanovdb deprecates allocateProcessedRoot.
6. Benchmark: single-grid occupancy sweep, build cost, and batches of 1-256 random or shell grids from 32³ to 256³, warm and cold.

Performance

active_voxel_iteration_benchmark, RTX PRO 6000 Blackwell (sm_120), CUDA 13.4, median of 5.

Batches

Each grid is either a 25% random subsample of a dim³ box (every leaf present, about a quarter full) or a 3-voxel-thick spherical shell inscribed in it (the narrow-band case). Warm reuses the batch's VBMs. Cold runs on a freshly built batch, so the VBM path includes the build. Speedup is leaf scan ÷ VBM.

Approximate voxels per grid:

dim	random	shell
32	8k	5k
64	66k	19k
128	524k	79k
256	4.2M	316k

Random:

grids	dim	warm leaf scan	warm VBM	warm ×	cold leaf scan	cold VBM	cold ×
1	32	8.59 µs	8.96 µs	0.96×	9.18 µs	16.4 µs	0.56×
64	32	22.7 µs	13.9 µs	1.63×	28.1 µs	26.1 µs	1.08×
256	32	65.7 µs	29.7 µs	2.21×	97.6 µs	45.8 µs	2.13×
4	64	15.1 µs	10.4 µs	1.45×	15.2 µs	18.1 µs	0.84×
16	64	35.7 µs	17.1 µs	2.09×	36.0 µs	25.3 µs	1.42×
256	64	628 µs	205 µs	3.06×	628 µs	208 µs	3.02×
1	128	20.3 µs	11.9 µs	1.71×	22.3 µs	19.7 µs	1.13×
4	128	57.1 µs	23.9 µs	2.39×	74.3 µs	32.6 µs	2.28×
64	128	1.12 ms	369 µs	3.04×	1.12 ms	367 µs	3.05×
256	128	4.45 ms	1.52 ms	2.94×	4.45 ms	1.52 ms	2.93×
1	256	105 µs	37.1 µs	2.83×	138 µs	46.3 µs	2.98×
16	256	2.20 ms	696 µs	3.16×	2.20 ms	695 µs	3.16×
64	256	8.79 ms	2.88 ms	3.05×	8.76 ms	2.86 ms	3.06×

Shell:

grids	dim	warm leaf scan	warm VBM	warm ×	cold leaf scan	cold VBM	cold ×
1	32	8.72 µs	9.06 µs	0.96×	9.07 µs	15.6 µs	0.58×
16	64	19.8 µs	12.1 µs	1.64×	20.1 µs	20.9 µs	0.96×
64	64	55.3 µs	22.7 µs	2.44×	73.3 µs	34.7 µs	2.11×
1	128	10.5 µs	9.28 µs	1.13×	12.6 µs	17.1 µs	0.74×
4	128	18.7 µs	11.3 µs	1.65×	23.6 µs	19.6 µs	1.20×
16	128	51.7 µs	21.2 µs	2.44×	68.6 µs	31.2 µs	2.20×
256	128	1.02 ms	305 µs	3.36×	1.03 ms	306 µs	3.35×
1	256	17.8 µs	11.2 µs	1.59×	22.1 µs	21.2 µs	1.04×
4	256	45.7 µs	19.9 µs	2.30×	60.9 µs	28.4 µs	2.14×
64	256	901 µs	286 µs	3.15×	901 µs	284 µs	3.17×

Warm, the VBM reaches ~3× from a few hundred thousand voxels per batch for both shapes. Below ~30k voxels both paths are launch-bound at ~9 µs and the VBM is up to 4% slower.

Cold, the build adds a fixed ~7 µs per batch over the leaf scan (VbmBuild alone, including its sync, is 11.8-12.0 µs at any occupancy). It pays for itself at about 300-500k active voxels per batch, for both shapes. We accept that cost: it is paid once per batch, it does not grow with grid count, and large or reused batches win.

One grid, by leaf occupancy

Random subsamples of a 160³ box, warm cache:

leaf occupancy	leaf scan	VBM	speedup
100%	32.8 µs	33.3 µs	0.98×
95%	33.3 µs	35.0 µs	0.95×
90%	33.0 µs	33.4 µs	0.99×
80%	32.5 µs	30.7 µs	1.06×
70%	32.3 µs	27.9 µs	1.16×
50%	32.1 µs	22.3 µs	1.44×
30%	31.9 µs	16.8 µs	1.90×
20%	31.9 µs	14.1 µs	2.26×
10%	31.6 µs	11.5 µs	2.75×
5%	31.2 µs	9.97 µs	3.13×
narrow shell (22%)	17.5 µs	11.1 µs	1.58×

Dense single grids are within 5% of the leaf scan, and the VBM wins from 80% occupancy down. The per-grid kernel this replaces was 1.1× faster on dense grids; the batched kernel's grid lookup costs that back, and buys the batch wins above.

End to end

reinitialize_sdf (order 3, repeated calls, 396k-voxel narrow shell) took 889 µs on main at f655b758 and 715 µs here (1.24×), measured before the batched rework. Its kernels are unchanged by the rework.

Notes

• Block width. The decode's rank step reads 2^L/64 jumpMap words per voxel. At L=9 (512) the VBM was ~12% slower than the leaf scan on dense grids. L=7 keeps it at parity.
• Launch shape. VBM kernels use 256-thread blocks. A <<<blockCount, 512>>> launch with __launch_bounds__(512) silently dropped work from a register-heavy callback on sm_120 / CUDA 13.0.
• GatherScatterDefault stays on the leaf scan. Its topology build is bound by counter-atomic contention, which Warp-aggregate the per-tap counter atomics in the conv topology build #838 addresses first. Routing it through the VBM is tracked in Extend VoxelBlockManager iteration across CUDA ops (follow-ups to #747) #817.

Testing

• VbmCacheTest (8): decode parity vs the leaf scan, value index == slot, views match their parent bitwise, empty grids, block-edge sizes, cross-stream lifetime, CPU builders sequential, and decode parity on 13 CUDA builders.
• pytest: .ijk/morton CPU-vs-CUDA parity including sliced batches and large batches with leading, interior and trailing empty grids; reinitialize_sdf sliced-view parity; fork safety with a warm cache.

./build.sh install gtests benchmarks
build/<build-dir>/gtests/VbmCacheTest                               # 8 passed
cd tests && pytest unit/test_active_grid_coords_vbm.py unit/test_fork_safety.py unit/test_sdf.py \
    unit/test_morton_hilbert.py unit/test_sliced_batch.py unit/test_basic_ops.py unit/test_batching.py \
    unit/test_dual.py unit/test_batched_topology_builder.py unit/test_conv_semantics_integration.py \
    unit/test_sample.py unit/test_ray_marching.py                   # 1021 passed, 5 skipped
build/<build-dir>/gbenchmarks/active_voxel_iteration_benchmark --benchmark_repetitions=5

🤖 Generated with Claude Code

[matthewdcong]

prefer something as close to firstLeafID.const_data_ptr<uint32_t> as possible here (even if it's just removing the const cast portion if the tensor can't be made native uint32)

[swahtz]

firstLeafID is now a native kUInt32 tensor (and jumpMap is kUInt64), so the view uses mFirstLeafID.const_data_ptr<uint32_t>() and mJumpMap.const_data_ptr<uint64_t>() with no casts.

[swahtz]

Actually, this changed a bit after the cold-build optimization work. The VBM's two arrays now share one allocation, a kUInt64 tensor, so a cold build makes one torch::empty and one record_stream instead of two. jumpMap reads through mutable_data_ptr<uint64_t>(). firstLeafID is packed after it and takes one reinterpret_cast<uint32_t *>(jumpMap + jumpMapWords). I tried narrow(...).view(torch::kUInt32) to avoid the cast, but it added about 1.5 µs of tensor dispatch to every cold build, roughly a fifth of the remaining fixed cost. So one commented cast stays.

[matthewdcong]

Do we have multiple host threads modifying this cache? Is the mutex necessary?

[swahtz]

Not from Python today, since no fvdb binding releases the GIL. I'd still make an argument for keeping it for two reasons:

1. The cache is built lazily behind a const GridBatchData &, and several Python grid objects and views can share one GridBatchData. The usual C++ contract is that const calls are safe to make concurrently, and a lazy build that writes members breaks that without a lock.
2. Autograd runs C++ backward functions on its own per-device threads without the GIL. Nothing in backward touches the VBM yet, but the channel-folded ops in Extend VoxelBlockManager iteration across CUDA ops (follow-ups to #747) #817 (pools, refine, inject) will, and their backward could then race the forward of another op on the same grid.

An uncontended std::mutex costs tens of nanoseconds, which is small next to the kernel launch it guards. In the reworked cache, the lock also protects the list of streams that have waited on the build (the fix for @harrism's repeated cudaStreamWaitEvent comment). So std::call_once alone wouldn't cover everything.

[harrism]

Thanks for this. The occupancy-independent iteration is a clear win, and caching the VBM on GridBatchData with view-stable byte-offset keys looks right to me. I checked that views resolve grids by physical offset the same way, and that featureIdx = cumVoxelsAt(bi) + slot - 1 matches the leaf scan's baseOffset + getValue - 1.

Two things I think need fixing before merge, both inline: the cold build's header read isn't ordered after the current stream, and ~VbmCache can call std::terminate (including in forked DataLoader workers, which #815 just fixed for grid teardown). I'd also like a measurement on a batched workload, since the path changes from one batch-wide launch to per-grid launches with a host sync per new grid.

Smaller: the cold build's blocking copy also makes .ijk/morton on a fresh grid fail under CUDA-graph capture. The PR notes this and openvdb#2291 removes it, but since these were capturable before, a line in the docs (or a clear error) would help until then.

[harrism]

At the pinned NanoVDB (0ab0a81e), buildVoxelBlockManager reads lowerCount through DeviceGridTraits::getTreeData, which is a plain cudaMemcpy on the legacy default stream. fvdb isn't built with per-thread default streams, and torch's pool streams are created cudaStreamNonBlocking, so that copy isn't ordered after stream.

Scenario: build a grid under torch.cuda.stream(s) and call .ijk on it immediately. The header read can run before the builder's writes on s land. That gives a wrong lowerCount (a short VBM, or an out-of-bounds node read for a garbage count), and the bad entry is cached for the life of the grid. The leaf-scan path never read device memory from the host, so it didn't have this exposure. (BuildPaddedGrid/PadGrid already make the same read, so it's not new to fvdb, but this PR puts it on the .ijk/morton path.)

Since the build blocks anyway, a cheap interim fix is to C10_CUDA_CHECK(cudaStreamSynchronize(stream.stream())) before calling buildVoxelBlockManager, until openvdb#2291 removes the read.

[swahtz]

Fixed by removing the read rather than syncing around it. The cache now builds all of a batch's VBMs at once through buildVoxelBlockManagerBatch from openvdb#2291, which this PR pins to for now. The build is sized entirely from host metadata (voxel and leaf counts, block bases), and each grid's entry is derived inside the build kernel from the batch's device-side voxel offsets. Nothing is read back from the device, so there is no host sync and no unordered copy on the .ijk/morton path.

BuildPaddedGrid and PadGrid still make the getTreeData read you mentioned. That predates this PR, so I've left it out of scope here.

[harrism]

If constructing this CUDAGuard throws, the exception escapes an implicitly noexcept destructor and calls std::terminate; the _WARN on cudaEventDestroy doesn't help with that. It can throw in a child forked after CUDA initialized, which is the case #815 fixed for grid teardown (merged after this branch was cut). With this PR, any grid whose .ijk, morton or reinitialize_sdf ran before a DataLoader fork would abort the worker when the child frees it. The same applies to teardown at process exit after the CUDA runtime has started unloading.

The merged tests/unit/test_fork_safety.py won't catch it after a rebase, because it never touches .ijk, so the cache is empty and the loop body never runs. Suggest dropping the guard (or catching around it, so the destroy is best-effort), and extending the fork test to warm the cache, e.g. call grid.ijk before the fork.

[swahtz]

Fixed. The destructor no longer takes a device guard, and it destroys the event with C10_CUDA_IGNORE_ERROR, so teardown in a forked child or at process exit is best-effort and can't throw. test_fork_safety.py now calls grid.ijk before the fork, so the child frees a grid whose cache is built.

One limit remains, and it's PyTorch's rather than ours. If a batch's VBMs are acquired from a second stream, the cache record_stream()s its buffer for that stream. Freeing any record_stream()ed tensor in a forked child aborts, because PyTorch's allocator records an event on the recorded stream at free time, and that throws in the child. I reproduced it with a plain tensor and no fVDB. Single-stream use never calls record_stream(), and the build event is now created only for a second stream, so the common DataLoader case frees normally. The limit is documented in VbmCache.h (e184a04).

[harrism]

This turns one batch-wide launch (the leaf scan over totalLeaves) into, per grid: a cache lookup under the mutex, a kernel launch, and on first use a blocking build. Grids are immutable, so code that makes new grids pays the host sync on every new grid. For example, _sdf.py's band loop calls dilated.ijk on a fresh grid each iteration, as does a training step that builds a new batch of grids.

The benchmarks in the description are all single-grid. Could you add a data point for a batch of many small, freshly built grids (say 64-256 grids, cold cache) against main? If it regresses, keeping the leaf scan for cold or large batches, or building all of a batch's VBMs with a single sync, might be worth it.

[swahtz]

This and making those benchmarks made me make some deeper changes to openvdb#2291 to allow for batched VBM creation. This was reworked along the lines you suggested.

Each batch now gets one build (one memset and one kernel, with no host sync) and one iteration launch over all of its grids, in place of one launch and one blocking build per grid. The description has the data point you asked for: batches of 1-256 grids, random or narrow-shell, from 32³ to 256³, warm and cold. They're compared against the leaf-scan kernel that main uses, in active_voxel_iteration_benchmark, not as a Python timing against main itself.

Cold batches still pay a fixed cost of about 7 µs over the leaf scan. The VBM pays that back at about 300-500k active voxels per batch, and wins up to ~3x above that. Below it, small fresh batches are slower. One fresh 32³ grid takes about 16 µs against 9 µs with the leaf scan, which is roughly the _sdf.py band-loop case. We decided to accept that rather than keep the leaf scan for cold batches: it's a fixed cost once per batch, it doesn't grow with grid count, and it no longer syncs the host.

[harrism]

Minor: every hit from a stream other than the build stream repeats cudaStreamWaitEvent plus two record_stream calls, per grid and per op call, under mMutex, even when that stream has already waited. Keeping a small set of already-synchronized streams in the Entry would make repeat hits free for apps that run on a non-default stream.

[swahtz]

Fixed. The cache remembers which streams have already waited on the build, so a repeat hit from one of them skips cudaStreamWaitEvent and record_stream entirely. The VBM now lives in one buffer, so the first hit from a new stream makes one record_stream call instead of two. The build event is also created lazily, on the first acquire from a second stream, so single-stream use never creates one.

[harrism]

Nit: get takes any GridBatchData, but entries are keyed only by cumBytesAt(bi). So a.vbmCache().get(b, i) with a batch that doesn't share this cache's storage would silently return another grid's VBM. Every call site also passes the same batch twice (batch.vbmCache().get(batch, bi)). A GridBatchData::vbmAt(bi) accessor, with get private or a TORCH_CHECK that the batch owns this cache, would close that off.

[swahtz]

Fixed. GridBatchData::vbm() and vbmAt(bi) replace vbmCache().get(batch, bi). The cache records the batch that built it, and acquire checks it's called with that batch. Views no longer share their parent's cache keyed by byte offset: every GridBatchData, views included, owns its own.

[harrism]

Nit: the slot-equals-value-index shortcut relies on the grid being sequential, which NanoVDB only checks with NANOVDB_ASSERT (compiled out in release). Grids from fvdb's builders are fine. A grid that's ill-formed already breaks the leaf scan too (e.g. a ValueIndex grid retyped by LoadNanovdb's copyIndexGridToHandle without renumbering), so this isn't new. But a cheap isSequential check from host metadata when the cache entry is built would turn silent mis-ordering into an error.

[swahtz]

Agreed that silent mis-ordering should be an error, but I don't think an isSequential() check at cache build would catch the case you describe. isSequential() checks that the nodes are fixed-size, which is always true for OnIndexGrid, and that the grid header has its breadth-first flag set. A ValueIndex grid that copyIndexGridToHandle retypes to ValueOnIndex keeps its breadth-first layout and header flag, so the check passes. What breaks is the numbering: the leaf offsets still count every value slot, so the active-voxel indices have gaps. Checking the flag at cache build would also mean reading the grid header back from the device, the kind of host sync this PR removes.

Every foreign grid enters fVDB through copyIndexGridToHandle in LoadNanovdb.cpp, on the host before upload, so that's the place to catch it: reject a ValueIndex source, or renumber it to ValueOnIndex, plus a host-side isSequential() check while the header is in reach. It already breaks the leaf scan too, so I've filed it separately as #841 rather than folding it into this PR.