From 44d037534a4414067baaec36917022299435bf08 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Fri, 7 Aug 2026 17:03:33 +1200 Subject: [PATCH 1/8] Pin nanovdb to the injectable-memory-resource stack Points the nanovdb pin at the merge of AcademySoftwareFoundation/openvdb PRs #2268, #2269, #2270, #2272 and #2273 (tracking issue #2232), which give every CUDA builder fvdb uses a ResourceT injection seam. Temporary until the stack merges upstream; the pin is a fast-forward of the previous one (f9754140 is an ancestor of the stack's base). Co-Authored-By: Claude Fable 5 Signed-off-by: Jonathan Swartz --- src/cmake/get_nanovdb.cmake | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index d8ce343af..ef203d8f8 100644 --- a/src/cmake/get_nanovdb.cmake +++ b/src/cmake/get_nanovdb.cmake @@ -1,10 +1,16 @@ # Copyright Contributors to the OpenVDB Project # SPDX-License-Identifier: Apache-2.0 +# TEMPORARY: pinned to the merged nanovdb injectable-memory-resource stack +# (AcademySoftwareFoundation/openvdb PRs #2268, #2269, #2270, #2272, #2273 on +# top of upstream master d980ad92; tracking issue #2232). fvdb relies on the +# ResourceT seams these PRs add to route builder scratch through PyTorch's +# caching allocator (see src/fvdb/TorchResource.h). Repoint at +# AcademySoftwareFoundation/openvdb once the stack merges upstream. CPMAddPackage( NAME nanovdb - GITHUB_REPOSITORY AcademySoftwareFoundation/openvdb - GIT_TAG f9754140ba6031813b37d8e1b239ed0253ebd96d + GITHUB_REPOSITORY swahtz/openvdb + GIT_TAG 558bfb2ead7c993f980ba18e89fee4b69991fb39 SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) From 8aac07bd0df87d484de5ac2fb93c042b43a6fbfa Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Fri, 7 Aug 2026 17:03:33 +1200 Subject: [PATCH 2/8] Route nanovdb builder scratch through PyTorch's caching allocator Adds fvdb::TorchResource, a stream-ordered nanovdb memory resource backed by c10::cuda::CUDACachingAllocator (raw_alloc_with_stream / raw_delete), and passes it as the ResourceT template parameter at every builder call site: voxelsToGrid and the DilateGrid / MergeGrids / PruneGrid / RefineGrid / CoarsenGrid ops. Builder scratch (sort keys, CUB temp storage, topology mask buffers) now lives in the same pool as fvdb / PyTorch tensors instead of a second cudaMallocAsync pool that fragments VRAM against it. This supersedes the forked-header approach of #655: same allocator routing, but through upstream nanovdb's injection seams instead of shadowed copies of DeviceBuffer.h / DeviceResource.h. The FVDB_NANOVDB_TRACE_ALLOCS env var from that PR is preserved inside TorchResource (=1 traces allocs >= 256 KiB, =2 traces all). Not routed (no upstream seam yet, all off the hot paths): DistributedPointsToGrid multi-GPU scratch, indexToGrid scratch in SaveNanoVDB, and the builders' small dual-space mProcessedRoot / mData buffers. Co-Authored-By: Claude Fable 5 Signed-off-by: Jonathan Swartz --- src/fvdb/TorchResource.h | 86 +++++++++++++++++++ .../detail/ops/BuildCoarseGridFromFine.cu | 4 +- src/fvdb/detail/ops/BuildDenseGrid.cu | 6 +- src/fvdb/detail/ops/BuildDilatedGrid.cu | 4 +- .../detail/ops/BuildFineGridFromCoarse.cu | 4 +- src/fvdb/detail/ops/BuildGridForConv.cu | 5 +- .../detail/ops/BuildGridForConvTranspose.cu | 8 +- src/fvdb/detail/ops/BuildGridFromIjk.cu | 6 +- src/fvdb/detail/ops/BuildGridFromPoints.cu | 19 ++-- src/fvdb/detail/ops/BuildMergedGrids.cu | 4 +- src/fvdb/detail/ops/BuildPaddedGrid.cu | 4 +- src/fvdb/detail/ops/BuildPrunedGrid.cu | 4 +- 12 files changed, 131 insertions(+), 23 deletions(-) create mode 100644 src/fvdb/TorchResource.h diff --git a/src/fvdb/TorchResource.h b/src/fvdb/TorchResource.h new file mode 100644 index 000000000..baccf7425 --- /dev/null +++ b/src/fvdb/TorchResource.h @@ -0,0 +1,86 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#ifndef FVDB_TORCHRESOURCE_H +#define FVDB_TORCHRESOURCE_H + +#include + +#include + +#include +#include +#include + +namespace fvdb { + +/// @brief NanoVDB stream-ordered memory resource backed by PyTorch's CUDA caching +/// allocator (c10::cuda::CUDACachingAllocator). +/// +/// Passed as the ResourceT template parameter of NanoVDB's CUDA builders +/// (PointsToGrid / DilateGrid / MergeGrids / PruneGrid / RefineGrid / +/// CoarsenGrid), it routes their internal device scratch — O(N-points) sort +/// keys, CUB temp storage, topology mask buffers — through the same pool +/// that fvdb / PyTorch tensors use. Without this, nanoVDB's default +/// DeviceResource allocates from a second cudaMallocAsync pool that +/// partitions VRAM against torch's pool, and large workloads (e.g. +/// multi-frame TSDF integration) OOM even when the GPU has free memory in +/// aggregate. +/// +/// The resource is stateless, so builders can bind the shared instance +/// returned by nanovdb::cuda::default_resource() — naming +/// the template parameter at a call site is sufficient, no instance needs +/// to be threaded through. +/// +/// Set FVDB_NANOVDB_TRACE_ALLOCS=1 in the environment to trace allocations +/// of 256 KiB and larger to stderr (a value starting with '2' traces every +/// allocation). Useful for diagnosing topology-op memory blowup on large +/// scenes. +struct TorchResource : nanovdb::cuda::SyncFromAsync { + /// Alignment guaranteed by every allocation. The caching allocator returns + /// blocks aligned to at least 512 bytes, so advertising nanoVDB's + /// conventional 256 (matching cuda::DeviceResource) is always satisfied and + /// the alignment parameter below can be ignored. + static constexpr size_t DEFAULT_ALIGNMENT = 256; + + /// @brief Stream-ordered allocation from torch's caching allocator. + /// @note raw_alloc_with_stream records @p stream against the block so torch + /// defers reuse until work on it completes, matching the stream-ordered + /// semantics of the cudaMallocAsync call it replaces. Allocation + /// happens on the current device, like cudaMallocAsync. + void * + allocate_async(size_t bytes, size_t /*alignment*/, cudaStream_t stream) { + if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { + const size_t cutoff = (env[0] == '2') ? 0 : (1ull << 18); // '2' = trace all, else >= 256 KiB + if (bytes >= cutoff) { + std::fprintf(stderr, "[fvdb/nanovdb] TorchResource alloc %12zu bytes (%.3f MB)\n", + bytes, double(bytes) / 1e6); + } + } + void *p = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(bytes, stream); + if (!p) { + throw std::runtime_error("fvdb: TorchResource::allocate_async failed"); + } + return p; + } + + /// @brief Free through torch's caching allocator. + /// @note The stream argument is deliberately ignored: raw_delete relies on + /// the stream recorded at allocation time plus torch's per-stream event + /// tracking, so the free is safe without ordering on the caller's + /// stream. + void + deallocate_async(void *p, size_t /*bytes*/, size_t /*alignment*/, cudaStream_t /*stream*/) { + if (p == nullptr) { + return; + } + c10::cuda::CUDACachingAllocator::raw_delete(p); + } +}; + +static_assert(nanovdb::cuda::is_async_resource::value, + "TorchResource must model nanoVDB's stream-ordered AsyncResource concept"); + +} // namespace fvdb + +#endif // FVDB_TORCHRESOURCE_H diff --git a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu index 8600f8d83..ee60ddaee 100644 --- a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu +++ b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -84,7 +85,8 @@ coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch, TORCH_CHECK(grid, "Grid is null"); nanovdb::GridHandle handle; for (int p = 0; p < nPasses; p += 1) { - nanovdb::tools::cuda::CoarsenGrid op(grid, stream.stream()); + nanovdb::tools::cuda::CoarsenGrid op( + grid, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); handle = op.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildDenseGrid.cu b/src/fvdb/detail/ops/BuildDenseGrid.cu index f3705a8fe..da98cdaa4 100644 --- a/src/fvdb/detail/ops/BuildDenseGrid.cu +++ b/src/fvdb/detail/ops/BuildDenseGrid.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -143,8 +144,9 @@ dispatchCreateNanoGridFromDense(int64_t batchSize, handles.push_back(createEmptyGridHandle(guide.device())); } else if (i == 0) { handles.push_back( - nanovdb::tools::cuda::voxelsToGrid( - (nanovdb::Coord *)ijkData.data_ptr(), nVoxels, 1.0, guide)); + nanovdb::tools::cuda:: + voxelsToGrid( + (nanovdb::Coord *)ijkData.data_ptr(), nVoxels, 1.0, guide)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { handles.push_back(handles[0].copy(guide)); diff --git a/src/fvdb/detail/ops/BuildDilatedGrid.cu b/src/fvdb/detail/ops/BuildDilatedGrid.cu index 2139847dd..246cfbdbc 100644 --- a/src/fvdb/detail/ops/BuildDilatedGrid.cu +++ b/src/fvdb/detail/ops/BuildDilatedGrid.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -54,7 +55,8 @@ dispatchDilateGrid(const GridBatchData &gridBatch, TORCH_CHECK(grid, "Grid is null"); for (auto j = 0; j < dilationAmount[i]; j += 1) { - nanovdb::tools::cuda::DilateGrid dilateOp(grid, stream); + nanovdb::tools::cuda::DilateGrid dilateOp( + grid, stream); dilateOp.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); dilateOp.setChecksum(nanovdb::CheckMode::Default); dilateOp.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu index c2a0d83ee..a3298a647 100644 --- a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu +++ b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -424,7 +425,8 @@ fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl, TORCH_CHECK(grid, "Grid is null"); nanovdb::GridHandle handle; for (int p = 0; p < nPasses; p += 1) { - nanovdb::tools::cuda::RefineGrid op(grid, stream.stream()); + nanovdb::tools::cuda::RefineGrid op( + grid, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); handle = op.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildGridForConv.cu b/src/fvdb/detail/ops/BuildGridForConv.cu index 60f11070f..50a0896fe 100644 --- a/src/fvdb/detail/ops/BuildGridForConv.cu +++ b/src/fvdb/detail/ops/BuildGridForConv.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -225,8 +226,8 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, nanovdb::GridHandle handle; if (k % 2 == 1) { for (int p = 0; p < (k - 1) / 2; p += 1) { - nanovdb::tools::cuda::DilateGrid op(grid, - stream.stream()); + nanovdb::tools::cuda::DilateGrid op( + grid, stream.stream()); op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu index da1d8a3e3..56d5943fe 100644 --- a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu +++ b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -232,8 +233,8 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl nanovdb::GridHandle handle; if (k % 2 == 1) { for (int p = 0; p < (k - 1) / 2; p += 1) { - nanovdb::tools::cuda::DilateGrid op(grid, - stream.stream()); + nanovdb::tools::cuda::DilateGrid op( + grid, stream.stream()); op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); @@ -260,7 +261,8 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl // 2S (+) {0,1}^3, and one negative pad pass adds (+) {-1,0}^3, composing to (+) [-1,1]^3. if (stride == nanovdb::Coord(2) && uniformKernel && kernelSize[0] == 3) { return perItemGridHandle(baseGridHdl, guide, [&](nanovdb::OnIndexGrid *grid) { - nanovdb::tools::cuda::RefineGrid refineOp(grid, stream.stream()); + nanovdb::tools::cuda::RefineGrid refineOp( + grid, stream.stream()); refineOp.setChecksum(nanovdb::CheckMode::Default); refineOp.setVerbose(0); nanovdb::GridHandle refined = refineOp.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildGridFromIjk.cu b/src/fvdb/detail/ops/BuildGridFromIjk.cu index 0b4b393ff..e1fcae175 100644 --- a/src/fvdb/detail/ops/BuildGridFromIjk.cu +++ b/src/fvdb/detail/ops/BuildGridFromIjk.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -93,8 +94,9 @@ dispatchCreateNanoGridFromIJK(const JaggedTensor &ijk) { handles.push_back( nVoxels == 0 ? createEmptyGridHandle(guide.device()) - : nanovdb::tools::cuda::voxelsToGrid( - (nanovdb::Coord *)dataPtr, nVoxels, 1.0, guide)); + : nanovdb::tools::cuda:: + voxelsToGrid( + (nanovdb::Coord *)dataPtr, nVoxels, 1.0, guide)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/src/fvdb/detail/ops/BuildGridFromPoints.cu b/src/fvdb/detail/ops/BuildGridFromPoints.cu index 451981e7f..f2ff10ab4 100644 --- a/src/fvdb/detail/ops/BuildGridFromPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromPoints.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -196,17 +197,19 @@ dispatchBuildGridFromPoints(const JaggedTensor &points, } else if (pointsAreContiguous) { using PointPtrT = TransformedPointPtr; handles.push_back( - nanovdb::tools::cuda::voxelsToGrid( - PointPtrT(pointsPtr + 3 * startIdx, txs[i]), nPoints, 1.0, guide)); + nanovdb::tools::cuda:: + voxelsToGrid( + PointPtrT(pointsPtr + 3 * startIdx, txs[i]), nPoints, 1.0, guide)); } else { using PointPtrT = TransformedPointPtr; handles.push_back( - nanovdb::tools::cuda::voxelsToGrid( - PointPtrT( - pointsPtr + startIdx * rowStride, txs[i], rowStride, colStride), - nPoints, - 1.0, - guide)); + nanovdb::tools::cuda:: + voxelsToGrid( + PointPtrT( + pointsPtr + startIdx * rowStride, txs[i], rowStride, colStride), + nPoints, + 1.0, + guide)); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/src/fvdb/detail/ops/BuildMergedGrids.cu b/src/fvdb/detail/ops/BuildMergedGrids.cu index c04186e4f..2e00df666 100644 --- a/src/fvdb/detail/ops/BuildMergedGrids.cu +++ b/src/fvdb/detail/ops/BuildMergedGrids.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -45,7 +46,8 @@ dispatchMergeGrids(const GridBatchData &gridBatch1, const GridBatc nanovdb::OnIndexGrid *grid2 = gridBatch2.mGridHdl->deviceGrid(i); TORCH_CHECK(grid2, "Second Grid is null"); - nanovdb::tools::cuda::MergeGrids mergeOp(grid1, grid2, stream); + nanovdb::tools::cuda::MergeGrids mergeOp( + grid1, grid2, stream); mergeOp.setChecksum(nanovdb::CheckMode::Default); mergeOp.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildPaddedGrid.cu b/src/fvdb/detail/ops/BuildPaddedGrid.cu index e8508eb9d..b6a61ed0e 100644 --- a/src/fvdb/detail/ops/BuildPaddedGrid.cu +++ b/src/fvdb/detail/ops/BuildPaddedGrid.cu @@ -3,6 +3,7 @@ // #include #include +#include #include #include #include @@ -301,7 +302,8 @@ erodeOncePass(nanovdb::OnIndexGrid *grid, return createEmptyGridHandle(device); } - nanovdb::tools::cuda::PruneGrid pruneOp(grid, keepMasks, stream); + nanovdb::tools::cuda::PruneGrid pruneOp( + grid, keepMasks, stream); pruneOp.setChecksum(nanovdb::CheckMode::Default); pruneOp.setVerbose(0); auto handle = pruneOp.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildPrunedGrid.cu b/src/fvdb/detail/ops/BuildPrunedGrid.cu index b32f22f31..ecc9ae932 100644 --- a/src/fvdb/detail/ops/BuildPrunedGrid.cu +++ b/src/fvdb/detail/ops/BuildPrunedGrid.cu @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -73,7 +74,8 @@ dispatchPruneGrid(const GridBatchData &gridBatch, const JaggedTens maskI.data_ptr(), reinterpret_cast *>(maskBuffer.deviceData())); C10_CUDA_KERNEL_LAUNCH_CHECK(); - nanovdb::tools::cuda::PruneGrid pruneOp(grid, leafMask); + nanovdb::tools::cuda::PruneGrid pruneOp(grid, + leafMask); pruneOp.setChecksum(nanovdb::CheckMode::Default); pruneOp.setVerbose(0); From 2ef632b80f32add5622741c40fa01b53a5968f17 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Fri, 7 Aug 2026 17:12:01 +1200 Subject: [PATCH 3/8] Format fix Signed-off-by: Jonathan Swartz --- src/fvdb/TorchResource.h | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/fvdb/TorchResource.h b/src/fvdb/TorchResource.h index baccf7425..e07c46138 100644 --- a/src/fvdb/TorchResource.h +++ b/src/fvdb/TorchResource.h @@ -51,10 +51,13 @@ struct TorchResource : nanovdb::cuda::SyncFromAsync { void * allocate_async(size_t bytes, size_t /*alignment*/, cudaStream_t stream) { if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { - const size_t cutoff = (env[0] == '2') ? 0 : (1ull << 18); // '2' = trace all, else >= 256 KiB + const size_t cutoff = + (env[0] == '2') ? 0 : (1ull << 18); // '2' = trace all, else >= 256 KiB if (bytes >= cutoff) { - std::fprintf(stderr, "[fvdb/nanovdb] TorchResource alloc %12zu bytes (%.3f MB)\n", - bytes, double(bytes) / 1e6); + std::fprintf(stderr, + "[fvdb/nanovdb] TorchResource alloc %12zu bytes (%.3f MB)\n", + bytes, + double(bytes) / 1e6); } } void *p = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(bytes, stream); From 8b0289acfdc5d347c71c7001bebd4fb2e3b049a1 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Tue, 18 Aug 2026 18:22:15 +1200 Subject: [PATCH 4/8] Update NanoVDB code to latest on master to pickup all needed memory resource work Signed-off-by: Jonathan Swartz --- src/cmake/get_nanovdb.cmake | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index ef203d8f8..f7eed24dd 100644 --- a/src/cmake/get_nanovdb.cmake +++ b/src/cmake/get_nanovdb.cmake @@ -1,16 +1,10 @@ # Copyright Contributors to the OpenVDB Project # SPDX-License-Identifier: Apache-2.0 -# TEMPORARY: pinned to the merged nanovdb injectable-memory-resource stack -# (AcademySoftwareFoundation/openvdb PRs #2268, #2269, #2270, #2272, #2273 on -# top of upstream master d980ad92; tracking issue #2232). fvdb relies on the -# ResourceT seams these PRs add to route builder scratch through PyTorch's -# caching allocator (see src/fvdb/TorchResource.h). Repoint at -# AcademySoftwareFoundation/openvdb once the stack merges upstream. CPMAddPackage( NAME nanovdb - GITHUB_REPOSITORY swahtz/openvdb - GIT_TAG 558bfb2ead7c993f980ba18e89fee4b69991fb39 + GITHUB_REPOSITORY AcademySoftwareFoundation/openvdb + GIT_TAG 63d1878f30563eb9cde07d8b433fe4868944de44 SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) From 6e040223d969dbe226e46a66ceb7630df833e2a8 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Tue, 18 Aug 2026 18:34:52 +1200 Subject: [PATCH 5/8] Route PadGrid and CUB scratch through PyTorch's caching allocator Two remaining device-scratch sites bypassed TorchResource and allocated from nanoVDB's separate cudaMallocAsync pool, partitioning VRAM against torch's. fvdb's own morphology::PadGrid drives nanoVDB's TopologyBuilder (internal mask buffers, countNodes CUB scratch, TempPool) but hardcoded the default DeviceResource. Add a ResourceT template parameter mirroring the upstream DilateGrid signature and forward it to TopologyBuilder, then pass TorchResource at all seven call sites. This removes a split where the odd-kernel branch of the conv builders used DilateGrid<..., TorchResource> while the even-kernel branch immediately below used PadGrid on the rival pool. BuildFineGridFromCoarse allocated CUB DeviceSegmentedReduce temp storage with a bare cudaMallocAsync; route it through TorchResource and check both cub calls, which were previously unchecked. The ResourceT default keeps PadGrid source-compatible. Raise the from_nearest_voxels_to_points peak-memory bound from 150 to 200 MiB. PadGrid scratch (~78 MiB at 2M points) is now visible to torch.cuda.max_memory_allocated() where it previously was not. Total device consumption is unchanged -- only the accounting moved -- and the guard against the old >300 MiB coordinate-materialization path is retained. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Jonathan Swartz --- .../detail/ops/BuildFineGridFromCoarse.cu | 47 ++++++++++++------- src/fvdb/detail/ops/BuildGridForConv.cu | 4 +- .../detail/ops/BuildGridForConvTranspose.cu | 6 +-- .../ops/BuildGridFromNearestVoxelsToPoints.cu | 3 +- src/fvdb/detail/ops/BuildPaddedGrid.cu | 3 +- src/fvdb/detail/utils/nanovdb/PadGrid.cuh | 40 ++++++++++------ tests/unit/test_basic_ops.py | 9 +++- 7 files changed, 71 insertions(+), 41 deletions(-) diff --git a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu index a3298a647..d96e2b89f 100644 --- a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu +++ b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu @@ -306,24 +306,35 @@ dispatchFineIJKForCoarseGrid(const GridBatchData &batchHdl, void *dTempStorage = nullptr; size_t tempStorageBytes = 0; - cub::DeviceSegmentedReduce::Sum(dTempStorage, - tempStorageBytes, - mask.value().jdata().const_data_ptr(), - maskCounts, - deviceNumSegments, - beginOffsets, - endOffsets, - stream); - cudaMallocAsync(&dTempStorage, tempStorageBytes, stream); - cub::DeviceSegmentedReduce::Sum(dTempStorage, - tempStorageBytes, - mask.value().jdata().const_data_ptr(), - maskCounts, - deviceNumSegments, - beginOffsets, - endOffsets, - stream); - cudaFreeAsync(dTempStorage, stream); + C10_CUDA_CHECK( + cub::DeviceSegmentedReduce::Sum(dTempStorage, + tempStorageBytes, + mask.value().jdata().const_data_ptr(), + maskCounts, + deviceNumSegments, + beginOffsets, + endOffsets, + stream)); + + // Route the CUB scratch through torch's caching allocator rather than + // cudaMallocAsync, so it shares torch's pool instead of partitioning VRAM against + // it (same rationale as the nanoVDB builders -- see fvdb/TorchResource.h). + auto &resource = nanovdb::cuda::default_resource(); + dTempStorage = + resource.allocate_async(tempStorageBytes, TorchResource::DEFAULT_ALIGNMENT, stream); + + C10_CUDA_CHECK( + cub::DeviceSegmentedReduce::Sum(dTempStorage, + tempStorageBytes, + mask.value().jdata().const_data_ptr(), + maskCounts, + deviceNumSegments, + beginOffsets, + endOffsets, + stream)); + + resource.deallocate_async( + dTempStorage, tempStorageBytes, TorchResource::DEFAULT_ALIGNMENT, stream); } for (const auto deviceId: c10::irange(c10::cuda::device_count())) { diff --git a/src/fvdb/detail/ops/BuildGridForConv.cu b/src/fvdb/detail/ops/BuildGridForConv.cu index 0cacc87be..1056d8822 100644 --- a/src/fvdb/detail/ops/BuildGridForConv.cu +++ b/src/fvdb/detail/ops/BuildGridForConv.cu @@ -435,7 +435,7 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, } } else { for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/false, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -443,7 +443,7 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, grid = handle.deviceGrid(); } for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu index fdbd3e87b..b2f23ea8b 100644 --- a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu +++ b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu @@ -235,7 +235,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl } } else { for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/false, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -243,7 +243,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl grid = handle.deviceGrid(); } for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -268,7 +268,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl nanovdb::GridHandle refined = refineOp.getHandle(guide); C10_CUDA_KERNEL_LAUNCH_CHECK(); - morphology::PadGrid padOp( + morphology::PadGrid padOp( refined.deviceGrid(), /*positiveOctant=*/false, stream.stream()); diff --git a/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu index 066c66670..fdda2aac7 100644 --- a/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // #include +#include #include #include #include @@ -114,7 +115,7 @@ dispatchBuildGridFromNearestVoxelsToPoints( } nanovdb::OnIndexGrid *grid = baseHdl.deviceGrid(i); TORCH_CHECK(grid, "Grid is null"); - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handles.push_back(op.getHandle(guide)); diff --git a/src/fvdb/detail/ops/BuildPaddedGrid.cu b/src/fvdb/detail/ops/BuildPaddedGrid.cu index b6a61ed0e..2e908e014 100644 --- a/src/fvdb/detail/ops/BuildPaddedGrid.cu +++ b/src/fvdb/detail/ops/BuildPaddedGrid.cu @@ -249,7 +249,8 @@ padOncePass(nanovdb::OnIndexGrid *grid, bool positive, const TorchDeviceBuffer &guide, cudaStream_t stream) { - fvdb::detail::morphology::PadGrid op(grid, positive, stream); + fvdb::detail::morphology::PadGrid op( + grid, positive, stream); op.setChecksum(nanovdb::CheckMode::Default); auto handle = op.getHandle(guide); C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/src/fvdb/detail/utils/nanovdb/PadGrid.cuh b/src/fvdb/detail/utils/nanovdb/PadGrid.cuh index 709f0690d..04a4e574e 100644 --- a/src/fvdb/detail/utils/nanovdb/PadGrid.cuh +++ b/src/fvdb/detail/utils/nanovdb/PadGrid.cuh @@ -612,7 +612,10 @@ template struct ErodeKeepMaskFunctor { /// Modeled on `nanovdb::tools::cuda::DilateGrid`; the driver, root speculation and /// the TopologyBuilder pipeline are reused as-is, with the internal-node and /// leaf-node stages swapped for their one-sided (`Positive`-selected) variants. -template class PadGrid { +template class PadGrid { + static_assert(nanovdb::cuda::is_async_resource::value, + "PadGrid allocates stream-ordered scratch and requires an AsyncResource"); + using GridT = NanoGrid; using TreeT = NanoTree; using RootT = NanoRoot; @@ -622,8 +625,15 @@ template class PadGrid { /// @param d_srcGrid source device grid to be padded /// @param positiveOctant true -> pad by {0,1}^3, false -> pad by {-1,0}^3 /// @param stream optional CUDA stream - PadGrid(const GridT *d_srcGrid, bool positiveOctant, cudaStream_t stream = 0) - : mBuilder(stream), mStream(stream), mDeviceSrcGrid(d_srcGrid), mPositive(positiveOctant) {} + /// @param resource resource instance all device scratch is allocated from; + /// must outlive this operator (defaults to the per-type default + /// resource) + PadGrid(const GridT *d_srcGrid, + bool positiveOctant, + cudaStream_t stream = 0, + ResourceT &resource = nanovdb::cuda::default_resource()) + : mBuilder(stream, resource), mStream(stream), mDeviceSrcGrid(d_srcGrid), + mPositive(positiveOctant) {} void setChecksum(CheckMode mode = CheckMode::Disable) { @@ -639,17 +649,17 @@ template class PadGrid { void processGridTreeRoot(); void padLeafNodes(); - tools::cuda::TopologyBuilder mBuilder; + tools::cuda::TopologyBuilder mBuilder; cudaStream_t mStream{0}; const GridT *mDeviceSrcGrid; bool mPositive; TreeData mSrcTreeData; -}; +}; // morphology::PadGrid -template +template template GridHandle -PadGrid::getHandle(const BufferT &pool) { +PadGrid::getHandle(const BufferT &pool) { // Copy TreeData from GPU -> CPU cudaStreamSynchronize(mStream); mSrcTreeData = util::cuda::DeviceGridTraits::getTreeData(mDeviceSrcGrid); @@ -683,9 +693,9 @@ PadGrid::getHandle(const BufferT &pool) { return GridHandle(std::move(buffer)); } -template +template void -PadGrid::padRoot() { +PadGrid::padRoot() { // Conservatively and speculatively expands the root tile table to accommodate any new // root nodes introduced by the padding. This mirrors `DilateGrid::dilateRoot` verbatim // (a symmetric 26-connected speculation): although a one-sided pass only spills into @@ -755,9 +765,9 @@ PadGrid::padRoot() { mBuilder.mProcessedRoot.deviceUpload(device, mStream, false); } -template +template void -PadGrid::padInternalNodes() { +PadGrid::padInternalNodes() { if (mSrcTreeData.mNodeCount[1]) { // Unless it's an empty grid if (mPositive) { using Op = PadInternalNodesFunctor; @@ -783,9 +793,9 @@ PadGrid::padInternalNodes() { } } -template +template void -PadGrid::processGridTreeRoot() { +PadGrid::processGridTreeRoot() { // Copy GridData from source grid (duplicates grid name and map; others reset later) cudaCheck(cudaMemcpyAsync(&mBuilder.data()->getGrid(), mDeviceSrcGrid->data(), @@ -799,9 +809,9 @@ PadGrid::processGridTreeRoot() { cudaCheckError(); } -template +template void -PadGrid::padLeafNodes() { +PadGrid::padLeafNodes() { if (mBuilder.data()->nodeCount[1]) { // Unless output grid is empty if (mPositive) { using Op = PadLeafNodesFunctor; diff --git a/tests/unit/test_basic_ops.py b/tests/unit/test_basic_ops.py index 173d29935..3234ef37c 100644 --- a/tests/unit/test_basic_ops.py +++ b/tests/unit/test_basic_ops.py @@ -793,10 +793,17 @@ def test_nearest_voxels_to_points_peak_memory(self): peak_extra = torch.cuda.max_memory_allocated() - base # The old path peaked at > 300 MiB of torch tensors for 2M points (8N int32 coords + two # 8N int32 jidx arrays); the mask path allocates ~one N-coord list plus the output grid. + # + # The threshold accounts for nanoVDB builder scratch being torch-visible: PadGrid now + # routes its TopologyBuilder scratch through TorchResource (torch's caching allocator) + # rather than nanoVDB's separate cudaMallocAsync pool, so ~78 MiB that this measurement + # previously could not see is now counted here. Total device consumption is unchanged -- + # only the accounting moved -- so the bound is raised rather than the routing reverted. + # Measured ~159 MiB; 200 MiB keeps the guard against a return to the >300 MiB path. self.assertGreater(grid.total_voxels, 0) self.assertLess( peak_extra, - 150 * 1024 * 1024, + 200 * 1024 * 1024, f"from_nearest_voxels_to_points torch peak {peak_extra / 1024 / 1024:.1f} MiB too large", ) From f43839ad9c4dd3a23fea42bace2168b7fc4bb972 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Thu, 20 Aug 2026 10:34:52 +1200 Subject: [PATCH 6/8] Route the SaveNanoVDB CUDA path through PyTorch's caching allocator 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 Signed-off-by: Jonathan Swartz --- src/cmake/get_nanovdb.cmake | 2 +- src/fvdb/detail/io/SaveNanoVDB.cu | 22 +++++++++++++--------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index f7eed24dd..2dafcf56b 100644 --- a/src/cmake/get_nanovdb.cmake +++ b/src/cmake/get_nanovdb.cmake @@ -4,7 +4,7 @@ CPMAddPackage( NAME nanovdb GITHUB_REPOSITORY AcademySoftwareFoundation/openvdb - GIT_TAG 63d1878f30563eb9cde07d8b433fe4868944de44 + GIT_TAG 7946f17edb443fe46076a22ea933e52a23453c24 SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) diff --git a/src/fvdb/detail/io/SaveNanoVDB.cu b/src/fvdb/detail/io/SaveNanoVDB.cu index aee96ce0c..564b3bb42 100644 --- a/src/fvdb/detail/io/SaveNanoVDB.cu +++ b/src/fvdb/detail/io/SaveNanoVDB.cu @@ -1,6 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include +#include #include #include @@ -618,7 +620,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, } using HostGridHandle = nanovdb::GridHandle; - using DeviceGridHandle = nanovdb::GridHandle; + using DeviceGridHandle = nanovdb::GridHandle; using ValueT = typename nanovdb::BuildToValueMap::type; // Hoist tensor shape info out of the per-batch loop. The data tensor has shape @@ -647,7 +649,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, // Determine the device pointer to the source index grid buffer. CPU-resident grids normally // return through the host path above; the upload branch is kept as a defensive fallback if // this helper is reused without that dispatch. - nanovdb::cuda::DeviceBuffer tmpDevBuf; // empty unless we need to upload + TorchDeviceBuffer tmpDevBuf; // empty unless we need to upload const torch::Device gridDevice = gridBatchData.device(); const torch::Device cudaDevice = gridDevice.is_cuda() ? gridDevice @@ -662,7 +664,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, const uint64_t srcBufferSize = gridBatchData.nanoGridHandle().buffer().size(); const uint8_t *srcHostData = static_cast(gridBatchData.nanoGridHandle().buffer().data()); - tmpDevBuf = nanovdb::cuda::DeviceBuffer(srcBufferSize, cudaDevice.index(), stream.stream()); + tmpDevBuf = TorchDeviceBuffer(srcBufferSize, cudaDevice); cudaCheck(cudaMemcpyAsync(tmpDevBuf.deviceData(), srcHostData, srcBufferSize, @@ -685,7 +687,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, // on the same stream as the indexToGrid kernels so the GPU can run them back-to-back. std::vector deviceHandles; - std::vector perBatchValueBufs; + std::vector perBatchValueBufs; std::vector hostBuffers; std::vector origGridBytesPerBi; deviceHandles.reserve(gridBatchData.batchSize()); @@ -708,9 +710,8 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, dSrcBufferStart + gridBatchData.cumBytesAt(bi)); const uint64_t valueBufElems = static_cast(numVoxelsBi) + 1u; - nanovdb::cuda::DeviceBuffer valueBuf( - valueBufElems * sizeof(ValueT), cudaDevice.index(), stream.stream()); - ValueT *dValuesBufBase = static_cast(valueBuf.deviceData()); + TorchDeviceBuffer valueBuf(valueBufElems * sizeof(ValueT), cudaDevice); + ValueT *dValuesBufBase = reinterpret_cast(valueBuf.deviceData()); cudaCheck(cudaMemsetAsync(dValuesBufBase, 0, sizeof(ValueT), stream.stream())); if (numVoxelsBi > 0) { cudaCheck(cudaMemcpyAsync(dValuesBufBase + 1, @@ -720,8 +721,11 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, stream.stream())); } - DeviceGridHandle dh = nanovdb::tools::cuda::indexToGrid( - dSrcGrid, dValuesBufBase, nanovdb::cuda::DeviceBuffer(), stream.stream()); + // The guide buffer only communicates the target device; the output grid buffer and the + // builder's internal scratch both come from torch's caching allocator. + DeviceGridHandle dh = nanovdb::tools::cuda:: + indexToGrid( + dSrcGrid, dValuesBufBase, TorchDeviceBuffer(0, cudaDevice), stream.stream()); const uint64_t origGridBytes = dh.buffer().size(); const uint64_t totalBytes = origGridBytes + blindOverhead; From d87e85240f51399dcd8ff29543ee153be20f7a66 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Thu, 20 Aug 2026 17:56:42 +1200 Subject: [PATCH 7/8] Clarify TorchResource docs: dispatches to Torch's currently active CUDA allocator c10::cuda::CUDACachingAllocator::raw_alloc_with_stream / raw_delete are namespace-level dispatchers through CUDACachingAllocator::get(), so TorchResource follows whatever allocator the Torch runtime has installed: the native caching allocator, the cudaMallocAsync backend, or a CUDAPluggableAllocator installed via change_current_allocator. Reword the header comments to say so rather than implying a hardcoded binding to the native caching allocator. Co-Authored-By: Claude Fable 5 Signed-off-by: Jonathan Swartz --- src/fvdb/TorchResource.h | 40 ++++++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/src/fvdb/TorchResource.h b/src/fvdb/TorchResource.h index e07c46138..b8f811801 100644 --- a/src/fvdb/TorchResource.h +++ b/src/fvdb/TorchResource.h @@ -14,8 +14,18 @@ namespace fvdb { -/// @brief NanoVDB stream-ordered memory resource backed by PyTorch's CUDA caching -/// allocator (c10::cuda::CUDACachingAllocator). +/// @brief NanoVDB stream-ordered memory resource backed by PyTorch's currently +/// active CUDA allocator. +/// +/// c10::cuda::CUDACachingAllocator is a namespace, not a concrete +/// allocator: its free functions raw_alloc_with_stream / raw_delete +/// dispatch through CUDACachingAllocator::get(), the runtime-swappable +/// c10::cuda::CUDAAllocator* Torch itself allocates tensors from. This +/// resource 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 user-provided allocator installed via +/// torch.cuda.memory.change_current_allocator(CUDAPluggableAllocator(...)). /// /// Passed as the ResourceT template parameter of NanoVDB's CUDA builders /// (PointsToGrid / DilateGrid / MergeGrids / PruneGrid / RefineGrid / @@ -37,17 +47,21 @@ namespace fvdb { /// allocation). Useful for diagnosing topology-op memory blowup on large /// scenes. struct TorchResource : nanovdb::cuda::SyncFromAsync { - /// Alignment guaranteed by every allocation. The caching allocator returns - /// blocks aligned to at least 512 bytes, so advertising nanoVDB's - /// conventional 256 (matching cuda::DeviceResource) is always satisfied and - /// the alignment parameter below can be ignored. + /// Alignment guaranteed by every allocation. Torch's native caching + /// allocator returns blocks aligned to at least 512 bytes and the + /// cudaMallocAsync backend to at least 256, so advertising nanoVDB's + /// conventional 256 (matching cuda::DeviceResource) is satisfied and the + /// alignment parameter below can be ignored. A pluggable allocator wrapping + /// any cudaMalloc-family call satisfies 256 as well. static constexpr size_t DEFAULT_ALIGNMENT = 256; - /// @brief Stream-ordered allocation from torch's caching allocator. + /// @brief Stream-ordered allocation from torch's active CUDA allocator. /// @note raw_alloc_with_stream records @p stream against the block so torch /// defers reuse until work on it completes, matching the stream-ordered /// semantics of the cudaMallocAsync call it replaces. Allocation - /// happens on the current device, like cudaMallocAsync. + /// happens on the current device, like cudaMallocAsync. The call + /// dispatches to CUDACachingAllocator::get(), so a swapped-in backend + /// or pluggable allocator is honored. void * allocate_async(size_t bytes, size_t /*alignment*/, cudaStream_t stream) { if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { @@ -67,11 +81,13 @@ struct TorchResource : nanovdb::cuda::SyncFromAsync { return p; } - /// @brief Free through torch's caching allocator. + /// @brief Free through torch's active CUDA allocator. /// @note The stream argument is deliberately ignored: raw_delete relies on - /// the stream recorded at allocation time plus torch's per-stream event - /// tracking, so the free is safe without ordering on the caller's - /// stream. + /// the stream recorded at allocation time — the native backend's + /// per-stream event tracking, or the alloc-time stream Torch hands a + /// pluggable allocator's free function — so the free is safe without + /// ordering on the caller's stream. This is the same contract Torch's + /// own tensor frees rely on. void deallocate_async(void *p, size_t /*bytes*/, size_t /*alignment*/, cudaStream_t /*stream*/) { if (p == nullptr) { From 416fa43cc315f010d9068b07217cbcc7a196fb12 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Thu, 20 Aug 2026 18:13:59 +1200 Subject: [PATCH 8/8] Centralize the builder scratch allocator choice behind a BuilderResource alias MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback from @harrism: naming TorchResource at every leaf call site hard-wires the allocator policy the same way nanoVDB itself just un-hard-wired it. All 22 ResourceT bindings across the 13 ops files now name fvdb::BuilderResource, a one-line alias (BuilderResource.h) currently set to TorchResource. A build that must run the builders without torch (e.g. the ONNX Runtime EP planned in #579) retargets the alias in one place instead of touching every op. The seam is compile-time and relies on the resource being stateless; a stateful resource (per-session allocator handle) would additionally need an instance plumbed through the call sites — noted in the header. All 13 affected translation units compile cleanly against the branch's nanovdb pin. Co-Authored-By: Claude Fable 5 Signed-off-by: Jonathan Swartz --- src/fvdb/BuilderResource.h | 38 +++++++++++++++++++ src/fvdb/TorchResource.h | 4 +- src/fvdb/detail/io/SaveNanoVDB.cu | 4 +- .../detail/ops/BuildCoarseGridFromFine.cu | 4 +- src/fvdb/detail/ops/BuildDenseGrid.cu | 4 +- src/fvdb/detail/ops/BuildDilatedGrid.cu | 4 +- .../detail/ops/BuildFineGridFromCoarse.cu | 16 ++++---- src/fvdb/detail/ops/BuildGridForConv.cu | 8 ++-- .../detail/ops/BuildGridForConvTranspose.cu | 12 +++--- src/fvdb/detail/ops/BuildGridFromIjk.cu | 4 +- .../ops/BuildGridFromNearestVoxelsToPoints.cu | 4 +- src/fvdb/detail/ops/BuildGridFromPoints.cu | 6 +-- src/fvdb/detail/ops/BuildMergedGrids.cu | 4 +- src/fvdb/detail/ops/BuildPaddedGrid.cu | 6 +-- src/fvdb/detail/ops/BuildPrunedGrid.cu | 6 +-- 15 files changed, 82 insertions(+), 42 deletions(-) create mode 100644 src/fvdb/BuilderResource.h diff --git a/src/fvdb/BuilderResource.h b/src/fvdb/BuilderResource.h new file mode 100644 index 000000000..388914011 --- /dev/null +++ b/src/fvdb/BuilderResource.h @@ -0,0 +1,38 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#ifndef FVDB_BUILDERRESOURCE_H +#define FVDB_BUILDERRESOURCE_H + +#include + +namespace fvdb { + +/// @brief The memory resource fvdb's ops bind as the ResourceT template +/// parameter of nanoVDB's CUDA builders (and of fvdb's own PadGrid), +/// routing their internal device scratch. +/// +/// This alias is the single seam choosing that policy: call sites name +/// BuilderResource, never a concrete resource type. Today it is +/// TorchResource, which allocates from PyTorch's currently active CUDA +/// allocator (see TorchResource.h). A build that must run these +/// builders without torch (e.g. an ONNX Runtime execution provider, +/// where c10 is unavailable) retargets the alias here — behind a +/// build-time switch guarding the TorchResource include — instead of +/// touching every op. +/// +/// The alias covers the builders' scratch only. Buffer allocations that +/// are torch tensors by design (TorchDeviceBuffer, the SaveNanoVDB +/// staging buffers) name their types directly. +/// +/// Note the seam is compile-time and relies on the resource being +/// stateless: builders bind the shared instance from +/// nanovdb::cuda::default_resource() through their +/// defaulted constructor arguments. A stateful resource (e.g. one +/// holding a per-session allocator handle) additionally needs an +/// instance plumbed through the ops' call sites. +using BuilderResource = TorchResource; + +} // namespace fvdb + +#endif // FVDB_BUILDERRESOURCE_H diff --git a/src/fvdb/TorchResource.h b/src/fvdb/TorchResource.h index b8f811801..0f93a95c3 100644 --- a/src/fvdb/TorchResource.h +++ b/src/fvdb/TorchResource.h @@ -29,7 +29,9 @@ namespace fvdb { /// /// Passed as the ResourceT template parameter of NanoVDB's CUDA builders /// (PointsToGrid / DilateGrid / MergeGrids / PruneGrid / RefineGrid / -/// CoarsenGrid), it routes their internal device scratch — O(N-points) sort +/// CoarsenGrid) — always via the fvdb::BuilderResource alias +/// (BuilderResource.h), never named directly at call sites — it routes +/// their internal device scratch — O(N-points) sort /// keys, CUB temp storage, topology mask buffers — through the same pool /// that fvdb / PyTorch tensors use. Without this, nanoVDB's default /// DeviceResource allocates from a second cudaMallocAsync pool that diff --git a/src/fvdb/detail/io/SaveNanoVDB.cu b/src/fvdb/detail/io/SaveNanoVDB.cu index 564b3bb42..0bda691e7 100644 --- a/src/fvdb/detail/io/SaveNanoVDB.cu +++ b/src/fvdb/detail/io/SaveNanoVDB.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include @@ -724,7 +724,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, // The guide buffer only communicates the target device; the output grid buffer and the // builder's internal scratch both come from torch's caching allocator. DeviceGridHandle dh = nanovdb::tools::cuda:: - indexToGrid( + indexToGrid( dSrcGrid, dValuesBufBase, TorchDeviceBuffer(0, cudaDevice), stream.stream()); const uint64_t origGridBytes = dh.buffer().size(); diff --git a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu index ee60ddaee..fe357e683 100644 --- a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu +++ b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -85,7 +85,7 @@ coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch, TORCH_CHECK(grid, "Grid is null"); nanovdb::GridHandle handle; for (int p = 0; p < nPasses; p += 1) { - nanovdb::tools::cuda::CoarsenGrid op( + nanovdb::tools::cuda::CoarsenGrid op( grid, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildDenseGrid.cu b/src/fvdb/detail/ops/BuildDenseGrid.cu index da98cdaa4..93fd7ccc9 100644 --- a/src/fvdb/detail/ops/BuildDenseGrid.cu +++ b/src/fvdb/detail/ops/BuildDenseGrid.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -145,7 +145,7 @@ dispatchCreateNanoGridFromDense(int64_t batchSize, } else if (i == 0) { handles.push_back( nanovdb::tools::cuda:: - voxelsToGrid( + voxelsToGrid( (nanovdb::Coord *)ijkData.data_ptr(), nVoxels, 1.0, guide)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { diff --git a/src/fvdb/detail/ops/BuildDilatedGrid.cu b/src/fvdb/detail/ops/BuildDilatedGrid.cu index 246cfbdbc..b5c4960c2 100644 --- a/src/fvdb/detail/ops/BuildDilatedGrid.cu +++ b/src/fvdb/detail/ops/BuildDilatedGrid.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -55,7 +55,7 @@ dispatchDilateGrid(const GridBatchData &gridBatch, TORCH_CHECK(grid, "Grid is null"); for (auto j = 0; j < dilationAmount[i]; j += 1) { - nanovdb::tools::cuda::DilateGrid dilateOp( + nanovdb::tools::cuda::DilateGrid dilateOp( grid, stream); dilateOp.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); dilateOp.setChecksum(nanovdb::CheckMode::Default); diff --git a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu index d96e2b89f..441eb1a42 100644 --- a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu +++ b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -316,12 +316,12 @@ dispatchFineIJKForCoarseGrid(const GridBatchData &batchHdl, endOffsets, stream)); - // Route the CUB scratch through torch's caching allocator rather than + // Route the CUB scratch through the builder resource rather than bare // cudaMallocAsync, so it shares torch's pool instead of partitioning VRAM against - // it (same rationale as the nanoVDB builders -- see fvdb/TorchResource.h). - auto &resource = nanovdb::cuda::default_resource(); - dTempStorage = - resource.allocate_async(tempStorageBytes, TorchResource::DEFAULT_ALIGNMENT, stream); + // it (same rationale as the nanoVDB builders -- see fvdb/BuilderResource.h). + auto &resource = nanovdb::cuda::default_resource(); + dTempStorage = resource.allocate_async( + tempStorageBytes, BuilderResource::DEFAULT_ALIGNMENT, stream); C10_CUDA_CHECK( cub::DeviceSegmentedReduce::Sum(dTempStorage, @@ -334,7 +334,7 @@ dispatchFineIJKForCoarseGrid(const GridBatchData &batchHdl, stream)); resource.deallocate_async( - dTempStorage, tempStorageBytes, TorchResource::DEFAULT_ALIGNMENT, stream); + dTempStorage, tempStorageBytes, BuilderResource::DEFAULT_ALIGNMENT, stream); } for (const auto deviceId: c10::irange(c10::cuda::device_count())) { @@ -436,7 +436,7 @@ fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl, TORCH_CHECK(grid, "Grid is null"); nanovdb::GridHandle handle; for (int p = 0; p < nPasses; p += 1) { - nanovdb::tools::cuda::RefineGrid op( + nanovdb::tools::cuda::RefineGrid op( grid, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); op.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildGridForConv.cu b/src/fvdb/detail/ops/BuildGridForConv.cu index 1056d8822..2963adf2e 100644 --- a/src/fvdb/detail/ops/BuildGridForConv.cu +++ b/src/fvdb/detail/ops/BuildGridForConv.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -424,7 +424,7 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, nanovdb::GridHandle handle; if (k % 2 == 1) { for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - nanovdb::tools::cuda::DilateGrid op( + nanovdb::tools::cuda::DilateGrid op( grid, stream.stream()); op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); op.setChecksum(nanovdb::CheckMode::Default); @@ -435,7 +435,7 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, } } else { for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/false, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -443,7 +443,7 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, grid = handle.deviceGrid(); } for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); diff --git a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu index b2f23ea8b..c0cb63d36 100644 --- a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu +++ b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -224,7 +224,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl nanovdb::GridHandle handle; if (k % 2 == 1) { for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - nanovdb::tools::cuda::DilateGrid op( + nanovdb::tools::cuda::DilateGrid op( grid, stream.stream()); op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); op.setChecksum(nanovdb::CheckMode::Default); @@ -235,7 +235,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl } } else { for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/false, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -243,7 +243,7 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl grid = handle.deviceGrid(); } for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handle = op.getHandle(guide); @@ -261,14 +261,14 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl if (geometry.stride() == nanovdb::Coord(2) && isUniformKernel(geometry) && geometry.kernelSize()[0] == 3) { return perItemGridHandle(baseGridHdl, guide, [&](nanovdb::OnIndexGrid *grid) { - nanovdb::tools::cuda::RefineGrid refineOp( + nanovdb::tools::cuda::RefineGrid refineOp( grid, stream.stream()); refineOp.setChecksum(nanovdb::CheckMode::Default); refineOp.setVerbose(0); nanovdb::GridHandle refined = refineOp.getHandle(guide); C10_CUDA_KERNEL_LAUNCH_CHECK(); - morphology::PadGrid padOp( + morphology::PadGrid padOp( refined.deviceGrid(), /*positiveOctant=*/false, stream.stream()); diff --git a/src/fvdb/detail/ops/BuildGridFromIjk.cu b/src/fvdb/detail/ops/BuildGridFromIjk.cu index e1fcae175..23bf90551 100644 --- a/src/fvdb/detail/ops/BuildGridFromIjk.cu +++ b/src/fvdb/detail/ops/BuildGridFromIjk.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -95,7 +95,7 @@ dispatchCreateNanoGridFromIJK(const JaggedTensor &ijk) { nVoxels == 0 ? createEmptyGridHandle(guide.device()) : nanovdb::tools::cuda:: - voxelsToGrid( + voxelsToGrid( (nanovdb::Coord *)dataPtr, nVoxels, 1.0, guide)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu index fdda2aac7..31d97b3a1 100644 --- a/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -115,7 +115,7 @@ dispatchBuildGridFromNearestVoxelsToPoints( } nanovdb::OnIndexGrid *grid = baseHdl.deviceGrid(i); TORCH_CHECK(grid, "Grid is null"); - morphology::PadGrid op( + morphology::PadGrid op( grid, /*positiveOctant=*/true, stream.stream()); op.setChecksum(nanovdb::CheckMode::Default); handles.push_back(op.getHandle(guide)); diff --git a/src/fvdb/detail/ops/BuildGridFromPoints.cu b/src/fvdb/detail/ops/BuildGridFromPoints.cu index f2ff10ab4..c84d7b370 100644 --- a/src/fvdb/detail/ops/BuildGridFromPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromPoints.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -198,13 +198,13 @@ dispatchBuildGridFromPoints(const JaggedTensor &points, using PointPtrT = TransformedPointPtr; handles.push_back( nanovdb::tools::cuda:: - voxelsToGrid( + voxelsToGrid( PointPtrT(pointsPtr + 3 * startIdx, txs[i]), nPoints, 1.0, guide)); } else { using PointPtrT = TransformedPointPtr; handles.push_back( nanovdb::tools::cuda:: - voxelsToGrid( + voxelsToGrid( PointPtrT( pointsPtr + startIdx * rowStride, txs[i], rowStride, colStride), nPoints, diff --git a/src/fvdb/detail/ops/BuildMergedGrids.cu b/src/fvdb/detail/ops/BuildMergedGrids.cu index 2e00df666..36f22ce89 100644 --- a/src/fvdb/detail/ops/BuildMergedGrids.cu +++ b/src/fvdb/detail/ops/BuildMergedGrids.cu @@ -1,8 +1,8 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include -#include #include #include #include @@ -46,7 +46,7 @@ dispatchMergeGrids(const GridBatchData &gridBatch1, const GridBatc nanovdb::OnIndexGrid *grid2 = gridBatch2.mGridHdl->deviceGrid(i); TORCH_CHECK(grid2, "Second Grid is null"); - nanovdb::tools::cuda::MergeGrids mergeOp( + nanovdb::tools::cuda::MergeGrids mergeOp( grid1, grid2, stream); mergeOp.setChecksum(nanovdb::CheckMode::Default); mergeOp.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildPaddedGrid.cu b/src/fvdb/detail/ops/BuildPaddedGrid.cu index 2e908e014..d4d41859d 100644 --- a/src/fvdb/detail/ops/BuildPaddedGrid.cu +++ b/src/fvdb/detail/ops/BuildPaddedGrid.cu @@ -1,9 +1,9 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include -#include #include #include #include @@ -249,7 +249,7 @@ padOncePass(nanovdb::OnIndexGrid *grid, bool positive, const TorchDeviceBuffer &guide, cudaStream_t stream) { - fvdb::detail::morphology::PadGrid op( + fvdb::detail::morphology::PadGrid op( grid, positive, stream); op.setChecksum(nanovdb::CheckMode::Default); auto handle = op.getHandle(guide); @@ -303,7 +303,7 @@ erodeOncePass(nanovdb::OnIndexGrid *grid, return createEmptyGridHandle(device); } - nanovdb::tools::cuda::PruneGrid pruneOp( + nanovdb::tools::cuda::PruneGrid pruneOp( grid, keepMasks, stream); pruneOp.setChecksum(nanovdb::CheckMode::Default); pruneOp.setVerbose(0); diff --git a/src/fvdb/detail/ops/BuildPrunedGrid.cu b/src/fvdb/detail/ops/BuildPrunedGrid.cu index ecc9ae932..c56527d33 100644 --- a/src/fvdb/detail/ops/BuildPrunedGrid.cu +++ b/src/fvdb/detail/ops/BuildPrunedGrid.cu @@ -1,10 +1,10 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include -#include #include #include #include @@ -74,8 +74,8 @@ dispatchPruneGrid(const GridBatchData &gridBatch, const JaggedTens maskI.data_ptr(), reinterpret_cast *>(maskBuffer.deviceData())); C10_CUDA_KERNEL_LAUNCH_CHECK(); - nanovdb::tools::cuda::PruneGrid pruneOp(grid, - leafMask); + nanovdb::tools::cuda::PruneGrid pruneOp(grid, + leafMask); pruneOp.setChecksum(nanovdb::CheckMode::Default); pruneOp.setVerbose(0);