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

src/cmake/get_nanovdb.cmake

@@ -4,7 +4,7 @@
 CPMAddPackage(
     NAME nanovdb
     GITHUB_REPOSITORY AcademySoftwareFoundation/openvdb
-    GIT_TAG e538a0646b14125a043f623f205fcf218c5070a0
+    GIT_TAG 7946f17edb443fe46076a22ea933e52a23453c24
     SOURCE_SUBDIR nanovdb/nanovdb
     DOWNLOAD_ONLY YES
 )

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 <fvdb/TorchResource.h>
+
+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<BuilderResource>() 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

src/fvdb/TorchResource.h

@@ -0,0 +1,107 @@
+// Copyright Contributors to the OpenVDB Project
+// SPDX-License-Identifier: Apache-2.0
+//
+#ifndef FVDB_TORCHRESOURCE_H
+#define FVDB_TORCHRESOURCE_H
+
+#include <nanovdb/cuda/DeviceResource.h>
+
+#include <c10/cuda/CUDACachingAllocator.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <stdexcept>
+
+namespace fvdb {
+
+/// @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 /
+///        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
+///        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<TorchResource>() — 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<TorchResource> {
+    /// 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 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. 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")) {
+            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 active CUDA allocator.
+    /// @note The stream argument is deliberately ignored: raw_delete relies on
+    ///       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) {
+            return;
+        }
+        c10::cuda::CUDACachingAllocator::raw_delete(p);
+    }
+};
+
+static_assert(nanovdb::cuda::is_async_resource<TorchResource>::value,
+              "TorchResource must model nanoVDB's stream-ordered AsyncResource concept");
+
+} // namespace fvdb
+
+#endif // FVDB_TORCHRESOURCE_H

src/fvdb/detail/io/SaveNanoVDB.cu

@@ -1,6 +1,8 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
+#include <fvdb/TorchDeviceBuffer.h>
 #include <fvdb/detail/io/SaveNanoVDB.h>
 #include <fvdb/detail/utils/Utils.h>
 
@@ -618,7 +620,7 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData,
     }
 
     using HostGridHandle   = nanovdb::GridHandle<nanovdb::HostBuffer>;
-    using DeviceGridHandle = nanovdb::GridHandle<nanovdb::cuda::DeviceBuffer>;
+    using DeviceGridHandle = nanovdb::GridHandle<TorchDeviceBuffer>;
     using ValueT           = typename nanovdb::BuildToValueMap<OutBuildT>::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<const uint8_t *>(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<DeviceGridHandle> deviceHandles;
-    std::vector<nanovdb::cuda::DeviceBuffer> perBatchValueBufs;
+    std::vector<TorchDeviceBuffer> perBatchValueBufs;
     std::vector<nanovdb::HostBuffer> hostBuffers;
     std::vector<uint64_t> origGridBytesPerBi;
     deviceHandles.reserve(gridBatchData.batchSize());
@@ -708,9 +710,8 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData,
             dSrcBufferStart + gridBatchData.cumBytesAt(bi));
 
         const uint64_t valueBufElems = static_cast<uint64_t>(numVoxelsBi) + 1u;
