From 2d0eb61a6ee90634bff82452e21167b7bbfda869 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Mon, 24 Aug 2026 11:13:45 +1200 Subject: [PATCH] Advance NanoVDB pin past single-space GridHandle and widen Torch-pool coverage Bump the NanoVDB pin from 7946f17e to e679862f, the merge commit of openvdb #2288 (single-space device GridHandle>). The new hasDeviceSingle/hasHostSingle buffer traits are detected with false defaults, so TorchDeviceBuffer's dual-trait specialization is unaffected. On top of the bump, three allocator-coverage improvements: - ReinitializeSdf: the VoxelBlockManager's firstLeafID/jumpMap buffers now allocate through BuilderResource (torch's active CUDA allocator) via a local TorchVbmBuffer adapter, instead of DeviceBuffer's separate pool. The upstream VBM handle still requires the legacy dual-buffer static interface (hasDeviceDual gates + BufferT::create + a void*-returning deviceData()), so the single-space cuda::Buffer is wrapped rather than used directly; the adapter documents the upstream follow-up (openvdb #2232). - SaveNanoVDB: the device staging buffers (the defensive host-grid upload and the per-batch (N+1)-element value buffer) become nanovdb::cuda::Buffer over BuilderResource, stream-ordered on the save stream. The indexToGrid output handle stays TorchDeviceBuffer: the builder's output path still requires a dual-trait buffer upstream. - TorchDeviceBuffer: CUDA allocations can now be associated with an explicit stream (raw_alloc_with_stream), and create() forwards the stream nanovdb builders pass instead of discarding it. Previously the allocation was silently associated with the device's current torch stream, which is only correct when that coincides with the builder's stream. Verified: full rebuild (112 targets, sm_120, -Werror=all-warnings clean); tests/unit/test_sdf.py 10/10, test_io.py 622/622, test_basic_ops.py 276/276 (+1 skip) on an RTX PRO 6000 Blackwell. Co-Authored-By: Claude Fable 5 Signed-off-by: Jonathan Swartz --- src/cmake/get_nanovdb.cmake | 2 +- src/fvdb/BuilderResource.h | 7 ++- src/fvdb/TorchDeviceBuffer.cpp | 19 ++++-- src/fvdb/TorchDeviceBuffer.h | 9 ++- src/fvdb/detail/io/SaveNanoVDB.cu | 24 ++++---- src/fvdb/detail/ops/ReinitializeSdf.cu | 84 +++++++++++++++++++++++++- 6 files changed, 122 insertions(+), 23 deletions(-) diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index 2dafcf56b..7ffbde4f9 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 7946f17edb443fe46076a22ea933e52a23453c24 + GIT_TAG e679862fca22e844e4254ecdf1e2183e1d944a05 SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) diff --git a/src/fvdb/BuilderResource.h b/src/fvdb/BuilderResource.h index 388914011..35739f9f0 100644 --- a/src/fvdb/BuilderResource.h +++ b/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 diff --git a/src/fvdb/TorchDeviceBuffer.cpp b/src/fvdb/TorchDeviceBuffer.cpp index 9115c8179..147e995b8 100644 --- a/src/fvdb/TorchDeviceBuffer.cpp +++ b/src/fvdb/TorchDeviceBuffer.cpp @@ -60,7 +60,8 @@ GridHandle::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(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(c10::cuda::CUDACachingAllocator::raw_alloc(size)); + mData = reinterpret_cast( + stream ? c10::cuda::CUDACachingAllocator::raw_alloc_with_stream( + size, static_cast(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"); } diff --git a/src/fvdb/TorchDeviceBuffer.h b/src/fvdb/TorchDeviceBuffer.h index 7187a621e..dd8ec0492 100644 --- a/src/fvdb/TorchDeviceBuffer.h +++ b/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; diff --git a/src/fvdb/detail/io/SaveNanoVDB.cu b/src/fvdb/detail/io/SaveNanoVDB.cu index 0bda691e7..0f1920102 100644 --- a/src/fvdb/detail/io/SaveNanoVDB.cu +++ b/src/fvdb/detail/io/SaveNanoVDB.cu @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -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 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(gridBatchData.nanoGridHandle().buffer().data()); - tmpDevBuf = TorchDeviceBuffer(srcBufferSize, cudaDevice); - cudaCheck(cudaMemcpyAsync(tmpDevBuf.deviceData(), - srcHostData, - srcBufferSize, - cudaMemcpyHostToDevice, - stream.stream())); - dSrcBufferStart = static_cast(tmpDevBuf.deviceData()); + tmpDevBuf = nanovdb::cuda::Buffer( + stream.stream(), srcBufferSize, nanovdb::cuda::noInit); + cudaCheck(cudaMemcpyAsync( + tmpDevBuf.data(), srcHostData, srcBufferSize, cudaMemcpyHostToDevice, stream.stream())); + dSrcBufferStart = reinterpret_cast(tmpDevBuf.data()); } const ValueT *dDataValuesBase = reinterpret_cast(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; std::vector deviceHandles; - std::vector perBatchValueBufs; + std::vector perBatchValueBufs; std::vector hostBuffers; std::vector origGridBytesPerBi; deviceHandles.reserve(gridBatchData.batchSize()); @@ -710,8 +712,8 @@ fvdbToNanovdbGridWithValues(const GridBatchData &gridBatchData, dSrcBufferStart + gridBatchData.cumBytesAt(bi)); const uint64_t valueBufElems = static_cast(numVoxelsBi) + 1u; - TorchDeviceBuffer valueBuf(valueBufElems * sizeof(ValueT), cudaDevice); - ValueT *dValuesBufBase = reinterpret_cast(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, diff --git a/src/fvdb/detail/ops/ReinitializeSdf.cu b/src/fvdb/detail/ops/ReinitializeSdf.cu index 2780efecd..df9ed0e42 100644 --- a/src/fvdb/detail/ops/ReinitializeSdf.cu +++ b/src/fvdb/detail/ops/ReinitializeSdf.cu @@ -1,11 +1,12 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include -#include +#include #include #include #include @@ -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::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 mBuffer; + + public: + TorchVbmBuffer() = default; + TorchVbmBuffer(TorchVbmBuffer &&) = default; + TorchVbmBuffer &operator=(TorchVbmBuffer &&) = default; + + void * + deviceData() const { + return const_cast(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(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 { + static const bool hasDeviceDual = true; +}; +} // namespace nanovdb + +namespace fvdb { +namespace detail { +namespace ops { + namespace { using OnIndexGridT = nanovdb::NanoGrid; -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;