NanoVDB: own the small builders' device scratch with cuda::Buffer (CUDA) - #2286
Conversation
Route the last raw mallocAsync/freeAsync pairs in the GPU tools through RAII cuda::Buffer scratch over an injectable resource, completing the audit that PointsToGrid, TopologyBuilder, MeshToGrid and the topology-op consumers already went through: IndexToGrid (node accessor + grid name), addBlindData (byte-size scratch), GridStats (per-node statistics scratch) and SignedFloodFill (node counts). Each gains a defaulted ResourceT template parameter following the established conventions -- classes take a must-outlive resource instance in the constructor, free functions take the template parameter and route through the per-type default instance. Every replaced free keeps its position in stream order via destroy(stream) at the old freeAsync site, and GridStats' mode-dependent allocation becomes a zero-element buffer when the mode has no average, so the resource sees identical traffic. IndexToGrid's nested NodeAccessor is hoisted to the namespace-scope IndexToGridNodeAccessor so the kernels that name it do not acquire a spurious ResourceT dependency; a public alias preserves the nested spelling. Seam tests with exact allocation counts cover all four (including the grid-name path, the Extrema zero-allocation path, and free-function forwarding to the default instance). Part of AcademySoftwareFoundation#2232 (step 2). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Mark Harris <mharris@nvidia.com>
The temporary NodeManager in GridStats::update allocates through the dual-space DeviceBuffer, which cannot yet take a resource; make that visible on update() so users injecting a resource for accounting or limits know this allocation bypasses it until the single-space handle work lands. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Mark Harris <mharris@nvidia.com>
kmuseth
left a comment
There was a problem hiding this comment.
Seems like minor changes to adopt the new API for remaining tools. I am however confused as to why one member struct was changed to a namespaced struct
|
|
||
| template<typename SrcBuildT> | ||
| struct IndexToGrid<SrcBuildT>::NodeAccessor | ||
| struct IndexToGridNodeAccessor |
There was a problem hiding this comment.
why is this no longer a member class of NodeAccesspr?
There was a problem hiding this comment.
Nested types inherit all of the enclosing class's template parameters, so keeping it nested would make the kernels' (below) argument type vary with ResourceT and that would cause every kernel to be re-instantiated for every resource type. This would impact compile time and binary size.
Hoisting the struct outside the class keeps its type keyed on SrcBuildT only, with a member alias preserving the old spelling inside the class.
This is the same as TopologyBuilderData in #2268.
Bump the nanovdb pin to openvdb master 7946f17e to pick up the small-builder ResourceT seams (AcademySoftwareFoundation/openvdb#2286), plus the synchronous resource adapters (#2272) and the MeshToGrid CALL_CUBS undef fix (#2284). The CUDA save path allocated its largest buffers from nanoVDB's default DeviceBuffer 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 (its device NodeAccessor) routes through TorchResource via the new #2286 seam. Stream-ordering semantics are preserved: the replaced DeviceBuffer constructors were stream-ordered cudaMallocAsync on the current stream, and TorchDeviceBuffer allocates via raw_alloc, which torch orders on the current stream -- identical here since every construction sits under the existing CUDAGuard with the same current stream the copies and kernels are queued on. The host path (indexToGridHost) and the HostBuffer file-staging buffers are unchanged. Verified: tests/unit/test_io.py (622 passed) and a traced save (FVDB_NANOVDB_TRACE_ALLOCS=2) showing the indexToGrid scratch flowing through TorchResource. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>
… via upstream memory-resource seams (#732) ## 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](AcademySoftwareFoundation/openvdb#2268), [#2269](AcademySoftwareFoundation/openvdb#2269), [#2270](AcademySoftwareFoundation/openvdb#2270), [#2272](AcademySoftwareFoundation/openvdb#2272), [#2273](AcademySoftwareFoundation/openvdb#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](AcademySoftwareFoundation/openvdb#2286)), the synchronous resource adapters ([#2272](AcademySoftwareFoundation/openvdb#2272)), and the `MeshToGrid` `CALL_CUBS` `#undef` fix ([#2284](AcademySoftwareFoundation/openvdb#2284)). 2. **`fvdb::TorchResource`.** A ~40-line stateless resource ([`src/fvdb/TorchResource.h`](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`](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 #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 #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 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 - [x] `./build.sh install` succeeds on a clean tree against the new pin (full CUDA build, `-Werror`). - [x] 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. - [x] 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). - [x] `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 (#656) onto this branch in place of #655, threading `TorchResource` through the new ops it adds. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Jonathan Swartz <jonathan@jswartz.info> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Part of #2232 (step 2). Completes the raw-allocation audit that PointsToGrid (#2270), TopologyBuilder/MeshToGrid (#2268/#2269) and the topology-op consumers (#2273) already went through: the last ten
mallocAsync/freeAsyncpairs in the GPU tools become RAIIcuda::Bufferscratch over an injectable resource.What
ResourceTtemplate parameter following the established conventions: classes take a must-outlive resource instance in the constructor (as DilateGrid et al.), free functions take the template parameter and route through the per-type default instance (aspointsToGrid). Existing code is unaffected — every call site in the tree was traced and remains source-compatible.IndexToGrid's nestedNodeAccessoris hoisted to a namespace-scopeIndexToGridNodeAccessor<SrcBuildT>so the kernels that name it do not acquire a spuriousResourceTdependency; a public alias preserves the nested spelling.destroy(stream)at the oldfreeAsyncsite), and GridStats' mode-dependent allocation becomes a zero-element buffer when the mode has no average — the resource sees identical traffic, asserted exactly by the tests.Tests
New seam tests in
TestMemoryResource.cuwith exact allocation counts per builder (IndexToGrid: 2, including the grid-name buffer; GridStats: 1, and 0 for the Extrema mode; SignedFloodFill: 1; addBlindData: 1 through the default-instance path), plus free-functionResourceTforwarding coverage.Two observations for the reviewer
IndexToGrid::setGridNameuploads the name to the device, but no kernel readsd_gridName— the name never reaches the output grid. This PR keeps that behavior (the upload now goes through the resource); fixing or removing it should be its own change.GridStats::update's temporary NodeManager still allocates through the dual-spaceDeviceBuffer(createNodeManagerishasDeviceDual-gated), so that allocation bypasses an injected resource. Now called out in an@noteonupdate(); the real fix belongs to the single-space handle work on NanoVDB: injectable CUDA memory resources — design & roadmap #2232.Verification
--Werror=all-warnings; full build of tests, tools and examples atCMAKE_CUDA_ARCHITECTURES=80.gpt-5.6-terra): one finding (the NodeManager boundary above), resolved; second pass clean with the reviewer independently building and running the test binary.🤖 Generated with Claude Code