diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c12e46ea0..386c07387 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -61,6 +61,7 @@ set(FVDB_CPP_FILES set(FVDB_CU_FILES fvdb/GridBatchData.cu fvdb/detail/GridBatchDataFactory.cu + fvdb/detail/VbmCache.cu fvdb/detail/io/SaveNanoVDB.cu fvdb/detail/ops/ActiveGridCoords.cu fvdb/detail/ops/ActiveVoxelsInBoundsMask.cu diff --git a/src/benchmarks/CMakeLists.txt b/src/benchmarks/CMakeLists.txt index 6214629fe..ee11c9407 100644 --- a/src/benchmarks/CMakeLists.txt +++ b/src/benchmarks/CMakeLists.txt @@ -137,6 +137,16 @@ ConfigureDispatchBench(for_each_benchmark dispatch/for_each_benchmark.cu ) +# Active-voxel iteration benchmark — legacy leaf-scan vs cached-VBM decode across occupancies +ConfigureDispatchBench(active_voxel_iteration_benchmark + dispatch/active_voxel_iteration_benchmark.cu +) +# NanoVDB's host-side VoxelBlockManager uses `#pragma omp simd`, unknown to the host compiler +# when OpenMP is off (the fvdb library build suppresses it the same way). +target_compile_options(active_voxel_iteration_benchmark PRIVATE + $<$:-Xcompiler=-Wno-unknown-pragmas> + $<$:-Wno-unknown-pragmas>) + # GatherScatterDefault sparse convolution benchmark ConfigureDispatchBench(gather_scatter_conv_benchmark convolution/gather_scatter_conv_benchmark.cu diff --git a/src/benchmarks/dispatch/active_voxel_iteration_benchmark.cu b/src/benchmarks/dispatch/active_voxel_iteration_benchmark.cu new file mode 100644 index 000000000..2bda41ff4 --- /dev/null +++ b/src/benchmarks/dispatch/active_voxel_iteration_benchmark.cu @@ -0,0 +1,255 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +// Active-voxel iteration benchmarks: legacy leaf-scan (forEachVoxelCUDA, one thread per +// 512-slot leaf position regardless of occupancy) vs VBM-backed iteration +// (forEachActiveVoxelVbmCUDA, one thread per active voxel via the register-only +// VoxelBlockManager inverse-map decode). +// +// The workload is the ActiveGridCoords kernel body (decode + 3 int32 stores), the most +// launch-bound per-active-voxel op in fVDB, across leaf occupancies from dense (100%) to a +// sparse spherical shell (~a few % of leaf slots active) — occupancy is the leaf scan's cost +// driver and the VBM decode is occupancy-independent. +// +// A separate benchmark measures the one-time VbmCache build cost (paid once per grid +// lifetime; the iteration benchmarks run against a warm cache). +// + +#ifdef __NVCC__ +#pragma nv_diag_suppress 177 +#endif + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include + +namespace { + +using namespace fvdb; + +// ============================================================================ +// Grid construction at controlled leaf occupancy +// ============================================================================ + +// Dense cube: every voxel of a dim^3 box is active (occupancy ~100%). +torch::Tensor +denseCubeIjk(int dim) { + auto r = torch::arange(dim, torch::kInt32); + auto grid = torch::meshgrid({r, r, r}, "ij"); + return torch::stack({grid[0].reshape(-1), grid[1].reshape(-1), grid[2].reshape(-1)}, 1); +} + +// Random subsample of a dense cube: occupancy ~= fraction. +torch::Tensor +randomIjk(int dim, double fraction, uint64_t seed) { + auto cube = denseCubeIjk(dim); + auto gen = at::detail::createCPUGenerator(seed); + int64_t count = int64_t(double(cube.size(0)) * fraction); + auto perm = torch::randperm(cube.size(0), gen, torch::kInt64).slice(0, 0, count); + return cube.index_select(0, perm); +} + +// Spherical shell of ~thickness voxels: the narrow-band case, sparsest leaf occupancy. +torch::Tensor +shellIjk(int dim, double thickness) { + auto cube = denseCubeIjk(dim); + auto center = double(dim - 1) / 2.0; + auto radius = double(dim) * 0.4; + auto d = (cube.to(torch::kFloat64) - center).norm(2, 1) - radius; + return cube.index({d.abs() < thickness / 2.0}); +} + +c10::intrusive_ptr +makeGrid(const torch::Tensor &ijk) { + JaggedTensor jt(ijk.to(torch::Device(torch::kCUDA, 0))); + return fvdb::detail::ops::createNanoGridFromIJK(jt, {{1.0, 1.0, 1.0}}, {{0.0, 0.0, 0.0}}); +} + +double +leafOccupancy(const GridBatchData &grid) { + return double(grid.totalVoxels()) / (double(grid.totalLeaves()) * 512.0); +} + +// ============================================================================ +// Workload: the ActiveGridCoords kernel body (3 int32 stores per active voxel) +// ============================================================================ + +struct WriteCoordsFunctor { + int32_t *out; // [totalVoxels, 3] + + // VBM path entry point (forEachActiveVoxelVbmCUDA contract) + __device__ void + perActiveVoxel(nanovdb::Coord const &ijk, int64_t featureIdx) const { + out[featureIdx * 3 + 0] = ijk[0]; + out[featureIdx * 3 + 1] = ijk[1]; + out[featureIdx * 3 + 2] = ijk[2]; + } + + // Legacy leaf-scan entry point (forEachVoxelCUDA contract) + __device__ void + operator()(int64_t batchIdx, + int64_t leafIdx, + int64_t voxelIdx, + int64_t, + GridBatchData::Accessor acc) const { + auto const *grid = acc.grid(batchIdx); + auto const &leaf = grid->tree().template getFirstNode<0>()[leafIdx]; + if (leaf.isActive(voxelIdx)) { + auto const ijk = leaf.offsetToGlobalCoord(voxelIdx); + perActiveVoxel(ijk, acc.voxelOffset(batchIdx) + leaf.getValue(voxelIdx) - 1); + } + } +}; + +// ============================================================================ +// Benchmarks +// ============================================================================ + +enum class GridShape { Shell }; + +c10::intrusive_ptr +makeShapedGrid(GridShape shape) { + switch (shape) { + case GridShape::Shell: return makeGrid(shellIjk(256, 3.0)); + } + return nullptr; +} + +// Random subsample of a dense box at `percent`% occupancy (100 -> fully dense). Because the +// dense box covers whole leaves, leaf occupancy tracks the subsample fraction, letting the +// sweep locate the crossover point where the VBM path stops paying off. +c10::intrusive_ptr +makeSweepGrid(int percent) { + constexpr int kDim = 160; // ~4.1M-voxel dense box + if (percent >= 100) { + return makeGrid(denseCubeIjk(kDim)); + } + return makeGrid(randomIjk(kDim, double(percent) / 100.0, /*seed=*/percent)); +} + +// The leaf-scan baseline launches the header-defined forEachVoxelCUDAKernel directly (the +// forEachVoxelCUDA wrapper's optional ultra-sparse path references a kernel that is not +// exported from libfvdb, so it cannot be linked from a benchmark executable). +void +leafScanIteration(const GridBatchData &grid, WriteCoordsFunctor func) { + constexpr int kNumThreads = 1024; + const int64_t VOXELS_PER_LEAF = nanovdb::OnIndexTree::LeafNodeType::NUM_VALUES; + const int64_t numBlocks = + (grid.totalLeaves() * VOXELS_PER_LEAF + kNumThreads - 1) / kNumThreads; + fvdb::_private::forEachVoxelCUDAKernel + <<>>( + grid.deviceAccessor(), true, 1, func); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void +runIterationBenchmark(benchmark::State &state, c10::intrusive_ptr grid) { + auto out = torch::empty({grid->totalVoxels(), 3}, + torch::TensorOptions().dtype(torch::kInt32).device(grid->device())); + WriteCoordsFunctor func{out.data_ptr()}; + + // Warmup (also builds the VBM cache so the timed loop measures iteration only). + if constexpr (UseVbm) { + forEachActiveVoxelVbmCUDA(*grid, func); + } else { + leafScanIteration(*grid, func); + } + c10::cuda::getCurrentCUDAStream().synchronize(); + + for (auto _: state) { + if constexpr (UseVbm) { + forEachActiveVoxelVbmCUDA(*grid, func); + } else { + leafScanIteration(*grid, func); + } + c10::cuda::getCurrentCUDAStream().synchronize(); + } + state.counters["voxels"] = double(grid->totalVoxels()); + state.counters["leaf_occupancy"] = leafOccupancy(*grid); + state.SetItemsProcessed(state.iterations() * grid->totalVoxels()); +} + +template +void +BM_ActiveVoxelIteration(benchmark::State &state) { + if (!torch::cuda::is_available()) { + state.SkipWithError("CUDA not available"); + return; + } + runIterationBenchmark(state, makeShapedGrid(Shape)); +} + +// Occupancy sweep: state.range(0) is the subsample percentage of a dense box. +template +void +BM_ActiveVoxelIterationSweep(benchmark::State &state) { + if (!torch::cuda::is_available()) { + state.SkipWithError("CUDA not available"); + return; + } + runIterationBenchmark(state, makeSweepGrid(int(state.range(0)))); +} + +// One-time VBM build cost (a fresh cache every iteration); state.range(0) is the subsample +// percentage of a dense box. +void +BM_VbmCacheBuild(benchmark::State &state) { + if (!torch::cuda::is_available()) { + state.SkipWithError("CUDA not available"); + return; + } + auto grid = makeSweepGrid(int(state.range(0))); + for (auto _: state) { + fvdb::detail::VbmCache cache; + benchmark::DoNotOptimize(cache.get(*grid, 0)); + c10::cuda::getCurrentCUDAStream().synchronize(); + } + state.counters["voxels"] = double(grid->totalVoxels()); +} + +#define OCCUPANCY_SWEEP_ARGS \ + ->Arg(100) \ + ->Arg(95) \ + ->Arg(90) \ + ->Arg(85) \ + ->Arg(80) \ + ->Arg(70) \ + ->Arg(60) \ + ->Arg(50) \ + ->Arg(40) \ + ->Arg(30) \ + ->Arg(20) \ + ->Arg(10) \ + ->Arg(5) + +BENCHMARK_TEMPLATE(BM_ActiveVoxelIterationSweep, false) + ->Name("LeafScan/occupancy_pct") + ->Unit(benchmark::kMicrosecond) OCCUPANCY_SWEEP_ARGS; +BENCHMARK_TEMPLATE(BM_ActiveVoxelIterationSweep, true) + ->Name("Vbm/occupancy_pct") + ->Unit(benchmark::kMicrosecond) OCCUPANCY_SWEEP_ARGS; +BENCHMARK_TEMPLATE(BM_ActiveVoxelIteration, GridShape::Shell, false) + ->Name("LeafScan/shell") + ->Unit(benchmark::kMicrosecond); +BENCHMARK_TEMPLATE(BM_ActiveVoxelIteration, GridShape::Shell, true) + ->Name("Vbm/shell") + ->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_VbmCacheBuild) + ->Name("VbmBuild/occupancy_pct") + ->Unit(benchmark::kMicrosecond) + ->Arg(100) + ->Arg(50) + ->Arg(10); + +} // namespace diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index 2dafcf56b..1e96c692c 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 c3ee2009c0ab2801fb8eb2148e6375383fa5162e SOURCE_SUBDIR nanovdb/nanovdb DOWNLOAD_ONLY YES ) diff --git a/src/fvdb/GridBatchData.cu b/src/fvdb/GridBatchData.cu index 4e826e555..4ee3ba899 100644 --- a/src/fvdb/GridBatchData.cu +++ b/src/fvdb/GridBatchData.cu @@ -3,10 +3,31 @@ // #include #include +#include #include namespace fvdb { +GridBatchData::GridBatchData(std::shared_ptr> gridHdl, + GridMetadata *hostGridMetadata, + GridMetadata *deviceGridMetadata, + int64_t batchSize, + GridBatchMetadata batchMetadata, + torch::Tensor leafBatchIndices, + torch::Tensor batchOffsets, + torch::Tensor listIndices, + std::shared_ptr vbmCache) + : mHostGridMetadata(hostGridMetadata), mDeviceGridMetadata(deviceGridMetadata), + mBatchSize(batchSize), mBatchMetadata(std::move(batchMetadata)), mGridHdl(std::move(gridHdl)), + mLeafBatchIndices(std::move(leafBatchIndices)), mBatchOffsets(std::move(batchOffsets)), + mListIndices(std::move(listIndices)), + mVbmCache(vbmCache ? std::move(vbmCache) : std::make_shared()) {} + +detail::VbmCache & +GridBatchData::vbmCache() const { + return *mVbmCache; +} + // ----------------------------------------------------------------------- // Methods that dereference mGridHdl (require complete TorchDeviceBuffer) // ----------------------------------------------------------------------- diff --git a/src/fvdb/GridBatchData.h b/src/fvdb/GridBatchData.h index 9754031dc..ffb822024 100644 --- a/src/fvdb/GridBatchData.h +++ b/src/fvdb/GridBatchData.h @@ -14,6 +14,7 @@ #include #include +#include #include #if !defined(__CUDACC__) && !defined(__restrict__) @@ -22,6 +23,10 @@ namespace fvdb { +namespace detail { +class VbmCache; +} // namespace detail + struct GridBatchData : public torch::CustomClassHolder { static constexpr int64_t MAX_GRIDS_PER_BATCH = 1024; // Maximum number of grids in a batch @@ -89,10 +94,20 @@ struct GridBatchData : public torch::CustomClassHolder { torch::Tensor mBatchOffsets; // Batch indices for grid torch::Tensor mListIndices; // List indices for grid (same as JaggedTensor) + // Lazily-built per-grid VoxelBlockManager handles (pure derived state; never serialized). + // Shared with sliced/indexed views of this batch, which alias the same grid buffer. Grid + // topology is immutable after construction, so cached entries never need invalidation. + std::shared_ptr mVbmCache; + // ----------------------------------------------------------------------- // Single constructor: bundles pre-computed fields (takes ownership of // metadata pointers). All computation happens outside, in factory - // functions, before this constructor is called. + // functions, before this constructor is called. Defined in GridBatchData.cu + // (requires the complete detail::VbmCache type). + // + // vbmCache is passed only by views that share this batch's grid buffer (so the + // parent's cached VBMs are reused); everyone else leaves it null and gets a + // fresh, empty cache. // ----------------------------------------------------------------------- GridBatchData(std::shared_ptr> gridHdl, GridMetadata *hostGridMetadata, @@ -101,11 +116,8 @@ struct GridBatchData : public torch::CustomClassHolder { GridBatchMetadata batchMetadata, torch::Tensor leafBatchIndices, torch::Tensor batchOffsets, - torch::Tensor listIndices) - : mHostGridMetadata(hostGridMetadata), mDeviceGridMetadata(deviceGridMetadata), - mBatchSize(batchSize), mBatchMetadata(std::move(batchMetadata)), - mGridHdl(std::move(gridHdl)), mLeafBatchIndices(std::move(leafBatchIndices)), - mBatchOffsets(std::move(batchOffsets)), mListIndices(std::move(listIndices)) {} + torch::Tensor listIndices, + std::shared_ptr vbmCache = nullptr); ~GridBatchData(); @@ -335,6 +347,10 @@ struct GridBatchData : public torch::CustomClassHolder { nanovdb::OnIndexGrid *deviceGridPtrAt(int64_t bi) const; nanovdb::OnIndexGrid *hostGridPtrAt(int64_t bi) const; + // The batch's lazily-built per-grid VoxelBlockManager cache (see detail/VbmCache.h). + // Defined in GridBatchData.cu (requires the complete detail::VbmCache type). + detail::VbmCache &vbmCache() const; + const VoxelCoordTransform & primalTransformAt(int64_t bi) const { bi = negativeToPositiveIndexWithRangecheck(bi); diff --git a/src/fvdb/detail/VbmCache.cu b/src/fvdb/detail/VbmCache.cu new file mode 100644 index 000000000..6ba66725e --- /dev/null +++ b/src/fvdb/detail/VbmCache.cu @@ -0,0 +1,89 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#include +#include + +#include + +#include +#include +#include + +namespace fvdb { +namespace detail { + +VbmCache::~VbmCache() { + for (auto &[key, entry]: mEntries) { + if (entry.builtOn) { + c10::cuda::CUDAGuard guard(entry.deviceIdx); + C10_CUDA_CHECK_WARN(cudaEventDestroy(entry.builtOn)); + } + } +} + +VbmCache::GridVbm +VbmCache::get(const GridBatchData &batch, int64_t bi) { + TORCH_CHECK(batch.device().is_cuda(), "VbmCache is only supported for CUDA grids"); + const uint64_t key = batch.cumBytesAt(bi); // view-stable identity into the shared buffer + + std::lock_guard lock(mMutex); + const auto stream = at::cuda::getCurrentCUDAStream(batch.device().index()); + + auto it = mEntries.find(key); + if (it != mEntries.end()) { + Entry &entry = it->second; + if (entry.builtOn && stream.stream() != entry.builtStream) { + // Execution ordering: this consumer stream must see the finished build. + C10_CUDA_CHECK(cudaStreamWaitEvent(stream.stream(), entry.builtOn, 0)); + // Lifetime ordering: register the consumer stream with the caching allocator so + // that destroying the owning GridBatchData while this stream's kernels are still + // in flight cannot recycle the buffers under them. + entry.firstLeafID.record_stream(stream.unwrap()); + entry.jumpMap.record_stream(stream.unwrap()); + } + return entry.view; + } + + Entry entry; + entry.deviceIdx = batch.device().index(); + const int64_t numVoxels = batch.numVoxelsAt(bi); + if (numVoxels > 0) { + c10::cuda::CUDAGuard guard(batch.device().index()); + const int64_t nBlocks = (numVoxels + kBlockWidth - 1) >> kLog2BlockWidth; + + // Tensor-backed buffers: the caching allocator ties frees to the allocation stream and + // to any record_stream()ed consumer streams, which no raw allocation would. The + // NanoVDB handle only sees non-owning views, and only for the duration of the build. + // Sizing the VBM from host metadata (firstOffset = 1, lastOffset = numVoxels) selects + // the in-place build overload, avoiding the blocking device read of activeVoxelCount + // that the allocating overload performs. + auto opts = torch::TensorOptions().device(batch.device()); + entry.firstLeafID = torch::empty({nBlocks}, opts.dtype(torch::kInt32)); + entry.jumpMap = torch::empty({nBlocks * kJumpMapWordCount}, opts.dtype(torch::kInt64)); + + nanovdb::tools::VoxelBlockManagerHandle handle( + VbmBufferView(entry.firstLeafID.data_ptr()), + VbmBufferView(entry.jumpMap.data_ptr()), + uint64_t(nBlocks), + /*firstOffset=*/1, + /*lastOffset=*/uint64_t(numVoxels)); + nanovdb::tools::cuda::buildVoxelBlockManager( + batch.deviceGridPtrAt(bi), handle, stream.stream()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + C10_CUDA_CHECK(cudaEventCreateWithFlags(&entry.builtOn, cudaEventDisableTiming)); + C10_CUDA_CHECK(cudaEventRecord(entry.builtOn, stream.stream())); + entry.builtStream = stream.stream(); + + entry.view = GridVbm{reinterpret_cast(entry.firstLeafID.data_ptr()), + reinterpret_cast(entry.jumpMap.data_ptr()), + uint32_t(nBlocks), + /*firstOffset=*/1, + /*lastOffset=*/uint64_t(numVoxels)}; + } + return mEntries.emplace(key, std::move(entry)).first->second.view; +} + +} // namespace detail +} // namespace fvdb diff --git a/src/fvdb/detail/VbmCache.h b/src/fvdb/detail/VbmCache.h new file mode 100644 index 000000000..20898fac0 --- /dev/null +++ b/src/fvdb/detail/VbmCache.h @@ -0,0 +1,136 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#ifndef FVDB_DETAIL_VBMCACHE_H +#define FVDB_DETAIL_VBMCACHE_H + +#include // for nanovdb::BufferTraits + +#include +#include + +#include + +#include +#include +#include + +namespace fvdb { + +struct GridBatchData; + +namespace detail { + +/// @brief Non-owning device-buffer view satisfying the NanoVDB buffer concept +/// (data()/deviceData()/clear() returning void*). Used only to run the in-place +/// buildVoxelBlockManager over memory whose ownership and stream-safety are managed by the +/// PyTorch caching allocator (torch::Tensor storage held by VbmCache::Entry). +class VbmBufferView { + void *mDevicePtr = nullptr; + + public: + VbmBufferView() = default; + explicit VbmBufferView(void *devicePtr) : mDevicePtr(devicePtr) {} + VbmBufferView(VbmBufferView &&) = default; + VbmBufferView &operator=(VbmBufferView &&) = default; + + void * + data() const { + return nullptr; // host side unused: the cache is CUDA-only + } + void * + deviceData() const { + return mDevicePtr; + } + void + clear() { + mDevicePtr = nullptr; // non-owning + } +}; + +} // namespace detail +} // namespace fvdb + +namespace nanovdb { +template <> struct BufferTraits { + static const bool hasDeviceDual = true; +}; +} // namespace nanovdb + +namespace fvdb { +namespace detail { + +/// @brief Lazily-built, per-grid NanoVDB VoxelBlockManager (VBM) handles for a GridBatchData. +/// +/// The VBM partitions a grid's active voxels into fixed-width blocks of BlockWidth sequential +/// value indices and stores, per block, the ID of the first leaf overlapping the block plus a +/// bitmask (jumpMap) of the in-block positions where subsequent leaves begin. Kernels can then +/// decode any active voxel's (leafIndex, voxelOffset) in registers via +/// nanovdb::tools::cuda::VoxelBlockManager::decodeInverseMap, giving occupancy-independent +/// per-active-voxel iteration (one decode per active voxel instead of one thread per 512-slot +/// leaf position). +/// +/// Grid topology is immutable after GridBatchData construction (every topology op returns a new +/// GridBatchData), so entries never need invalidation; the cache is pure derived state and must +/// never be serialized. It is shared between a GridBatchData and any sliced/indexed views of it +/// (which share the same underlying grid buffer); entries are keyed by the grid's byte offset +/// into that shared buffer, which is the only view-stable per-grid identity. +/// +/// Stream safety: entry buffers are torch::Tensor storage from the CUDA caching allocator. +/// get() makes a consumer stream that differs from the build stream wait on the recorded build +/// event (execution ordering) and record_stream()s the buffers for it (lifetime ordering), so +/// destroying the owning GridBatchData while a cross-stream consumer kernel is still in flight +/// cannot recycle the buffers under it. Callers must invoke get() on the same current stream +/// they subsequently launch consuming kernels on. +/// +/// CUDA-only: grids on other devices must use the legacy leaf-scan iteration paths. +class VbmCache { + public: + static constexpr int kLog2BlockWidth = 7; // 128 active voxels per VBM block + static constexpr int kBlockWidth = 1 << kLog2BlockWidth; + static constexpr int kJumpMapWordCount = kBlockWidth / 64; + + /// @brief POD view of one grid's VBM metadata. All pointers are device pointers valid for + /// the lifetime of the owning GridBatchData (and of any views sharing its grid buffer). + struct GridVbm { + const uint32_t *firstLeafID = nullptr; // [blockCount] + const uint64_t *jumpMap = nullptr; // [blockCount * kJumpMapWordCount] + uint32_t blockCount = 0; + uint64_t firstOffset = 0; // always 1 when built (value index 0 = background) + uint64_t lastOffset = 0; // == number of active voxels in the grid + }; + + VbmCache() = default; + ~VbmCache(); + + VbmCache(const VbmCache &) = delete; + VbmCache &operator=(const VbmCache &) = delete; + VbmCache(VbmCache &&) = delete; + VbmCache &operator=(VbmCache &&) = delete; + + /// @brief Return the VBM for logical grid @p bi of @p batch, building it on the current + /// CUDA stream of the batch's device on first access. Thread-safe. If a later call arrives + /// on a different stream than the one the entry was built on, that stream is made to wait + /// on the recorded build event and is registered with the caching allocator as a user of + /// the entry's buffers before this returns. + /// @return The grid's VBM view, or a zero GridVbm (blockCount == 0) for an empty grid. + GridVbm get(const GridBatchData &batch, int64_t bi); + + private: + struct Entry { + torch::Tensor firstLeafID; // int32 [blockCount] (holds uint32 values) + torch::Tensor jumpMap; // int64 [blockCount * kJumpMapWordCount] (holds uint64 bits) + GridVbm view; + cudaEvent_t builtOn = nullptr; // recorded on the build stream + cudaStream_t builtStream = nullptr; + c10::DeviceIndex deviceIdx = -1; + }; + + std::mutex mMutex; + std::unordered_map mEntries; // keyed by GridBatchData::cumBytesAt(bi) +}; + +} // namespace detail +} // namespace fvdb + +#endif // FVDB_DETAIL_VBMCACHE_H diff --git a/src/fvdb/detail/ops/IndexGrid.cu b/src/fvdb/detail/ops/IndexGrid.cu index 8f6ebda99..f065fd62a 100644 --- a/src/fvdb/detail/ops/IndexGrid.cu +++ b/src/fvdb/detail/ops/IndexGrid.cu @@ -96,6 +96,8 @@ indexGridInternal(const fvdb::GridBatchData &grid, const Indexable &idx, int64_t listIndices = grid.mListIndices; } + // The view shares the parent's grid buffer, so it also shares the parent's VBM cache + // (entries are keyed by byte offset into the shared buffer, which is view-stable). return c10::make_intrusive(grid.mGridHdl, hostMeta, deviceMeta, @@ -103,7 +105,8 @@ indexGridInternal(const fvdb::GridBatchData &grid, const Indexable &idx, int64_t std::move(batchMeta), std::move(leafBatchIndices), std::move(batchOffsets), - std::move(listIndices)); + std::move(listIndices), + grid.mVbmCache); } struct RangeAccessor { diff --git a/src/fvdb/detail/ops/ReinitializeSdf.cu b/src/fvdb/detail/ops/ReinitializeSdf.cu index 2780efecd..d2439c4e7 100644 --- a/src/fvdb/detail/ops/ReinitializeSdf.cu +++ b/src/fvdb/detail/ops/ReinitializeSdf.cu @@ -1,11 +1,11 @@ // Copyright Contributors to the OpenVDB Project // SPDX-License-Identifier: Apache-2.0 // +#include #include #include #include -#include #include #include #include @@ -24,40 +24,44 @@ namespace ops { namespace { using OnIndexGridT = nanovdb::NanoGrid; -using VbmBuffer = nanovdb::cuda::DeviceBuffer; // log2 of the VoxelBlockManager block width: each VBM block spans 2^9 = 512 active voxels. -static constexpr int kLog2BlockWidth = 9; +static constexpr int kLog2BlockWidth = VbmCache::kLog2BlockWidth; // ------------------------- VBM fused 6-face stencil preamble ------------------------- -// The VBM decode gives the centre coord / value-index for free; we then read just the 6 FACE -// neighbours through a cached ReadAccessor. Yields `centerIndex` (the centre voxel's value index) -// and `faceIndex[6]` (the 6 face-neighbour value indices, -x,+x,-y,+y,-z,+z; 0 = -// inactive/background). Kernels using it take the grid/firstLeafID/jumpMap/firstOffset parameters -// by these exact names. -#define VBM_FACES_BEGIN() \ - constexpr int blockWidth = 1 << kLog2BlockWidth, jumpMapWordCount = blockWidth / 64; \ - using VoxelBlockManagerT = nanovdb::tools::cuda::VoxelBlockManager; \ - __shared__ uint32_t sharedLeafIndex[blockWidth]; \ - __shared__ uint16_t sharedVoxelOffset[blockWidth]; \ - VoxelBlockManagerT::template decodeInverseMaps( \ - grid, \ - firstLeafID[blockIdx.x], \ - &jumpMap[uint64_t(blockIdx.x) * jumpMapWordCount], \ - firstOffset + uint64_t(blockIdx.x) * blockWidth, \ - sharedLeafIndex, \ - sharedVoxelOffset); \ - if (sharedLeafIndex[threadIdx.x] == VoxelBlockManagerT::UnusedLeafIndex) \ - return; \ - const auto &leaf = grid->tree().template getFirstNode<0>()[sharedLeafIndex[threadIdx.x]]; \ - const nanovdb::Coord centerCoord = leaf.offsetToGlobalCoord(sharedVoxelOffset[threadIdx.x]); \ - const uint64_t centerIndex = leaf.getValue(sharedVoxelOffset[threadIdx.x]); \ - auto accessor = grid->getAccessor(); \ - const uint64_t faceIndex[6] = {accessor.getValue(centerCoord.offsetBy(-1, 0, 0)), \ - accessor.getValue(centerCoord.offsetBy(1, 0, 0)), \ - accessor.getValue(centerCoord.offsetBy(0, -1, 0)), \ - accessor.getValue(centerCoord.offsetBy(0, 1, 0)), \ - accessor.getValue(centerCoord.offsetBy(0, 0, -1)), \ +// The register-only VBM decode gives the centre coord / value-index for free (no shared memory, +// no barrier); we then read just the 6 FACE neighbours through a cached ReadAccessor. Yields +// `centerIndex` (the centre voxel's value index) and `faceIndex[6]` (the 6 face-neighbour value +// indices, -x,+x,-y,+y,-z,+z; 0 = inactive/background). Kernels using it take the +// grid/firstLeafID/jumpMap/firstOffset/lastOffset parameters by these exact names and launch as +// <<>> (blockWidth = the VBM block width, 1 << kLog2BlockWidth). +#define VBM_FACES_BEGIN() \ + constexpr int blockWidth = 1 << kLog2BlockWidth, jumpMapWordCount = blockWidth / 64; \ + using VoxelBlockManagerT = nanovdb::tools::cuda::VoxelBlockManager; \ + const uint64_t blockFirstOffset = firstOffset + uint64_t(blockIdx.x) * blockWidth; \ + if (blockFirstOffset + threadIdx.x > lastOffset) \ + return; \ + uint32_t decodedLeafIndex; \ + uint16_t decodedVoxelOffset; \ + VoxelBlockManagerT::template decodeInverseMap( \ + grid, \ + firstLeafID[blockIdx.x], \ + &jumpMap[uint64_t(blockIdx.x) * jumpMapWordCount], \ + blockFirstOffset, \ + int(threadIdx.x), \ + decodedLeafIndex, \ + decodedVoxelOffset); \ + if (decodedLeafIndex == VoxelBlockManagerT::UnusedLeafIndex) \ + return; /* defensive; the lastOffset guard above makes this unreachable */ \ + const auto &leaf = grid->tree().template getFirstNode<0>()[decodedLeafIndex]; \ + const nanovdb::Coord centerCoord = leaf.offsetToGlobalCoord(decodedVoxelOffset); \ + const uint64_t centerIndex = leaf.getValue(decodedVoxelOffset); \ + auto accessor = grid->getAccessor(); \ + const uint64_t faceIndex[6] = {accessor.getValue(centerCoord.offsetBy(-1, 0, 0)), \ + accessor.getValue(centerCoord.offsetBy(1, 0, 0)), \ + accessor.getValue(centerCoord.offsetBy(0, -1, 0)), \ + accessor.getValue(centerCoord.offsetBy(0, 1, 0)), \ + accessor.getValue(centerCoord.offsetBy(0, 0, -1)), \ accessor.getValue(centerCoord.offsetBy(0, 0, 1))}; // ===================== fused stencil kernels ==================================================== @@ -68,6 +72,7 @@ signFusedKernel(const OnIndexGridT *grid, const uint32_t *firstLeafID, const uint64_t *jumpMap, uint64_t firstOffset, + uint64_t lastOffset, const ScalarT *field, ScalarT voxelSize, ScalarT *sign) { @@ -105,6 +110,7 @@ godunovFusedKernel(const OnIndexGridT *grid, const uint32_t *firstLeafID, const uint64_t *jumpMap, uint64_t firstOffset, + uint64_t lastOffset, const ScalarT *field, const ScalarT *sign, ScalarT voxelSize, @@ -128,6 +134,7 @@ smoothFusedKernel(const OnIndexGridT *grid, const uint32_t *firstLeafID, const uint64_t *jumpMap, uint64_t firstOffset, + uint64_t lastOffset, const ScalarT *in, ScalarT weight, ScalarT *out) { @@ -189,28 +196,6 @@ heunKernel(ScalarT *out, out[i] = nanovdb::math::Min(nanovdb::math::Max(value, -bandWidth), bandWidth); } -// small VBM helper: build once, expose the block count + the firstLeafID/jumpMap device pointers. -struct VBMHelper { - nanovdb::tools::VoxelBlockManagerHandle handle; - uint32_t blockCount{0}; - uint64_t firstOffset{0}, valueCount{1}; - VBMHelper(OnIndexGridT *grid, cudaStream_t stream) { - handle = nanovdb::tools::cuda::buildVoxelBlockManager( - grid, 0, 0, 0, stream); - blockCount = (uint32_t)handle.blockCount(); - firstOffset = handle.firstOffset(); - valueCount = handle.lastOffset() + 1; - } - const uint32_t * - firstLeafID() const { - return handle.deviceFirstLeafID(); - } - const uint64_t * - jumpMap() const { - return handle.deviceJumpMap(); - } -}; - // Redistance (|grad phi| = 1) + optional de-staircase one grid's value-indexed buffer, in place. // `phi`/scratch are length `valueCount` with slot 0 holding the +bandWidth background; the // stencil/combiner kernels never write slot 0 (so inactive-neighbour reads always see the boundary @@ -218,7 +203,7 @@ struct VBMHelper { template void runReinit(OnIndexGridT *grid, - const VBMHelper &vbm, + const VbmCache::GridVbm &vbm, ScalarT *phi, ScalarT *sign, ScalarT *phiBase, @@ -236,22 +221,23 @@ runReinit(OnIndexGridT *grid, cudaStream_t stream) { const uint32_t blockCount = vbm.blockCount; constexpr int blockWidth = 1 << kLog2BlockWidth; - const uint32_t *firstLeafID = vbm.firstLeafID(); - const uint64_t *jumpMap = vbm.jumpMap(); + const uint32_t *firstLeafID = vbm.firstLeafID; + const uint64_t *jumpMap = vbm.jumpMap; const uint64_t firstOffset = vbm.firstOffset; + const uint64_t lastOffset = vbm.lastOffset; const ScalarT timeStep = ScalarT(0.4) * voxelSize; auto godunov = [&](const ScalarT *field, ScalarT *out) { if (blockCount) { godunovFusedKernel<<>>( - grid, firstLeafID, jumpMap, firstOffset, field, sign, voxelSize, out); + grid, firstLeafID, jumpMap, firstOffset, lastOffset, field, sign, voxelSize, out); C10_CUDA_KERNEL_LAUNCH_CHECK(); } }; auto redistance = [&](int iters) { if (blockCount) { signFusedKernel<<>>( - grid, firstLeafID, jumpMap, firstOffset, phi, voxelSize, sign); + grid, firstLeafID, jumpMap, firstOffset, lastOffset, phi, voxelSize, sign); C10_CUDA_KERNEL_LAUNCH_CHECK(); } for (int it = 0; it < iters; ++it) { @@ -320,7 +306,7 @@ runReinit(OnIndexGridT *grid, auto pass = [&](ScalarT weight) { if (blockCount) { smoothFusedKernel<<>>( - grid, firstLeafID, jumpMap, firstOffset, cur, weight, other); + grid, firstLeafID, jumpMap, firstOffset, lastOffset, cur, weight, other); C10_CUDA_KERNEL_LAUNCH_CHECK(); } std::swap(cur, other); @@ -361,14 +347,16 @@ reinitializeSdfCuda(const GridBatchData &batchHdl, const int64_t numVoxels = batchHdl.numVoxelsAt(batchIdx); if (numVoxels == 0) continue; - OnIndexGridT *grid = - batchHdl.mGridHdl->deviceGrid((uint32_t)batchIdx); + // deviceGridPtrAt resolves the *logical* grid by byte offset, which is correct for + // sliced/indexed views (mGridHdl->deviceGrid(i) indexes physically and reads the wrong + // grid for a view). + OnIndexGridT *grid = batchHdl.deviceGridPtrAt(batchIdx); const int64_t voxelOffset = batchHdl.cumVoxelsAt(batchIdx); const ScalarT voxelSize = (ScalarT)batchHdl.voxelSizeAt(batchIdx)[0]; const ScalarT bandWidth = (ScalarT)band * voxelSize; // narrow-band half-width, world units - VBMHelper vbm(grid, stream); - const int64_t valueCount = (int64_t)vbm.valueCount; // numVoxels + 1 (slot 0 = background) + const VbmCache::GridVbm vbm = batchHdl.vbmCache().get(batchHdl, batchIdx); + const int64_t valueCount = int64_t(vbm.lastOffset) + 1; // numVoxels + 1 (slot 0 = bg) torch::Tensor phiBuf = torch::empty({valueCount}, opts); torch::Tensor signBuf = torch::empty({valueCount}, opts); diff --git a/src/fvdb/detail/utils/SimpleOpHelper.h b/src/fvdb/detail/utils/SimpleOpHelper.h index aa04e0f36..8e3bad1a7 100644 --- a/src/fvdb/detail/utils/SimpleOpHelper.h +++ b/src/fvdb/detail/utils/SimpleOpHelper.h @@ -9,6 +9,7 @@ #include #include #include +#include #include @@ -219,8 +220,10 @@ struct BasePerActiveVoxelProcessor { makeOutTensorFromGridBatch(grid_batch, out_element); auto out_accessor = makeAccessor(out_tensor); if constexpr (DeviceTag == torch::kCUDA) { - forEachVoxelCUDA( - 1, grid_batch, *static_cast(this), out_accessor); + // VBM-backed iteration: one thread per active voxel via the cached per-grid + // VoxelBlockManager decode, instead of one thread per 512-slot leaf position. + forEachActiveVoxelVbmCUDA( + grid_batch, *static_cast(this), out_accessor); } else if constexpr (DeviceTag == torch::kPrivateUse1) { forEachVoxelPrivateUse1( 1, grid_batch, *static_cast(this), out_accessor); diff --git a/src/fvdb/detail/utils/cuda/ForEachVbmCUDA.cuh b/src/fvdb/detail/utils/cuda/ForEachVbmCUDA.cuh new file mode 100644 index 000000000..812b6eef5 --- /dev/null +++ b/src/fvdb/detail/utils/cuda/ForEachVbmCUDA.cuh @@ -0,0 +1,117 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#ifndef FVDB_DETAIL_UTILS_CUDA_FOREACHVBMCUDA_CUH +#define FVDB_DETAIL_UTILS_CUDA_FOREACHVBMCUDA_CUH + +#include +#include + +#include + +#include +#include +#include + +namespace fvdb { + +namespace _private { + +constexpr int kVbmForEachBlockDim = 256; + +/// Grid-stride over one grid's active voxels: each iteration decodes one voxel's +/// (leafIndex, voxelOffset) in registers via the VoxelBlockManager's rank+select inverse-map +/// decode -- no shared memory, no barriers, and no threads wasted on inactive leaf slots +/// (unlike the totalLeaves*512 leaf scan in forEachVoxelCUDAKernel). decodeInverseMap takes +/// the in-block slot position explicitly, so the launch shape is decoupled from the VBM's +/// 512-slot blocks; a 256-thread grid-stride kernel leaves the per-thread register budget to +/// the callback. (A one-thread-per-slot <<>> shape with +/// __launch_bounds__(512) was observed to silently drop work from a register-heavy callback.) +template +__global__ void __launch_bounds__(kVbmForEachBlockDim) +forEachActiveVoxelVbmKernel(const nanovdb::OnIndexGrid *__restrict__ grid, + const uint32_t *__restrict__ firstLeafID, + const uint64_t *__restrict__ jumpMap, + uint64_t firstOffset, + uint64_t lastOffset, + int64_t baseVoxelOffset, // cumVoxelsAt(bi) of this grid in the batch + Func func, + Args... args) { + using VbmT = nanovdb::tools::cuda::VoxelBlockManager; + + const uint64_t numSlots = lastOffset - firstOffset + 1; + const uint64_t stride = uint64_t(gridDim.x) * blockDim.x; + + for (uint64_t i = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x; i < numSlots; i += stride) { + const uint32_t vbmBlock = uint32_t(i >> detail::VbmCache::kLog2BlockWidth); + const int blockOffset = int(i & (detail::VbmCache::kBlockWidth - 1)); + const uint64_t blockFirstOffset = + firstOffset + uint64_t(vbmBlock) * detail::VbmCache::kBlockWidth; + uint32_t leafIndex; + uint16_t voxelOffset; + VbmT::decodeInverseMap(grid, + firstLeafID[vbmBlock], + jumpMap + uint64_t(vbmBlock) * detail::VbmCache::kJumpMapWordCount, + blockFirstOffset, + blockOffset, + leafIndex, + voxelOffset); + if (leafIndex == VbmT::UnusedLeafIndex) { + continue; // defensive; i < numSlots makes this unreachable + } + const auto &leaf = grid->tree().template getFirstNode<0>()[leafIndex]; + const nanovdb::Coord ijk = leaf.offsetToGlobalCoord(voxelOffset); + // Sequential OnIndex invariant: the decoded voxel's value index equals its slot + // (firstOffset + i), so the batch-wide feature index is baseVoxelOffset + + // (firstOffset + i) - 1 (matching the legacy path's baseOffset + getValue - 1). + const int64_t featureIdx = baseVoxelOffset + int64_t(firstOffset + i) - 1; + func.perActiveVoxel(ijk, featureIdx, args...); + } +} + +} // namespace _private + +/// @brief Run func.perActiveVoxel(ijk, featureIdx, args...) for every active voxel of every +/// grid in the batch, using the batch's cached per-grid VoxelBlockManagers: exactly one +/// grid-stride iteration per active voxel, independent of leaf occupancy. +/// +/// The callback contract matches BasePerActiveVoxelProcessor::perActiveVoxel: +/// void perActiveVoxel(nanovdb::Coord const &ijk, int64_t featureIdx, Args...) const +/// where featureIdx is the batch-wide linear voxel index (cumVoxelsAt(bi) + in-grid index). +/// +/// One kernel launch per grid in the batch. CUDA-only; callers must fall back to the legacy +/// leaf-scan paths for CPU and PrivateUse1 grids. +template +void +forEachActiveVoxelVbmCUDA(const fvdb::GridBatchData &batchHdl, Func func, Args... args) { + TORCH_CHECK(batchHdl.device().is_cuda(), "Grid batch must be on a CUDA device"); + TORCH_CHECK(batchHdl.device().has_index(), "Grid batch device must have an index"); + c10::cuda::CUDAGuard deviceGuard(batchHdl.device()); + const at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(batchHdl.device().index()); + + for (int64_t bi = 0; bi < batchHdl.batchSize(); ++bi) { + const auto vbm = batchHdl.vbmCache().get(batchHdl, bi); + if (vbm.blockCount == 0) { + continue; // empty grid + } + const uint64_t numSlots = vbm.lastOffset - vbm.firstOffset + 1; + const uint32_t numBlocks = uint32_t((numSlots + _private::kVbmForEachBlockDim - 1) / + _private::kVbmForEachBlockDim); + _private::forEachActiveVoxelVbmKernel<<>>(batchHdl.deviceGridPtrAt(bi), + vbm.firstLeafID, + vbm.jumpMap, + vbm.firstOffset, + vbm.lastOffset, + batchHdl.cumVoxelsAt(bi), + func, + args...); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } +} + +} // namespace fvdb + +#endif // FVDB_DETAIL_UTILS_CUDA_FOREACHVBMCUDA_CUH diff --git a/src/tests/CMakeLists.txt b/src/tests/CMakeLists.txt index f7d538194..295625c80 100644 --- a/src/tests/CMakeLists.txt +++ b/src/tests/CMakeLists.txt @@ -167,6 +167,12 @@ ConfigureTest(GaussianMCMCAddNoiseTest "GaussianMCMCAddNoiseTest.cpp") ConfigureTest(GaussianMCMCRelocationTest "GaussianMCMCRelocationTest.cpp") ConfigureTest(GatherScatterDefaultConvTest "GatherScatterDefaultConvTest.cu") ConfigureTest(PredGatherIGemmTest "PredGatherIGemmTest.cu") +ConfigureTest(VbmCacheTest "VbmCacheTest.cu") +# NanoVDB's host-side VoxelBlockManager uses `#pragma omp simd`, which -Wall -Werror rejects +# as an unknown pragma when OpenMP is off (the fvdb library build suppresses it the same way). +target_compile_options(VbmCacheTest_obj PRIVATE + $<$:-Xcompiler=-Wno-unknown-pragmas> + $<$:-Wno-unknown-pragmas>) if(NANOVDB_EDITOR_INSTALLED) #ConfigureTest(ViewerTest "ViewerTest.cpp") endif() diff --git a/src/tests/VbmCacheTest.cu b/src/tests/VbmCacheTest.cu new file mode 100644 index 000000000..5bd65279d --- /dev/null +++ b/src/tests/VbmCacheTest.cu @@ -0,0 +1,310 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +// VbmCacheTest.cu -- Tests for the lazily-built per-grid VoxelBlockManager cache. +// +// The VBM decode contract these tests pin down: +// - decoding slot s of grid bi yields the leaf/voxel whose sequential value index is s, +// i.e. leaf.getValue(voxelOffset) == s and leaf.offsetToGlobalCoord(voxelOffset) is the +// s-th active coordinate in the grid's sequential order (== activeGridCoords row s-1); +// - cache entries are built once per grid and shared between a GridBatchData and its +// sliced/indexed views (same underlying buffer); +// - empty grids yield a zero GridVbm and never touch the device; +// - every fVDB grid production path emits sequential (breadth-first, fixed-size) grids, +// which the decode requires (grid->isSequential()). +// +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +#include + +#include + +using namespace fvdb; +using namespace fvdb::detail; + +namespace { + +bool +cudaIsAvailable() { + int count = 0; + auto err = cudaGetDeviceCount(&count); + return err == cudaSuccess && count > 0; +} + +// Deterministic pseudo-random unique ijk coordinates: a seeded permutation of a dense box, +// truncated to numVoxels rows. Coordinates are unique by construction so the built grid has +// exactly numVoxels active voxels. +torch::Tensor +randomIjk(int64_t numVoxels, int boxDim, uint64_t seed) { + const int64_t boxVolume = int64_t(boxDim) * boxDim * boxDim; + TORCH_CHECK(numVoxels <= boxVolume, "box too small for requested voxel count"); + auto gen = at::detail::createCPUGenerator(seed); + auto perm = torch::randperm(boxVolume, gen, torch::kInt64).slice(0, 0, numVoxels); + auto ijk = torch::empty({numVoxels, 3}, torch::kInt32); + ijk.select(1, 0).copy_(perm.div(boxDim * boxDim, "floor")); + ijk.select(1, 1).copy_(perm.div(boxDim, "floor").remainder(boxDim)); + ijk.select(1, 2).copy_(perm.remainder(boxDim)); + return ijk; +} + +c10::intrusive_ptr +makeBatch(const std::vector &ijkPerGrid, torch::Device device) { + std::vector onDevice; + std::vector voxelSizes, origins; + for (const auto &ijk: ijkPerGrid) { + onDevice.push_back(ijk.to(device)); + voxelSizes.push_back({1.0, 1.0, 1.0}); + origins.push_back({0.0, 0.0, 0.0}); + } + JaggedTensor jt(onDevice); + return ops::createNanoGridFromIJK(jt, voxelSizes, origins); +} + +using VbmT = nanovdb::tools::cuda::VoxelBlockManager; + +// One thread per VBM slot: decode and record the coordinate and value index of the decoded +// voxel so the host can compare against the legacy (leaf-scan) ground truth. +__global__ void +decodeAllSlotsKernel(const nanovdb::OnIndexGrid *grid, + const uint32_t *firstLeafID, + const uint64_t *jumpMap, + uint64_t firstOffset, + uint64_t lastOffset, + int32_t *outIjk, // [numVoxels, 3] + int64_t *outValueIdx) // [numVoxels] +{ + const uint64_t blockFirstOffset = firstOffset + uint64_t(blockIdx.x) * VbmCache::kBlockWidth; + const uint64_t slot = blockFirstOffset + threadIdx.x; + if (slot > lastOffset) { + return; + } + uint32_t leafIndex; + uint16_t voxelOffset; + VbmT::decodeInverseMap(grid, + firstLeafID[blockIdx.x], + jumpMap + uint64_t(blockIdx.x) * VbmCache::kJumpMapWordCount, + blockFirstOffset, + int(threadIdx.x), + leafIndex, + voxelOffset); + if (leafIndex == VbmT::UnusedLeafIndex) { + return; + } + const auto &leaf = grid->tree().template getFirstNode<0>()[leafIndex]; + const nanovdb::Coord ijk = leaf.offsetToGlobalCoord(voxelOffset); + const uint64_t row = slot - 1; + outIjk[row * 3 + 0] = ijk[0]; + outIjk[row * 3 + 1] = ijk[1]; + outIjk[row * 3 + 2] = ijk[2]; + outValueIdx[row] = int64_t(leaf.getValue(voxelOffset)); +} + +// Decode every slot of grid `bi` through the cache and assert coordinate and value-index +// parity against the legacy leaf-scan ground truth (activeGridCoords on a CPU twin). +void +expectDecodeParity(GridBatchData &batch, int64_t bi, const torch::Tensor &expectedIjkCpu) { + auto vbm = batch.vbmCache().get(batch, bi); + ASSERT_EQ(int64_t(vbm.lastOffset), batch.numVoxelsAt(bi)); + ASSERT_EQ(vbm.firstOffset, 1u); + ASSERT_EQ( + vbm.blockCount, + uint32_t((batch.numVoxelsAt(bi) + VbmCache::kBlockWidth - 1) >> VbmCache::kLog2BlockWidth)); + + const int64_t numVoxels = batch.numVoxelsAt(bi); + auto opts = torch::TensorOptions().device(batch.device()); + auto outIjk = torch::full({numVoxels, 3}, -12345, opts.dtype(torch::kInt32)); + auto outValueIdx = torch::zeros({numVoxels}, opts.dtype(torch::kInt64)); + + decodeAllSlotsKernel<<>>( + batch.deviceGridPtrAt(bi), + vbm.firstLeafID, + vbm.jumpMap, + vbm.firstOffset, + vbm.lastOffset, + outIjk.data_ptr(), + outValueIdx.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + EXPECT_TRUE(torch::equal(outIjk.cpu(), expectedIjkCpu)); + // Sequential OnIndex invariant: the decoded voxel's value index equals its slot. + EXPECT_TRUE(torch::equal(outValueIdx.cpu(), torch::arange(1, numVoxels + 1, torch::kInt64))); +} + +} // namespace + +TEST(VbmCacheTest, DecodeParity) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + auto ijk = randomIjk(50000, 64, /*seed=*/42); + auto cudaBatch = makeBatch({ijk}, torch::Device(torch::kCUDA, 0)); + auto cpuBatch = makeBatch({ijk}, torch::Device(torch::kCPU)); + ASSERT_EQ(cudaBatch->numVoxelsAt(0), 50000); + + auto expected = ops::activeGridCoords(*cpuBatch).jdata(); + expectDecodeParity(*cudaBatch, 0, expected); +} + +TEST(VbmCacheTest, MultiBatchDecodeParity) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + // Different sizes, deliberately including exact multiples of the block width (512) and + // off-by-one sizes around it. + std::vector ijks = {randomIjk(100, 16, 1), + randomIjk(512, 32, 2), + randomIjk(513, 32, 3), + randomIjk(9000, 48, 4)}; + auto cudaBatch = makeBatch(ijks, torch::Device(torch::kCUDA, 0)); + auto cpuBatch = makeBatch(ijks, torch::Device(torch::kCPU)); + auto expected = ops::activeGridCoords(*cpuBatch); + + for (int64_t bi = 0; bi < cudaBatch->batchSize(); ++bi) { + SCOPED_TRACE("grid " + std::to_string(bi)); + expectDecodeParity(*cudaBatch, bi, expected.index(bi).jdata()); + } +} + +TEST(VbmCacheTest, ViewSharing) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + std::vector ijks = { + randomIjk(600, 16, 5), randomIjk(700, 16, 6), randomIjk(800, 16, 7)}; + auto parent = makeBatch(ijks, torch::Device(torch::kCUDA, 0)); + auto slice = ops::indexGrid(*parent, 1, 3, 1); + + // The slice shares the parent's cache object, and the same logical grid resolves to the + // same cached entry (same device pointers) through either. + ASSERT_EQ(&slice->vbmCache(), &parent->vbmCache()); + auto fromSlice = slice->vbmCache().get(*slice, 0); + auto fromParent = parent->vbmCache().get(*parent, 1); + EXPECT_EQ(fromSlice.firstLeafID, fromParent.firstLeafID); + EXPECT_EQ(fromSlice.jumpMap, fromParent.jumpMap); + EXPECT_EQ(fromSlice.blockCount, fromParent.blockCount); + + // Decode parity through the view. + auto cpuBatch = makeBatch(ijks, torch::Device(torch::kCPU)); + auto expected = ops::activeGridCoords(*cpuBatch); + expectDecodeParity(*slice, 0, expected.index(1).jdata()); + expectDecodeParity(*slice, 1, expected.index(2).jdata()); +} + +TEST(VbmCacheTest, EmptyGrid) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + std::vector ijks = {torch::empty({0, 3}, torch::kInt32), randomIjk(100, 16, 8)}; + auto batch = makeBatch(ijks, torch::Device(torch::kCUDA, 0)); + ASSERT_EQ(batch->numVoxelsAt(0), 0); + + auto vbm = batch->vbmCache().get(*batch, 0); + EXPECT_EQ(vbm.blockCount, 0u); + EXPECT_EQ(vbm.firstLeafID, nullptr); + EXPECT_EQ(vbm.jumpMap, nullptr); + + auto cpuBatch = makeBatch(ijks, torch::Device(torch::kCPU)); + auto expected = ops::activeGridCoords(*cpuBatch); + expectDecodeParity(*batch, 1, expected.index(1).jdata()); +} + +TEST(VbmCacheTest, RepeatedGetIsCached) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + auto batch = makeBatch({randomIjk(2000, 32, 9)}, torch::Device(torch::kCUDA, 0)); + auto first = batch->vbmCache().get(*batch, 0); + auto again = batch->vbmCache().get(*batch, 0); + EXPECT_EQ(first.firstLeafID, again.firstLeafID); + EXPECT_EQ(first.jumpMap, again.jumpMap); +} + +// A consumer on a different stream than the build must (a) observe the finished build (the +// build event wait) and (b) keep the cached buffers alive until its kernels drain, even if the +// owning GridBatchData is destroyed while they are still in flight (record_stream with the +// caching allocator). Decode parity through the side stream verifies (a); destroying the batch +// immediately after the async launch exercises (b) -- a lifetime bug here surfaces as corrupt +// output or as an invalid access under compute-sanitizer. +TEST(VbmCacheTest, CrossStreamConsumerLifetime) { + if (!cudaIsAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + auto ijk = randomIjk(50000, 64, 12); + auto cpuBatch = makeBatch({ijk}, torch::Device(torch::kCPU)); + auto expected = ops::activeGridCoords(*cpuBatch).jdata(); + + auto device = torch::Device(torch::kCUDA, 0); + auto opts = torch::TensorOptions().device(device); + const int64_t numVoxels = 50000; + auto outIjk = torch::full({numVoxels, 3}, -12345, opts.dtype(torch::kInt32)); + auto outValueIdx = torch::zeros({numVoxels}, opts.dtype(torch::kInt64)); + + { + auto batch = makeBatch({ijk}, device); + ASSERT_EQ(batch->numVoxelsAt(0), numVoxels); + // Build the cache entry on the default stream. + (void)batch->vbmCache().get(*batch, 0); + + // Consume it from a side stream, then destroy the batch (and with it the cache) while + // the side stream's kernel may still be running. + c10::cuda::CUDAStream sideStream = c10::cuda::getStreamFromPool(false, device.index()); + { + c10::cuda::CUDAStreamGuard streamGuard(sideStream); + auto vbm = batch->vbmCache().get(*batch, 0); + decodeAllSlotsKernel<<>>( + batch->deviceGridPtrAt(0), + vbm.firstLeafID, + vbm.jumpMap, + vbm.firstOffset, + vbm.lastOffset, + outIjk.data_ptr(), + outValueIdx.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + } // batch (and VbmCache) destroyed here, side-stream kernel possibly in flight + + C10_CUDA_CHECK(cudaDeviceSynchronize()); + EXPECT_TRUE(torch::equal(outIjk.cpu(), expected)); + EXPECT_TRUE(torch::equal(outValueIdx.cpu(), torch::arange(1, numVoxels + 1, torch::kInt64))); +} + +// The register decode requires grid->isSequential() (fixed-size, breadth-first leaves). Pin +// that invariant for every grid production path, on CPU grids where the header is readable. +TEST(VbmCacheTest, ProductionPathsAreSequential) { + auto cpu = torch::Device(torch::kCPU); + + std::vector ijks = {randomIjk(1000, 24, 10), randomIjk(1500, 24, 11)}; + auto built = makeBatch(ijks, cpu); + auto padded = ops::buildPaddedGrid(*built, -1, 1, false, false); + auto sliced = ops::indexGrid(*built, 1, 2, 1); + auto concatenated = ops::concatenateGrids({built, padded}); + auto contiguous = ops::makeContiguous(sliced); + + for (const auto &[name, batch]: + std::vector>>{ + {"createNanoGridFromIJK", built}, + {"buildPaddedGrid", padded}, + {"indexGrid", sliced}, + {"concatenateGrids", concatenated}, + {"makeContiguous", contiguous}}) { + for (int64_t bi = 0; bi < batch->batchSize(); ++bi) { + EXPECT_TRUE(batch->hostGridPtrAt(bi)->isSequential()) + << name << " grid " << bi << " is not sequential"; + } + } +} diff --git a/tests/unit/test_active_grid_coords_vbm.py b/tests/unit/test_active_grid_coords_vbm.py new file mode 100644 index 000000000..abc7487eb --- /dev/null +++ b/tests/unit/test_active_grid_coords_vbm.py @@ -0,0 +1,59 @@ +# Copyright Contributors to the OpenVDB Project +# SPDX-License-Identifier: Apache-2.0 +# +# Parity tests for the VBM-backed CUDA per-active-voxel iteration path used by +# Grid.ijk / GridBatch.ijk (ActiveGridCoords) and morton/hilbert (SerializeEncode): +# the CUDA results must match the CPU leaf-scan path bitwise, including through +# sliced GridBatch views (which share the parent's cached VBM handles) and for +# batches containing empty grids. +import unittest + +import torch + +import fvdb + + +def _random_ijk(num: int, box: int, seed: int) -> torch.Tensor: + gen = torch.Generator().manual_seed(seed) + perm = torch.randperm(box * box * box, generator=gen)[:num] + return torch.stack([perm // (box * box), (perm // box) % box, perm % box], dim=1).to(torch.int32) + + +class ActiveGridCoordsVbmTests(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available(): + self.skipTest("requires a CUDA device") + self.device = torch.device("cuda:0") + + def _make_batches(self, ijks): + cpu = fvdb.GridBatch.from_ijk(fvdb.JaggedTensor([i for i in ijks])) + cuda = fvdb.GridBatch.from_ijk(fvdb.JaggedTensor([i.to(self.device) for i in ijks])) + return cpu, cuda + + def test_ijk_matches_cpu(self): + ijks = [_random_ijk(100, 16, 1), _random_ijk(512, 32, 2), _random_ijk(513, 32, 3), _random_ijk(9000, 48, 4)] + cpu, cuda = self._make_batches(ijks) + self.assertTrue(torch.equal(cuda.ijk.jdata.cpu(), cpu.ijk.jdata)) + + def test_ijk_on_sliced_batch(self): + ijks = [_random_ijk(600, 16, 5), _random_ijk(700, 16, 6), _random_ijk(800, 16, 7)] + cpu, cuda = self._make_batches(ijks) + # Repeated access exercises the cached VBM entries shared between parent and view. + _ = cuda.ijk + self.assertTrue(torch.equal(cuda[1:3].ijk.jdata.cpu(), cpu[1:3].ijk.jdata)) + self.assertTrue(torch.equal(cuda[1:3].ijk.jdata.cpu(), cpu.ijk[1:3].jdata)) + + def test_ijk_with_empty_grid(self): + ijks = [torch.empty(0, 3, dtype=torch.int32), _random_ijk(100, 16, 8)] + cpu, cuda = self._make_batches(ijks) + self.assertEqual(cuda.ijk[0].jdata.shape[0], 0) + self.assertTrue(torch.equal(cuda.ijk.jdata.cpu(), cpu.ijk.jdata)) + + def test_morton_matches_cpu(self): + ijks = [_random_ijk(1000, 24, 9), _random_ijk(1500, 24, 10)] + cpu, cuda = self._make_batches(ijks) + self.assertTrue(torch.equal(cuda.morton().jdata.cpu(), cpu.morton().jdata)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_sdf.py b/tests/unit/test_sdf.py index 8a157034f..bbf5defce 100644 --- a/tests/unit/test_sdf.py +++ b/tests/unit/test_sdf.py @@ -150,6 +150,32 @@ def _cube_points(self, vx: float, half: int) -> torch.Tensor: ii, jj, kk = torch.meshgrid(rng, rng, rng, indexing="ij") return torch.stack([ii, jj, kk], dim=-1).reshape(-1, 3) * vx + def test_sliced_batch_matches_standalone(self): + """reinitialize_sdf on a sliced GridBatch view must read the view's grids, not the + physical grids of the underlying batch (regression test for a physical-vs-logical + grid indexing bug).""" + vx = self.vx + halves = [12, 10, 8] + gb = fvdb.GridBatch.from_points( + fvdb.JaggedTensor([self._cube_points(vx, h) for h in halves]), + voxel_sizes=vx, + ) + view = gb[1:3] + analytic = (view.ijk.jdata.float() * vx).norm(dim=1) - self.R + field = view.jagged_like(analytic.clamp(-self.bw, self.bw)) + phi_view = view.reinitialize_sdf(field, band=self.band, order=3) + + standalone = fvdb.GridBatch.from_points( + fvdb.JaggedTensor([self._cube_points(vx, h) for h in halves[1:]]), + voxel_sizes=vx, + ) + analytic_sa = (standalone.ijk.jdata.float() * vx).norm(dim=1) - self.R + field_sa = standalone.jagged_like(analytic_sa.clamp(-self.bw, self.bw)) + phi_sa = standalone.reinitialize_sdf(field_sa, band=self.band, order=3) + + self.assertTrue(torch.equal(view.ijk.jdata, standalone.ijk.jdata)) + self.assertTrue(torch.allclose(phi_view.jdata, phi_sa.jdata, atol=1e-5)) + if __name__ == "__main__": unittest.main()