diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index a780e8e94..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 e538a0646b14125a043f623f205fcf218c5070a0 + GIT_TAG 7946f17edb443fe46076a22ea933e52a23453c24 SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) 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 new file mode 100644 index 000000000..0f93a95c3 --- /dev/null +++ b/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 + +#include + +#include +#include +#include + +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() — 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. 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::value, + "TorchResource must model nanoVDB's stream-ordered AsyncResource concept"); + +} // namespace fvdb + +#endif // FVDB_TORCHRESOURCE_H diff --git a/src/fvdb/detail/io/SaveNanoVDB.cu b/src/fvdb/detail/io/SaveNanoVDB.cu index aee96ce0c..0bda691e7 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; diff --git a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu index 8600f8d83..fe357e683 100644 --- a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu +++ b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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..93fd7ccc9 100644 --- a/src/fvdb/detail/ops/BuildDenseGrid.cu +++ b/src/fvdb/detail/ops/BuildDenseGrid.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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..b5c4960c2 100644 --- a/src/fvdb/detail/ops/BuildDilatedGrid.cu +++ b/src/fvdb/detail/ops/BuildDilatedGrid.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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..441eb1a42 100644 --- a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu +++ b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include @@ -305,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 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(); + dTempStorage = resource.allocate_async( + tempStorageBytes, BuilderResource::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, 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 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 27c427d6f..2963adf2e 100644 --- a/src/fvdb/detail/ops/BuildGridForConv.cu +++ b/src/fvdb/detail/ops/BuildGridForConv.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include @@ -423,8 +424,8 @@ 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(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); @@ -434,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); @@ -442,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 c633e5d56..c0cb63d36 100644 --- a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu +++ b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include @@ -223,8 +224,8 @@ 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(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); @@ -234,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); @@ -242,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); @@ -260,13 +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(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); 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 0b4b393ff..23bf90551 100644 --- a/src/fvdb/detail/ops/BuildGridFromIjk.cu +++ b/src/fvdb/detail/ops/BuildGridFromIjk.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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/BuildGridFromNearestVoxelsToPoints.cu b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu index 066c66670..31d97b3a1 100644 --- a/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromNearestVoxelsToPoints.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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/BuildGridFromPoints.cu b/src/fvdb/detail/ops/BuildGridFromPoints.cu index 451981e7f..c84d7b370 100644 --- a/src/fvdb/detail/ops/BuildGridFromPoints.cu +++ b/src/fvdb/detail/ops/BuildGridFromPoints.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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..36f22ce89 100644 --- a/src/fvdb/detail/ops/BuildMergedGrids.cu +++ b/src/fvdb/detail/ops/BuildMergedGrids.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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..d4d41859d 100644 --- a/src/fvdb/detail/ops/BuildPaddedGrid.cu +++ b/src/fvdb/detail/ops/BuildPaddedGrid.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include @@ -248,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(); @@ -301,7 +303,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..c56527d33 100644 --- a/src/fvdb/detail/ops/BuildPrunedGrid.cu +++ b/src/fvdb/detail/ops/BuildPrunedGrid.cu @@ -1,6 +1,7 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#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); 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", )