Update NanoVDB for single-space GridHandle and widen Torch-pool coverage

src/cmake/get_nanovdb.cmake

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

src/fvdb/BuilderResource.h

@@ -21,9 +21,10 @@ namespace fvdb {
 ///        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.
+///        The alias covers the builders' scratch and the device staging
+///        buffers feeding them (nanovdb::cuda::Buffer<..., BuilderResource> in
+///        SaveNanoVDB and ReinitializeSdf). Grid buffers that are torch-device
+///        aware by design (TorchDeviceBuffer) name their allocator directly.
 ///
 ///        Note the seam is compile-time and relies on the resource being
 ///        stateless: builders bind the shared instance from

src/fvdb/TorchDeviceBuffer.cpp

@@ -60,7 +60,8 @@ GridHandle<fvdb::TorchDeviceBuffer>::copy(const fvdb::TorchDeviceBuffer &guide)
 namespace fvdb {
 
 TorchDeviceBuffer::TorchDeviceBuffer(uint64_t size /* = 0*/,
-                                     const torch::Device &device /* = torch::kCPU*/)
+                                     const torch::Device &device /* = torch::kCPU*/,
+                                     void *stream /* = nullptr*/)
     : mSize(size), mData(nullptr), mDevice(device) {
     if (!mSize) {
         return;
@@ -70,9 +71,15 @@ TorchDeviceBuffer::TorchDeviceBuffer(uint64_t size /* = 0*/,
         // Initalize on the host
         mData = reinterpret_cast<uint8_t *>(malloc(size));
     } else if (mDevice.is_cuda()) {
-        // Initalize on the device
+        // Initalize on the device. With an explicit stream the allocation is associated with
+        // that stream in the caching allocator (the stream nanovdb builders order their work
+        // on); raw_alloc would silently associate it with the device's current torch stream
+        // instead, which is only correct when the two coincide.
         c10::cuda::CUDAGuard deviceGuard(mDevice);
-        mData = reinterpret_cast<uint8_t *>(c10::cuda::CUDACachingAllocator::raw_alloc(size));
+        mData = reinterpret_cast<uint8_t *>(
+            stream ? c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(
+                         size, static_cast<cudaStream_t>(stream))
+                   : c10::cuda::CUDACachingAllocator::raw_alloc(size));
         checkPtr(mData, "failed to allocate device data");
     } else if (mDevice.is_privateuseone()) {
         auto allocator = c10::GetAllocator(c10::DeviceType::PrivateUse1);
@@ -204,15 +211,17 @@ TorchDeviceBuffer::clear() {
 
 TorchDeviceBuffer
 TorchDeviceBuffer::create(uint64_t size, const TorchDeviceBuffer *proto, int device, void *stream) {
+    // The stream a nanovdb builder passes here is the one it orders its writes into the buffer
+    // on; forward it so the allocation is associated with that stream (see the constructor).
     if (proto) {
         // This is a hack to pass in the device index when creating grids from nanovdb. Since we
         // can't pass arguments through nanovdb creation functions, we use a prototype grid to pass
         // in the device index.
-        return TorchDeviceBuffer(size, proto->device());
+        return TorchDeviceBuffer(size, proto->device(), stream);
     } else if (device == cudaCpuDeviceId) {
         return TorchDeviceBuffer(size, torch::kCPU);
     } else if (device > cudaCpuDeviceId) {
-        return TorchDeviceBuffer(size, torch::Device(torch::kCUDA, device));
+        return TorchDeviceBuffer(size, torch::Device(torch::kCUDA, device), stream);
     } else {
         TORCH_CHECK(false, "Invalid parameters specified for TorchDeviceBuffer::create");
     }

src/fvdb/TorchDeviceBuffer.h

@@ -29,7 +29,14 @@ class TorchDeviceBuffer {
     /// nanovdb::HostBuffer
     /// @param size The size (in bytes to allocate for this buffer)
     /// @param device Specifies the device to use for the buffer
-    TorchDeviceBuffer(uint64_t size = 0, const torch::Device &device = torch::kCPU);
+    /// @param stream For a CUDA device, the stream the allocation is associated with in torch's
+    /// caching allocator (raw_alloc_with_stream): torch defers reusing the block until work queued
+    /// on that stream at free time completes, so pass the stream the buffer is used on. Null (the
+    /// default, indistinguishable from the legacy default stream) associates the allocation with
+    /// the device's current torch stream. Ignored for CPU and PrivateUse1 devices.
+    TorchDeviceBuffer(uint64_t size               = 0,
+                      const torch::Device &device = torch::kCPU,
+                      void *stream                = nullptr);
 
     /// @brief Disallow copy-construction
     TorchDeviceBuffer(const TorchDeviceBuffer &) = delete;

src/fvdb/detail/io/SaveNanoVDB.cu

@@ -7,6 +7,7 @@
 #include <fvdb/detail/utils/Utils.h>
 
 #include <nanovdb/NanoVDB.h>
+#include <nanovdb/cuda/Buffer.h>
 #include <nanovdb/cuda/DeviceBuffer.h>
 #include <nanovdb/io/IO.h>
 #include <nanovdb/tools/GridChecksum.h>
@@ -649,7 +650,9 @@ 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.
-    TorchDeviceBuffer tmpDevBuf; // empty unless we need to upload
+    // Device staging goes through the builders' resource (torch's caching allocator),
+    // stream-ordered on the same stream as the copies and kernels that use it.
+    nanovdb::cuda::Buffer<std::byte, BuilderResource> tmpDevBuf; // empty unless we need to upload
     const torch::Device gridDevice = gridBatchData.device();
     const torch::Device cudaDevice = gridDevice.is_cuda()
                                          ? gridDevice
@@ -664,13 +667,11 @@ 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 = TorchDeviceBuffer(srcBufferSize, cudaDevice);
-        cudaCheck(cudaMemcpyAsync(tmpDevBuf.deviceData(),
-                                  srcHostData,
-                                  srcBufferSize,
-                                  cudaMemcpyHostToDevice,
-                                  stream.stream()));
-        dSrcBufferStart = static_cast<const uint8_t *>(tmpDevBuf.deviceData());
+        tmpDevBuf = nanovdb::cuda::Buffer<std::byte, BuilderResource>(
+            stream.stream(), srcBufferSize, nanovdb::cuda::noInit);
+        cudaCheck(cudaMemcpyAsync(
+            tmpDevBuf.data(), srcHostData, srcBufferSize, cudaMemcpyHostToDevice, stream.stream()));
+        dSrcBufferStart = reinterpret_cast<const uint8_t *>(tmpDevBuf.data());
     }
 
     const ValueT *dDataValuesBase = reinterpret_cast<const ValueT *>(cudaData.jdata().data_ptr());
@@ -686,8 +687,9 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData,
     // to zero, and D2D-copy the data slice into [1..N]. All allocations and copies are queued
     // on the same stream as the indexToGrid kernels so the GPU can run them back-to-back.
 
+    using ValueStagingBuffer = nanovdb::cuda::Buffer<ValueT, BuilderResource>;
     std::vector<DeviceGridHandle> deviceHandles;
-    std::vector<TorchDeviceBuffer> perBatchValueBufs;
+    std::vector<ValueStagingBuffer> perBatchValueBufs;
     std::vector<nanovdb::HostBuffer> hostBuffers;
     std::vector<uint64_t> origGridBytesPerBi;
     deviceHandles.reserve(gridBatchData.batchSize());
@@ -710,8 +712,8 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData,
             dSrcBufferStart + gridBatchData.cumBytesAt(bi));
 
         const uint64_t valueBufElems = static_cast<uint64_t>(numVoxelsBi) + 1u;
-        TorchDeviceBuffer valueBuf(valueBufElems * sizeof(ValueT), cudaDevice);
-        ValueT *dValuesBufBase = reinterpret_cast<ValueT *>(valueBuf.deviceData());
+        ValueStagingBuffer valueBuf(stream.stream(), valueBufElems, nanovdb::cuda::noInit);
+        ValueT *dValuesBufBase = valueBuf.data();
         cudaCheck(cudaMemsetAsync(dValuesBufBase, 0, sizeof(ValueT), stream.stream()));
         if (numVoxelsBi > 0) {
             cudaCheck(cudaMemcpyAsync(dValuesBufBase + 1,

src/fvdb/detail/ops/ReinitializeSdf.cu

@@ -1,11 +1,12 @@
 // Copyright Contributors to the OpenVDB Project
 // SPDX-License-Identifier: Apache-2.0
 //
+#include <fvdb/BuilderResource.h>
 #include <fvdb/detail/ops/ReinitializeSdf.h>
 #include <fvdb/detail/utils/cuda/GridDim.h>
 
 #include <nanovdb/NanoVDB.h>
-#include <nanovdb/cuda/DeviceBuffer.h>
+#include <nanovdb/cuda/Buffer.h>
 #include <nanovdb/math/Math.h>
 #include <nanovdb/tools/VoxelBlockManager.h>
 #include <nanovdb/tools/cuda/VoxelBlockManager.cuh>
@@ -21,10 +22,89 @@ namespace fvdb {
 namespace detail {
 namespace ops {
 
+/// @brief Device-only buffer backing the VoxelBlockManagerHandle, allocating
+///        through the same resource as the nanovdb builders (BuilderResource,
+///        i.e. torch's active CUDA allocator) instead of DeviceBuffer's
+///        separate cudaMallocAsync pool.
+///
+///        VoxelBlockManagerHandle's device accessors are gated on
+///        BufferTraits<BufferT>::hasDeviceDual and static_cast their pointers
+///        from a void*-returning deviceData(), and buildVoxelBlockManager
+///        allocates via the legacy BufferT::create(bytes, guide, device,
+///        stream) static interface -- so the single-space nanovdb::cuda::Buffer
+///        cannot be plugged in directly until the VBM handle joins the
+///        single-space GridHandle work upstream (openvdb #2232/#2288). This
+///        wrapper adapts it: there is no host side, and the device allocation
+///        is stream-ordered on create()'s stream and freed on it.
+class TorchVbmBuffer {
+    nanovdb::cuda::Buffer<std::byte, BuilderResource> mBuffer;
+
+  public:
+    TorchVbmBuffer()                             = default;
+    TorchVbmBuffer(TorchVbmBuffer &&)            = default;
+    TorchVbmBuffer &operator=(TorchVbmBuffer &&) = default;
+
+    void *
+    deviceData() const {
+        return const_cast<std::byte *>(mBuffer.data());
+    }
+    void *
+    data() const {
+        return nullptr; // no host-side allocation
+    }
+    uint64_t
+    size() const {
+        return mBuffer.size_bytes();
+    }
+    bool
+    empty() const {
+        return mBuffer.empty();
+    }
+    bool
+    isEmpty() const {
+        return this->empty();
+    }
+    void
+    clear() {
+        mBuffer.destroy();
+    }
+
+    /// @brief Matches the signature buildVoxelBlockManager allocates through.
+    ///        The device argument is ignored: like cudaMallocAsync, the torch
+    ///        allocator allocates on the current device, which is what the
+    ///        caller passes anyway (it reads it back from cudaGetDevice).
+    static TorchVbmBuffer
+    create(uint64_t bytes,
+           const TorchVbmBuffer * /*guide*/ = nullptr,
+           int /*device*/                   = -1,
+           cudaStream_t stream              = 0) {
+        TorchVbmBuffer buf;
+        buf.mBuffer =
+            nanovdb::cuda::Buffer<std::byte, BuilderResource>(stream, bytes, nanovdb::cuda::noInit);
+        return buf;
+    }
+}; // TorchVbmBuffer
+
+} // namespace ops
+} // namespace detail
+} // namespace fvdb
+
+namespace nanovdb {
+// Device-only in practice (data() reports no host side); the dual trait is what
+// VoxelBlockManagerHandle's device accessors are gated on -- see TorchVbmBuffer.
+template <> struct BufferTraits<fvdb::detail::ops::TorchVbmBuffer> {
+    static const bool hasDeviceDual = true;
+};
+} // namespace nanovdb
+
+namespace fvdb {
+namespace detail {
+namespace ops {
+
 namespace {
 
 using OnIndexGridT = nanovdb::NanoGrid<nanovdb::ValueOnIndex>;
-using VbmBuffer    = nanovdb::cuda::DeviceBuffer;
+using VbmBuffer    = TorchVbmBuffer;
 
 // log2 of the VoxelBlockManager block width: each VBM block spans 2^9 = 512 active voxels.
 static constexpr int kLog2BlockWidth = 9;