-        nanovdb::cuda::DeviceBuffer valueBuf(
-            valueBufElems * sizeof(ValueT), cudaDevice.index(), stream.stream());
-        ValueT *dValuesBufBase = static_cast<ValueT *>(valueBuf.deviceData());
+        TorchDeviceBuffer valueBuf(valueBufElems * sizeof(ValueT), cudaDevice);
+        ValueT *dValuesBufBase = reinterpret_cast<ValueT *>(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<OutBuildT>(
-            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<OutBuildT, nanovdb::ValueOnIndex, TorchDeviceBuffer, BuilderResource>(
+                dSrcGrid, dValuesBufBase, TorchDeviceBuffer(0, cudaDevice), stream.stream());
 
         const uint64_t origGridBytes = dh.buffer().size();
         const uint64_t totalBytes    = origGridBytes + blindOverhead;

src/fvdb/detail/ops/BuildCoarseGridFromFine.cu

@@ -1,6 +1,7 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
 #include <fvdb/GridBatchData.h>
 #include <fvdb/detail/GridBatchDataFactory.h>
 #include <fvdb/detail/ops/BuildCoarseGridFromFine.h>
@@ -84,7 +85,8 @@ coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch,
         TORCH_CHECK(grid, "Grid is null");
         nanovdb::GridHandle<TorchDeviceBuffer> handle;
         for (int p = 0; p < nPasses; p += 1) {
-            nanovdb::tools::cuda::CoarsenGrid<nanovdb::ValueOnIndex> op(grid, stream.stream());
+            nanovdb::tools::cuda::CoarsenGrid<nanovdb::ValueOnIndex, BuilderResource> op(
+                grid, stream.stream());
             op.setChecksum(nanovdb::CheckMode::Default);
             op.setVerbose(0);
             handle = op.getHandle(guide);

src/fvdb/detail/ops/BuildDenseGrid.cu

@@ -1,6 +1,7 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
 #include <fvdb/GridBatchData.h>
 #include <fvdb/detail/GridBatchDataFactory.h>
 #include <fvdb/detail/ops/BuildDenseGrid.h>
@@ -143,8 +144,9 @@ dispatchCreateNanoGridFromDense<torch::kCUDA>(int64_t batchSize,
             handles.push_back(createEmptyGridHandle(guide.device()));
         } else if (i == 0) {
             handles.push_back(
-                nanovdb::tools::cuda::voxelsToGrid<GridT, nanovdb::Coord *, TorchDeviceBuffer>(
-                    (nanovdb::Coord *)ijkData.data_ptr(), nVoxels, 1.0, guide));
+                nanovdb::tools::cuda::
+                    voxelsToGrid<GridT, nanovdb::Coord *, TorchDeviceBuffer, BuilderResource>(
+                        (nanovdb::Coord *)ijkData.data_ptr(), nVoxels, 1.0, guide));
             C10_CUDA_KERNEL_LAUNCH_CHECK();
         } else {
             handles.push_back(handles[0].copy(guide));

src/fvdb/detail/ops/BuildDilatedGrid.cu

@@ -1,6 +1,7 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
 #include <fvdb/TorchDeviceBuffer.h>
 #include <fvdb/detail/GridBatchDataFactory.h>
 #include <fvdb/detail/ops/BuildDilatedGrid.h>
@@ -54,7 +55,8 @@ dispatchDilateGrid<torch::kCUDA>(const GridBatchData &gridBatch,
             TORCH_CHECK(grid, "Grid is null");
 
             for (auto j = 0; j < dilationAmount[i]; j += 1) {
-                nanovdb::tools::cuda::DilateGrid<nanovdb::ValueOnIndex> dilateOp(grid, stream);
+                nanovdb::tools::cuda::DilateGrid<nanovdb::ValueOnIndex, BuilderResource> dilateOp(
+                    grid, stream);
                 dilateOp.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX);
                 dilateOp.setChecksum(nanovdb::CheckMode::Default);
                 dilateOp.setVerbose(0);

src/fvdb/detail/ops/BuildFineGridFromCoarse.cu

@@ -1,6 +1,7 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
 #include <fvdb/GridBatchData.h>
 #include <fvdb/detail/GridBatchDataFactory.h>
 #include <fvdb/detail/ops/BuildFineGridFromCoarse.h>
@@ -305,24 +306,35 @@ dispatchFineIJKForCoarseGrid<torch::kPrivateUse1>(const GridBatchData &batchHdl,
 
             void *dTempStorage      = nullptr;
             size_t tempStorageBytes = 0;
-            cub::DeviceSegmentedReduce::Sum(dTempStorage,
-                                            tempStorageBytes,
-                                            mask.value().jdata().const_data_ptr<bool>(),
-                                            maskCounts,
-                                            deviceNumSegments,
-                                            beginOffsets,
-                                            endOffsets,
-                                            stream);
-            cudaMallocAsync(&dTempStorage, tempStorageBytes, stream);
-            cub::DeviceSegmentedReduce::Sum(dTempStorage,
-                                            tempStorageBytes,
-                                            mask.value().jdata().const_data_ptr<bool>(),
-                                            maskCounts,
-                                            deviceNumSegments,
-                                            beginOffsets,
-                                            endOffsets,
-                                            stream);
-            cudaFreeAsync(dTempStorage, stream);
+            C10_CUDA_CHECK(
+                cub::DeviceSegmentedReduce::Sum(dTempStorage,
+                                                tempStorageBytes,
+                                                mask.value().jdata().const_data_ptr<bool>(),
+                                                maskCounts,
+                                                deviceNumSegments,
+                                                beginOffsets,
+                                                endOffsets,
+                                                stream));
+
+            // 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/BuilderResource.h).
+            auto &resource = nanovdb::cuda::default_resource<BuilderResource>();
+            dTempStorage   = resource.allocate_async(
+                tempStorageBytes, BuilderResource::DEFAULT_ALIGNMENT, stream);
+
+            C10_CUDA_CHECK(
+                cub::DeviceSegmentedReduce::Sum(dTempStorage,
+                                                tempStorageBytes,
+                                                mask.value().jdata().const_data_ptr<bool>(),
+                                                maskCounts,
+                                                deviceNumSegments,
+                                                beginOffsets,
+                                                endOffsets,
+                                                stream));
+
+            resource.deallocate_async(
+                dTempStorage, tempStorageBytes, BuilderResource::DEFAULT_ALIGNMENT, stream);
         }
 
         for (const auto deviceId: c10::irange(c10::cuda::device_count())) {
@@ -424,7 +436,8 @@ fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl,
         TORCH_CHECK(grid, "Grid is null");
         nanovdb::GridHandle<TorchDeviceBuffer> handle;
         for (int p = 0; p < nPasses; p += 1) {
-            nanovdb::tools::cuda::RefineGrid<nanovdb::ValueOnIndex> op(grid, stream.stream());
+            nanovdb::tools::cuda::RefineGrid<nanovdb::ValueOnIndex, BuilderResource> op(
+                grid, stream.stream());
             op.setChecksum(nanovdb::CheckMode::Default);
             op.setVerbose(0);
             handle = op.getHandle(guide);