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;