Route NanoVDB builder scratch through PyTorch's active CUDA allocator via upstream memory-resource seams

[swahtz]

Summary

fvdb's grid builders allocate their device scratch from nanoVDB's default DeviceResource — a second cudaMallocAsync pool that partitions VRAM against PyTorch's. Large workloads (e.g. multi-frame TSDF integration) then hit a clean OOM even when the GPU has free memory in aggregate.

This routes that scratch — O(N-points) sort keys, CUB temp storage, topology mask buffers — through PyTorch's CUDA allocator instead, so it shares one pool with fvdb / PyTorch tensors. 22 sites across 13 .cu files plus PadGrid.cuh.

Note this is not hardcoded to Torch's native caching allocator: c10::cuda::CUDACachingAllocator is a namespace, and its raw_alloc_with_stream / raw_delete free functions dispatch through CUDACachingAllocator::get() — the runtime-swappable allocator Torch itself allocates tensors from. fvdb's scratch therefore follows whatever allocator the user has installed: the native caching allocator (including PYTORCH_CUDA_ALLOC_CONF knobs), the cudaMallocAsync backend (PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync), or a custom allocator installed via torch.cuda.memory.change_current_allocator(CUDAPluggableAllocator(...)).

Supersedes #655, which vendored modified nanoVDB headers into the tree. This instead uses the injectable-memory-resource seams we developed upstream (AcademySoftwareFoundation/openvdb#2232; PRs #2268, #2269, #2270, #2272, #2273) — now merged, so the pin is plain upstream master. No fork, no include-path shadowing, no resync procedure.

What's in this PR

1. Pin nanovdb to upstream master. src/cmake/get_nanovdb.cmake → AcademySoftwareFoundation/openvdb @ 7946f17e, which includes the small-builder ResourceT seams (#2286), the synchronous resource adapters (#2272), and the MeshToGrid CALL_CUBS #undef fix (#2284).

2. fvdb::TorchResource. A ~40-line stateless resource (src/fvdb/TorchResource.h) modeling nanoVDB's stream-ordered AsyncResource concept over c10::cuda::CUDACachingAllocator::raw_alloc_with_stream / raw_delete — the dispatchers to Torch's currently active CUDA allocator (see Summary). Passed as the ResourceT template parameter at all 13 upstream builder call sites — voxelsToGrid, DilateGrid, MergeGrids, PruneGrid, RefineGrid, CoarsenGrid — always via the fvdb::BuilderResource alias (src/fvdb/BuilderResource.h), never named directly, so the allocator policy lives in a single line (a non-torch build, e.g. the ONNX Runtime EP planned in Proposal: ⚫ ONNX Runtime support for fVDB operations #579, retargets the alias there instead of touching every op). Being stateless, it binds through each builder's defaulted constructor argument, so no instance is plumbed through. Retains Fork a small set of nanoVDB headers to share PyTorch's CUDA allocator #655's FVDB_NANOVDB_TRACE_ALLOCS tracing (=1 traces ≥ 256 KiB, a value starting with 2 traces everything).

3. PadGrid gains a ResourceT seam. The conv builders used DilateGrid<..., TorchResource> for odd kernels and PadGrid — on the rival pool — for even ones: same loop, same grid, a different allocator depending on kernel parity. fvdb's own morphology::PadGrid drives nanoVDB's TopologyBuilder (internal mask buffers, countNodes CUB scratch, TempPool) but hardcoded DeviceResource. It now takes a ResourceT parameter mirroring the upstream DilateGrid signature and forwards it, with BuilderResource passed at all 7 call sites. The default keeps it source-compatible.

4. CUB scratch in BuildFineGridFromCoarse. cub::DeviceSegmentedReduce temp storage used a bare cudaMallocAsync; it now routes through BuilderResource. Both cub calls are also now C10_CUDA_CHECK-wrapped — previously unchecked, as was the allocation.

5. The SaveNanoVDB CUDA path. The save path allocated its largest device buffers from nanoVDB's default pool: the per-batch (N+1)-element value staging buffer, the indexToGrid output grid handle, and the defensive host-upload buffer. All three now use TorchDeviceBuffer, and indexToGrid's internal scratch routes through TorchResource via the #2286 seam. Stream-ordering is preserved: the replaced stream-ordered DeviceBuffer constructors become raw_alloc on the same current stream the copies and kernels are queued on. The host path (indexToGridHost) and the HostBuffer file-staging buffers are unchanged.

Not routed (no upstream seam yet; all off the hot paths)

• DistributedPointsToGrid multi-GPU scratch (deferred upstream behind NanoVDB: handle empty partitions in the distributed merge path search AcademySoftwareFoundation/openvdb#2248) — the most valuable remaining seam
• VoxelBlockManager / buildVoxelBlockManager scratch in ReinitializeSdf.cu
• the builders' small dual-space mProcessedRoot / mData buffers (upstream roadmap Step 3)

MeshToGrid is the one merged seam fvdb does not use: BuildGridFromMesh.cu does its own parametric surface sampling and goes through _createNanoGridFromIJK, so there is nothing to route.

Test plan

• ./build.sh install succeeds on a clean tree against the new pin (full CUDA build, -Werror).
• Injection verified live via FVDB_NANOVDB_TRACE_ALLOCS: a 500k-point Grid.from_points + dilated_grid(2) prints 42 TorchResource traces with correct results (484,631 → 8,461,871 voxels); from_nearest_voxels_to_points at 2M points shows PadGrid scratch routed.
• 584 tests passed — conv semantics + integration (203), conv/conv-transpose default + prune + empty grids (103), basic ops (276, 1 skipped), sliced batch (2, covering the BuildFineGridFromCoarse CUB path).
• test_io.py — 622 passed against the 7946f17e pin; a traced save_nanovdb (FVDB_NANOVDB_TRACE_ALLOCS=2) shows the indexToGrid scratch flowing through TorchResource.

Followups

• Rebase the TSDF / ESDF / Occupancy stack (Add TSDF / ESDF / Occupancy / Decay APIs and a fast marching-cubes variant #656) onto this branch in place of Fork a small set of nanoVDB headers to share PyTorch's CUDA allocator #655, threading TorchResource through the new ops it adds.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[harrism]

One suggestion on plumbing, not policy. The default here is right — raw_alloc_with_stream dispatches through CUDACachingAllocator::get(), i.e. the currently set torch allocator, so an application that installs a pluggable allocator (RMM's torch plug, etc.) already re-routes all of this scratch with no fvdb changes.

But naming TorchResource at ~22 leaf call sites hard-wires the choice the same way nanoVDB itself just spent five PRs un-hard-wiring. A one-line internal seam — namespace fvdb { using BuilderResource = TorchResource; } (or a member of the ops' context) — with the call sites naming BuilderResource gives the same binary and the same behavior, but the policy becomes a single line.

The future that needs it is already on the books: the ONNX plan (#579) runs these builders in a non-torch process where c10 doesn't exist, and an OrtAllocatorResource over ORT's arena would splice in at that alias (or graduate to a real injected instance through the ops' factory) instead of re-touching all 22 sites.

[swahtz]

Agreed — done in 416fa43. All 22 ResourceT bindings across the 13 ops files now name fvdb::BuilderResource, a one-line alias in src/fvdb/BuilderResource.h currently set to TorchResource.

Two notes on scope, called out in the header:

• The seam is compile-time and leans on the resource being stateless (builders bind nanovdb::cuda::default_resource<BuilderResource>() through their defaulted constructor args). An OrtAllocatorResource holding a per-session OrtAllocator* would additionally need an instance plumbed through the call sites — your "graduate to a real injected instance" — so the alias removes the naming churn but not that eventual plumbing. A non-torch build would also put the TorchResource include behind a build-time switch in this one header.
• The alias covers builder scratch only; allocations that are torch tensors by design (TorchDeviceBuffer, the SaveNanoVDB staging buffers) still name their types directly.

All 13 affected TUs compile cleanly against the branch's nanovdb pin.