diff --git a/fvdb/convolution_plan.py b/fvdb/convolution_plan.py index 3ae2a5cf2..d73cd7cda 100644 --- a/fvdb/convolution_plan.py +++ b/fvdb/convolution_plan.py @@ -716,6 +716,19 @@ def from_grid_batch( if backend_name == "pred_gather_igemm": _validate_pred_gather_igemm_admission(kernel_size, stride, channel_pairs, transposed=False) _validate_pred_gather_igemm_grid_admission(source_grid, target_grid) + elif backend_name in ("default", "gather_scatter"): + identity_plan = cls._identity_matmul_plan( + kernel_size, + stride, + source_grid, + target_grid, + channel_pairs, + resolved_policy, + topology_provenance, + transposed=False, + ) + if identity_plan is not None: + return identity_plan if target_grid is None: target_grid = source_grid.conv_grid(kernel_size, stride) @@ -808,6 +821,19 @@ def from_grid_batch_transposed( if backend_name == "pred_gather_igemm": _validate_pred_gather_igemm_admission(kernel_size, stride, channel_pairs, transposed=True) _validate_pred_gather_igemm_grid_admission(source_grid, target_grid) + elif backend_name in ("default", "gather_scatter"): + identity_plan = cls._identity_matmul_plan( + kernel_size, + stride, + source_grid, + target_grid, + channel_pairs, + resolved_policy, + topology_provenance, + transposed=True, + ) + if identity_plan is not None: + return identity_plan if target_grid is None: target_grid = source_grid.conv_transpose_grid(kernel_size, stride) @@ -1164,6 +1190,66 @@ def has_fixed_topology(self) -> bool: # Private methods # ============================================================ + @classmethod + def _identity_matmul_plan( + cls, + kernel_size: torch.Tensor, + stride: torch.Tensor, + source_grid: GridBatch, + target_grid: GridBatch | None, + channel_pairs: tuple[tuple[int, int], ...], + resolved_policy: ConvolutionTopologyPolicy, + topology_provenance: ConvolutionTopologyProvenance, + transposed: bool, + ) -> "ConvolutionPlan | None": + """Identity (K == S == 1, shared grid data) short circuit to the matmul backend. + + Per-iteration classifier heads rebuild identity plans constantly (issue #755); the general + path spends ~1 ms per call building tensor-valued transform diagnostics that are trivially + satisfied when source and target share their GridBatchData. Returns None when the fast + path does not apply (the caller then follows the general path, keeping full validation + for distinct-but-equal-looking grids and incompatible transforms). + """ + if kernel_size.tolist() != [1, 1, 1] or stride.tolist() != [1, 1, 1]: + return None + if target_grid is None: + # conv_grid / conv_transpose_grid are the identity at K == S == 1: the generated + # target is the source grid itself, preserving public and data identity. + target_grid = source_grid + elif not _get_grid_data(source_grid).is_same(_get_grid_data(target_grid)): + return None + # Preserve the general path's construction-time channel-pair validation. + for channel_pair in channel_pairs: + if len(channel_pair) != 2 or channel_pair[0] <= 0 or channel_pair[1] <= 0: + raise ValueError("channel_pair must be a tuple of two positive integers") + geometry = _fvdb_cpp.ConvolutionGeometry(kernel_size, stride) + # Shared grid data makes every registration diagnostic exact by construction. + compatibility = ConvolutionTransformCompatibility( + fine_grid_count=source_grid.grid_count, + coarse_grid_count=source_grid.grid_count, + same_batch_size=True, + same_device=True, + scale_compatible=True, + registration_integer=True, + registration_zero=True, + compatible=True, + registration_offset=torch.zeros((source_grid.grid_count, 3), dtype=torch.float64), + ) + backend = _MatmulBackend() + return cls( + source_grid, + target_grid, + geometry, + channel_pairs, + transposed, + backend, + compatibility, + resolved_policy, + topology_provenance, + _CoverageReportCache(backend, source_grid, target_grid), + False, + ) + @staticmethod def _build_backend( source_grid: GridBatch, diff --git a/src/benchmarks/convolution/benchmark_conv_grid_build.py b/src/benchmarks/convolution/benchmark_conv_grid_build.py new file mode 100644 index 000000000..da7cb8b37 --- /dev/null +++ b/src/benchmarks/convolution/benchmark_conv_grid_build.py @@ -0,0 +1,132 @@ +# Copyright Contributors to the OpenVDB Project +# SPDX-License-Identifier: Apache-2.0 +# +"""Benchmark generated-topology grid construction (issue #755). + +Times the per-call cost of conv_grid / conv_transpose_grid / refined_grid / coarsened_grid and of +ConvolutionPlan construction with generated targets, as a function of batch size. Before the +batched leaf-mask builder these scaled linearly in batch size (~1.5 ms fixed overhead per member); +after, they should be near-flat in B. + +Usage: + python src/benchmarks/convolution/benchmark_conv_grid_build.py [--json results.json] [--gso] + +--gso additionally runs the issue's verbatim repro on the GSO shoes dataset (requires the +dataset download used by fvdb.utils.examples.load_gso_shoes). +""" + +import argparse +import json +import time + +import torch + +import fvdb +from fvdb import ConvolutionPlan, GridBatch, JaggedTensor + + +def make_shell_batch(batch_size: int, resolution: int = 64, device: str = "cuda") -> GridBatch: + """B roughly-spherical shells of ~`resolution`^2*3 voxels each (surface-like sparsity, + similar occupancy statistics to meshes voxelized at `resolution`).""" + ijks = [] + for b in range(batch_size): + torch.manual_seed(1234 + b) + n = resolution * resolution * 6 + pts = torch.randn(n, 3, dtype=torch.float64) + pts = pts / pts.norm(dim=-1, keepdim=True) + radius = 0.35 + 0.05 * (b % 5) / 5.0 + ijk = ((pts * radius + 0.5) * resolution).floor().to(torch.int32) + ijks.append(torch.unique(ijk, dim=0)) + jt = JaggedTensor([t.to(device) for t in ijks]) + return GridBatch.from_ijk(jt, voxel_sizes=1.0 / resolution, origins=0.0) + + +def time_op(fn, warmup: int = 3, iters: int = 20) -> float: + """Median wall time of fn() in milliseconds, CUDA-event timed.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + times = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + times.append(start.elapsed_time(end)) + times.sort() + return times[len(times) // 2] + + +def bench_batch(grid: GridBatch) -> dict: + ops = { + "conv_transpose_grid k2s2": lambda: grid.conv_transpose_grid(kernel_size=2, stride=2), + "conv_grid k2s2": lambda: grid.conv_grid(kernel_size=2, stride=2), + "conv_grid k3s1": lambda: grid.conv_grid(kernel_size=3, stride=1), + "refined_grid x2": lambda: grid.refined_grid(2), + "coarsened_grid x2": lambda: grid.coarsened_grid(2), + "plan from_grid_batch k2s2": lambda: ConvolutionPlan.from_grid_batch(2, 2, grid), + "plan from_grid_batch_transposed k2s2": lambda: ConvolutionPlan.from_grid_batch_transposed(2, 2, grid), + } + return {name: time_op(fn) for name, fn in ops.items()} + + +def bench_pyramid(grid: GridBatch, levels: int = 4) -> float: + """The generative-training pattern: rebuild the full conv_grid pyramid + per-level plans.""" + + def build(): + g = grid + plans = [] + for _ in range(levels): + plans.append(ConvolutionPlan.from_grid_batch(3, 1, g)) + plans.append(ConvolutionPlan.from_grid_batch(2, 2, g)) + g = g.conv_grid(kernel_size=2, stride=2) + return plans + + return time_op(build, warmup=2, iters=10) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--json", type=str, default=None, help="write results to this JSON file") + parser.add_argument("--gso", action="store_true", help="also run the issue #755 GSO repro") + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 32, 48]) + args = parser.parse_args() + + assert torch.cuda.is_available(), "this benchmark requires CUDA" + device = torch.cuda.get_device_name() + print(f"device: {device}") + + results = {"device": device, "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "batches": {}} + + for batch_size in args.batch_sizes: + grid = make_shell_batch(batch_size) + entry = bench_batch(grid) + entry["pyramid 4-level (8 plans + 3 conv_grids)"] = bench_pyramid(grid) + entry["total_voxels"] = int(grid.total_voxels) + results["batches"][batch_size] = entry + print(f"\nbatch_size={batch_size} (total voxels {grid.total_voxels}):") + for name, ms in entry.items(): + if isinstance(ms, float): + print(f" {name:45s} {ms:8.3f} ms") + + if args.gso: + from fvdb.utils.examples import load_gso_shoes + + meshes = load_gso_shoes(limit=16) + v = JaggedTensor([(m[0] - m[0].amin(0)) / m[0].amax() * 0.96 + 0.02 for m in meshes]) + f = JaggedTensor([m[1].int() for m in meshes]) + g = GridBatch.from_mesh(v, f, voxel_sizes=1 / 64, origins=0.0) + ms = time_op(lambda: ConvolutionPlan.from_grid_batch_transposed(2, 2, g)) + results["gso_from_grid_batch_transposed_k2s2_ms"] = ms + print(f"\nGSO shoes B=16: ConvolutionPlan.from_grid_batch_transposed(2,2,g): {ms:.3f} ms") + + if args.json: + with open(args.json, "w") as fp: + json.dump(results, fp, indent=2) + print(f"\nwrote {args.json}") + + +if __name__ == "__main__": + main() diff --git a/src/fvdb/GridBatchData.cu b/src/fvdb/GridBatchData.cu index 4e826e555..1fff02b6e 100644 --- a/src/fvdb/GridBatchData.cu +++ b/src/fvdb/GridBatchData.cu @@ -176,11 +176,14 @@ const torch::Tensor GridBatchData::voxelSizesTensor() const { torch::Tensor retTorch = torch::empty({batchSize(), 3}, torch::TensorOptions().dtype(torch::kFloat64)); + // Direct accessor fill: per-element tensor indexing costs 6 ATen dispatches per grid, which + // dominates plan-construction-time transform validation (issue #755). + auto acc = retTorch.accessor(); for (int64_t bi = 0; bi < batchSize(); bi += 1) { const auto voxSize = voxelSizeAt(bi); - retTorch[bi][0] = voxSize[0]; - retTorch[bi][1] = voxSize[1]; - retTorch[bi][2] = voxSize[2]; + acc[bi][0] = voxSize[0]; + acc[bi][1] = voxSize[1]; + acc[bi][2] = voxSize[2]; } return retTorch; } @@ -189,11 +192,12 @@ const torch::Tensor GridBatchData::voxelOriginsTensor() const { torch::Tensor retTorch = torch::empty({batchSize(), 3}, torch::TensorOptions().dtype(torch::kFloat64)); + auto acc = retTorch.accessor(); for (int64_t bi = 0; bi < batchSize(); bi += 1) { const auto voxOrigin = voxelOriginAt(bi); - retTorch[bi][0] = voxOrigin[0]; - retTorch[bi][1] = voxOrigin[1]; - retTorch[bi][2] = voxOrigin[2]; + acc[bi][0] = voxOrigin[0]; + acc[bi][1] = voxOrigin[1]; + acc[bi][2] = voxOrigin[2]; } return retTorch; } diff --git a/src/fvdb/detail/GridBatchDataFactory.cu b/src/fvdb/detail/GridBatchDataFactory.cu index 67b0c5f28..10dfdd94e 100644 --- a/src/fvdb/detail/GridBatchDataFactory.cu +++ b/src/fvdb/detail/GridBatchDataFactory.cu @@ -178,15 +178,18 @@ makeGridBatchData(nanovdb::GridHandle &&gridHdl, TORCH_CHECK(listIndices.numel() == 0 || listIndices.size(0) == (batchOffsets.size(0) - 1), "Invalid list indices when building grid"); - std::vector leafBatchIdxs; - leafBatchIdxs.reserve(batchSize); + // One repeat_interleave instead of a torch::full + torch::cat per member: the per-member + // version costs B+1 kernel dispatches on every grid construction (issue #755). Leaf counts + // are already host-side in hostMeta. + torch::Tensor leafCounts = + torch::empty({batchSize}, torch::TensorOptions().dtype(torch::kInt64)); + auto leafCountsAcc = leafCounts.accessor(); for (int64_t i = 0; i < batchSize; i += 1) { - leafBatchIdxs.push_back( - torch::full({hostMeta[i].mNumLeaves}, - static_cast(i), - torch::TensorOptions().dtype(fvdb::JIdxScalarType).device(device))); + leafCountsAcc[i] = hostMeta[i].mNumLeaves; } - torch::Tensor leafBatchIndices = torch::cat(leafBatchIdxs, 0); + torch::Tensor leafBatchIndices = torch::repeat_interleave( + torch::arange(batchSize, torch::TensorOptions().dtype(fvdb::JIdxScalarType).device(device)), + leafCounts.to(device)); auto gridHdlPtr = std::make_shared>(std::move(gridHdl)); diff --git a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu index fe357e683..f3c2e6fad 100644 --- a/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu +++ b/src/fvdb/detail/ops/BuildCoarseGridFromFine.cu @@ -11,12 +11,11 @@ #include #include #include -#include +#include #include #include #include -#include #include #include @@ -53,10 +52,11 @@ dispatchBuildCoarseGridFromFine(const GridBatchData &fineGridBatch, nanovdb::GridHandle coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch, const nanovdb::Coord &branchingFactor) { - // fvdb coarsening maps fine voxel f to floor(f / factor); NanoVDB's CoarsenGrid maps f to - // floor(f / 2) per pass (its coarsenComponent is exactly floor(n/2) for all n, and it unions - // each 2^3 fine block). So a uniform power-of-two factor is that many CoarsenGrid passes -- no - // coordinate list, no radix sort. Non-power-of-two / non-uniform factors keep the coord path. + // fvdb coarsening maps fine voxel f to floor(f / factor); a factor-2 coarsen pass maps f to + // floor(f / 2) (coarsenCoord is exactly floor(n/2) for all n, and each 2^3 fine block is + // unioned). So a uniform power-of-two factor is that many batched leaf-mask coarsen passes + // over the whole batch (BatchedTopologyBuilder) -- no coordinate list, no per-member builds. + // Non-power-of-two / non-uniform factors keep the coordinate path. const int nPasses = uniformPowerOfTwoLog2(branchingFactor); if (nPasses < 0) { JaggedTensor coords = ops::coarseIJKForFineGrid(fineGridBatch, branchingFactor); @@ -65,7 +65,6 @@ coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch, c10::cuda::CUDAGuard deviceGuard(fineGridBatch.device()); at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(fineGridBatch.device().index()); - TorchDeviceBuffer guide(0, fineGridBatch.device()); if (nPasses == 0) { // Coarsening factor 1 is the identity: the coarse grid == the fine grid. Compact the @@ -73,31 +72,12 @@ coarseGridHandleFromFineCUDA(const GridBatchData &fineGridBatch, return ops::contiguousGridHandle(fineGridBatch); } - std::vector> handles; - handles.reserve(fineGridBatch.batchSize()); - for (int64_t i = 0; i < fineGridBatch.batchSize(); i += 1) { - if (fineGridBatch.numVoxelsAt(i) == 0) { - handles.push_back(createEmptyGridHandle(fineGridBatch.device())); - continue; - } - - nanovdb::OnIndexGrid *grid = fineGridBatch.deviceGridPtrAt(i); - 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()); - op.setChecksum(nanovdb::CheckMode::Default); - op.setVerbose(0); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - handles.push_back(std::move(handle)); - } - - return handles.size() == 1 ? std::move(handles[0]) - : nanovdb::cuda::mergeGridHandles(handles, &guide); + // All batch members are coarsened together, one batched pass per factor of 2: a single output + // buffer, one stream synchronization per pass, no per-member builds or handle merging + // (issue #755). Empty members become valid empty grids inline. + const std::vector passes(nPasses, + batched::TopologyPassSpec::coarsen()); + return batched::batchedTopologyHandle(fineGridBatch, passes, stream.stream()); } template <> diff --git a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu index 441eb1a42..53045dd26 100644 --- a/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu +++ b/src/fvdb/detail/ops/BuildFineGridFromCoarse.cu @@ -14,12 +14,11 @@ #include #include #include -#include +#include #include #include #include -#include #include #include @@ -391,11 +390,11 @@ nanovdb::GridHandle fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl, const nanovdb::Coord &factor, const std::optional &mask) { - // fvdb subdivision maps coarse voxel c to the fine block c*factor + [0, factor-1]^3; NanoVDB's - // RefineGrid maps c to 2c + {0,1}^3 per pass. So a uniform power-of-two factor is that many - // RefineGrid passes -- leaf-mask morphology, no coordinate list, no radix sort. A per-coarse - // -voxel mask is applied by pruning the coarse grid to it first (PruneGrid), then refining. - // Non-power-of-two / non-uniform factors keep the coordinate path. + // fvdb subdivision maps coarse voxel c to the fine block c*factor + [0, factor-1]^3; a factor-2 + // refine pass maps c to 2c + {0,1}^3. So a uniform power-of-two factor is that many batched + // leaf-mask refine passes over the whole batch (BatchedTopologyBuilder) -- no coordinate list, + // no per-member builds. A per-coarse-voxel mask is applied by pruning the coarse grid to it + // first (PruneGrid), then refining. Non-power-of-two / non-uniform factors keep the coord path. const int nPasses = subdivUniformPowerOfTwoLog2(factor); if (nPasses < 0) { JaggedTensor coords = @@ -405,7 +404,6 @@ fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl, c10::cuda::CUDAGuard deviceGuard(coarseBatchHdl.device()); at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(coarseBatchHdl.device().index()); - TorchDeviceBuffer guide(0, coarseBatchHdl.device()); // The grid to refine is the coarse grid, or -- for masked subdivision -- the coarse grid pruned // to the selected voxels. pruneGrid keeps the coarse transform and canonical order; only its @@ -424,31 +422,12 @@ fineGridHandleFromCoarseCUDA(const GridBatchData &coarseBatchHdl, return ops::contiguousGridHandle(*src); } - std::vector> handles; - handles.reserve(src->batchSize()); - for (int64_t i = 0; i < src->batchSize(); i += 1) { - if (src->numVoxelsAt(i) == 0) { - handles.push_back(createEmptyGridHandle(coarseBatchHdl.device())); - continue; - } - - nanovdb::OnIndexGrid *grid = src->deviceGridPtrAt(i); - 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()); - op.setChecksum(nanovdb::CheckMode::Default); - op.setVerbose(0); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - handles.push_back(std::move(handle)); - } - - return handles.size() == 1 ? std::move(handles[0]) - : nanovdb::cuda::mergeGridHandles(handles, &guide); + // All batch members are refined together, one batched pass per factor of 2: a single output + // buffer, one stream synchronization per pass, no per-member builds or handle merging + // (issue #755). Empty members become valid empty grids inline. + const std::vector passes(nPasses, + batched::TopologyPassSpec::refine()); + return batched::batchedTopologyHandle(*src, passes, stream.stream()); } template <> diff --git a/src/fvdb/detail/ops/BuildGridForConv.cu b/src/fvdb/detail/ops/BuildGridForConv.cu index 2963adf2e..88e088a5d 100644 --- a/src/fvdb/detail/ops/BuildGridForConv.cu +++ b/src/fvdb/detail/ops/BuildGridForConv.cu @@ -12,11 +12,9 @@ #include #include #include -#include -#include +#include #include -#include #include #include @@ -368,26 +366,6 @@ countThenFillConvIJKForGrid(const GridBatchData &batchHdl, ConvolutionGeometry c outIJK, outIJKBIdx, batchHdl.jlidx(), batchHdl.batchSize()); } -// Applies fn(grid) -> handle to each logical non-empty batch item, preserves empties, then merges. -template -nanovdb::GridHandle -perItemGridHandle(const GridBatchData &base, const TorchDeviceBuffer &guide, PerGridFn &&fn) { - std::vector> handles; - handles.reserve(base.batchSize()); - for (int64_t i = 0; i < base.batchSize(); i += 1) { - if (base.numVoxelsAt(i) == 0) { - handles.push_back(createEmptyGridHandle(base.device())); - continue; - } - - nanovdb::OnIndexGrid *grid = base.deviceGridPtrAt(i); - TORCH_CHECK(grid, "Grid is null"); - handles.push_back(fn(grid)); - } - return handles.size() == 1 ? std::move(handles[0]) - : nanovdb::cuda::mergeGridHandles(handles, &guide); -} - template <> nanovdb::GridHandle dispatchBuildGridForConv(const GridBatchData &baseGridHdl, @@ -410,49 +388,30 @@ dispatchBuildGridForConv(const GridBatchData &baseGridHdl, } // At stride one the canonical forward support is source (+) - // [-paddingAfter, paddingBefore]^3. Realize that box with NanoVDB morphology. + // [-paddingAfter, paddingBefore]^3. Realize that box with batched leaf-mask morphology + // (Minkowski sums by boxes compose: odd K = symmetric [-1,1]^3 passes, even K = one-sided + // {-1,0}^3 / {0,1}^3 passes), all batch members per pass at once (issue #755). if (geometry.stride() == nanovdb::Coord(1) && isUniformKernel(geometry) && geometry.kernelSize()[0] > 1) { + using PassSpec = batched::TopologyPassSpec; gLastBuildGridForConvResourceStats = morphologyStats(baseGridHdl.totalVoxels(), geometry.kernelVolume()); c10::cuda::CUDAGuard deviceGuard(baseGridHdl.device()); at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(baseGridHdl.device().index()); - TorchDeviceBuffer guide(0, baseGridHdl.device()); - const int k = geometry.kernelSize()[0]; - - return perItemGridHandle(baseGridHdl, guide, [&](nanovdb::OnIndexGrid *grid) { - 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()); - op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); - op.setChecksum(nanovdb::CheckMode::Default); - op.setVerbose(0); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - } else { - for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( - grid, /*positiveOctant=*/false, stream.stream()); - op.setChecksum(nanovdb::CheckMode::Default); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( - grid, /*positiveOctant=*/true, stream.stream()); - op.setChecksum(nanovdb::CheckMode::Default); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - } - return handle; - }); + const int k = geometry.kernelSize()[0]; + + std::vector passes; + if (k % 2 == 1) { + passes.assign(geometry.paddingBefore()[0], + PassSpec::boxDilate(nanovdb::Coord(-1), nanovdb::Coord(1))); + } else { + passes.assign(geometry.paddingAfter()[0], + PassSpec::boxDilate(nanovdb::Coord(-1), nanovdb::Coord(0))); + passes.insert(passes.end(), + geometry.paddingBefore()[0], + PassSpec::boxDilate(nanovdb::Coord(0), nanovdb::Coord(1))); + } + return batched::batchedTopologyHandle(baseGridHdl, passes, stream.stream()); } // Shifted K=S uses one exact quotient per input. Other geometries use exact-M count/fill. diff --git a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu index c0cb63d36..5aefcdc4b 100644 --- a/src/fvdb/detail/ops/BuildGridForConvTranspose.cu +++ b/src/fvdb/detail/ops/BuildGridForConvTranspose.cu @@ -11,12 +11,9 @@ #include #include #include -#include -#include +#include #include -#include -#include #include #include @@ -178,34 +175,15 @@ convTransposeIJKForGrid(const GridBatchData &batchHdl, ConvolutionGeometry const outIJK, outIJKBIdx, batchHdl.jlidx(), batchHdl.batchSize()); } -// Applies fn(grid) -> handle to each non-empty batch item, empties -> empty grid, then merges. -template -static nanovdb::GridHandle -perItemGridHandle(const GridBatchData &base, const TorchDeviceBuffer &guide, PerGridFn &&fn) { - std::vector> handles; - handles.reserve(base.batchSize()); - for (int64_t i = 0; i < base.batchSize(); i += 1) { - if (base.numVoxelsAt(i) == 0) { - handles.push_back(createEmptyGridHandle(base.device())); - continue; - } - - nanovdb::OnIndexGrid *grid = base.deviceGridPtrAt(i); - TORCH_CHECK(grid, "Grid is null"); - handles.push_back(fn(grid)); - } - return handles.size() == 1 ? std::move(handles[0]) - : nanovdb::cuda::mergeGridHandles(handles, &guide); -} - template <> nanovdb::GridHandle dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl, const nanovdb::Coord &kernelSize, const nanovdb::Coord &stride) { + using PassSpec = batched::TopologyPassSpec; ConvolutionGeometry const geometry(kernelSize, stride); - // NanoVDB realizes the unshifted K=S={1,2} subdivision directly on leaf masks. + // The unshifted K=S={1,2} subdivision is realized directly on leaf masks (batched). // Shifted K=S geometries retain the canonical -paddingBefore phase in the fallback below. if (supportsLeafMaskSubdivision(geometry)) { return fineGridHandleFromCoarseCUDA(baseGridHdl, geometry.stride(), std::nullopt); @@ -213,70 +191,36 @@ dispatchBuildGridForConvTranspose(const GridBatchData &baseGridHdl c10::cuda::CUDAGuard deviceGuard(baseGridHdl.device()); at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(baseGridHdl.device().index()); - TorchDeviceBuffer guide(0, baseGridHdl.device()); // At stride one the canonical transpose support is source (+) - // [-paddingBefore, paddingAfter]^3. Realize that box with NanoVDB morphology. + // [-paddingBefore, paddingAfter]^3. Realize that box with batched leaf-mask morphology + // (Minkowski sums by boxes compose: odd K = symmetric [-1,1]^3 passes, even K = one-sided + // {-1,0}^3 / {0,1}^3 passes), all batch members per pass at once (issue #755). if (geometry.stride() == nanovdb::Coord(1) && isUniformKernel(geometry) && geometry.kernelSize()[0] > 1) { const int k = geometry.kernelSize()[0]; - return perItemGridHandle(baseGridHdl, guide, [&](nanovdb::OnIndexGrid *grid) { - 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()); - op.setOperation(nanovdb::tools::morphology::NN_FACE_EDGE_VERTEX); - op.setChecksum(nanovdb::CheckMode::Default); - op.setVerbose(0); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - } else { - for (int p = 0; p < geometry.paddingBefore()[0]; p += 1) { - morphology::PadGrid op( - grid, /*positiveOctant=*/false, stream.stream()); - op.setChecksum(nanovdb::CheckMode::Default); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - for (int p = 0; p < geometry.paddingAfter()[0]; p += 1) { - morphology::PadGrid op( - grid, /*positiveOctant=*/true, stream.stream()); - op.setChecksum(nanovdb::CheckMode::Default); - handle = op.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - grid = handle.deviceGrid(); - } - } - return handle; - }); + std::vector passes; + if (k % 2 == 1) { + passes.assign(geometry.paddingBefore()[0], + PassSpec::boxDilate(nanovdb::Coord(-1), nanovdb::Coord(1))); + } else { + passes.assign(geometry.paddingBefore()[0], + PassSpec::boxDilate(nanovdb::Coord(-1), nanovdb::Coord(0))); + passes.insert(passes.end(), + geometry.paddingAfter()[0], + PassSpec::boxDilate(nanovdb::Coord(0), nanovdb::Coord(1))); + } + return batched::batchedTopologyHandle(baseGridHdl, passes, stream.stream()); } // Fast path 3: stride 2, kernel 3 (the classic upsampling conv-transpose). The output is - // 2S (+) [-1,1]^3 (dstIjk = 2*srcIjk + offset, offset in [-1,1]^3). RefineGrid gives + // 2S (+) [-1,1]^3 (dstIjk = 2*srcIjk + offset, offset in [-1,1]^3). The refine pass gives // 2S (+) {0,1}^3, and one negative pad pass adds (+) {-1,0}^3, composing to (+) [-1,1]^3. 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()); - refineOp.setChecksum(nanovdb::CheckMode::Default); - refineOp.setVerbose(0); - nanovdb::GridHandle refined = refineOp.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - morphology::PadGrid padOp( - refined.deviceGrid(), - /*positiveOctant=*/false, - stream.stream()); - padOp.setChecksum(nanovdb::CheckMode::Default); - nanovdb::GridHandle handle = padOp.getHandle(guide); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return handle; - }); + const std::vector passes = { + PassSpec::refine(), PassSpec::boxDilate(nanovdb::Coord(-1), nanovdb::Coord(0))}; + return batched::batchedTopologyHandle(baseGridHdl, passes, stream.stream()); } if (isUnshiftedSubdivision(geometry)) { diff --git a/src/fvdb/detail/utils/nanovdb/BatchedTopologyBuilder.cuh b/src/fvdb/detail/utils/nanovdb/BatchedTopologyBuilder.cuh new file mode 100644 index 000000000..de3a469b8 --- /dev/null +++ b/src/fvdb/detail/utils/nanovdb/BatchedTopologyBuilder.cuh @@ -0,0 +1,1208 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +// Batched leaf-mask topology construction: builds ALL grids of a batch in one pass. +// +// NanoVDB's morphology builders (RefineGrid / CoarsenGrid via TopologyBuilder) are single-grid by +// construction: each getHandle() call performs several stream synchronizations (source-tree +// readback, host-side speculative root refinement, node-count readback for allocation sizing) and +// produces one single-grid handle, which fvdb then merges per batch member -- with another +// synchronization per grid inside nanovdb::cuda::mergeGridHandles. For generated-topology +// workloads that rebuild grids every training iteration (issue #755), that per-member fixed +// overhead -- not the topology size -- dominates wall clock and serializes the GPU. +// +// This header rebuilds the factor-2 refine (subdivision) and coarsen passes so one invocation +// covers the whole batch: +// +// emit one kernel over all source leaves of all members emits candidate output leaves as +// (root-tile sort key, in-tile node key, leaf origin, 512-bit activity mask) slots, +// segmented per grid; +// sort two stable segmented radix sorts put each grid's slots in canonical NanoVDB node +// order (root tiles by the PointsToGrid offset-shifted key, then upper/lower child +// offsets); invalid slots sort to the segment tails; +// dedup head-flag + scan passes derive the unique leaf/lower/upper nodes, their per-grid +// counts, and parent linkage (coarsen additionally OR-combines duplicate leaf masks); +// size ONE host synchronization reads back the per-grid node counts; per-grid byte offsets +// are computed on the host and a single output buffer is allocated; +// build batched kernels write every grid's GridData/TreeData/RootData (mGridIndex = g, +// mGridCount = B), root tiles, upper/lower/leaf nodes, leaf mOffset/mPrefixSum, and +// bounding boxes. Empty members become valid empty grids inline (no host proxy grids). +// +// The mask bit math is NanoVDB's own (RefineLeafMasksFunctor::refineMask / +// CoarsenLeafMasksFunctor::coarsenMask); the header/bbox/prefix-sum stages are transcriptions of +// tools::cuda::TopologyBuilder's functors with (gridIndex, localIndex) indexing. Checksums are +// disabled on the output, matching ops::contiguousGridHandle and mergeGridHandles behavior. +// +// Scratch is allocated through torch (the caching allocator) on the caller's current stream. The +// only stream synchronization per pass is the node-count readback that sizes the output buffer. + +#ifndef FVDB_DETAIL_UTILS_NANOVDB_BATCHEDTOPOLOGYBUILDER_CUH +#define FVDB_DETAIL_UTILS_NANOVDB_BATCHEDTOPOLOGYBUILDER_CUH + +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace fvdb { +namespace detail { +namespace batched { + +using BuildT = nanovdb::ValueOnIndex; +using GridT = nanovdb::NanoGrid; +using TreeT = nanovdb::NanoTree; +using RootT = nanovdb::NanoRoot; +using UpperT = nanovdb::NanoUpper; +using LowerT = nanovdb::NanoLower; +using LeafT = nanovdb::NanoLeaf; + +/// One batched topology pass over the whole batch. `Refine` and `Coarsen` are the factor-2 +/// subdivision / coarsening passes; `BoxDilate` is a Minkowski sum with the axis-aligned box +/// [boxLo, boxHi] (components in {-1,0,1}), covering NanoVDB's 26-neighbor DilateGrid +/// (boxLo=-1, boxHi=1) and fvdb's one-sided PadGrid octants ({-1,0}^3 and {0,1}^3). Larger +/// boxes compose from multiple passes (Minkowski sums by boxes compose). +struct TopologyPassSpec { + enum class Op { Refine, Coarsen, BoxDilate }; + Op op; + nanovdb::Coord boxLo{0}, boxHi{0}; // BoxDilate only + + static TopologyPassSpec + refine() { + return {Op::Refine}; + } + static TopologyPassSpec + coarsen() { + return {Op::Coarsen}; + } + static TopologyPassSpec + boxDilate(const nanovdb::Coord &lo, const nanovdb::Coord &hi) { + return {Op::BoxDilate, lo, hi}; + } +}; + +/// Device-pointer view of the B source grids of one pass. The grids may live anywhere (a +/// GridBatchData buffer -- including sliced/non-contiguous views -- or the output buffer of a +/// previous pass); only per-member device grid pointers and leaf counts are needed. +struct BatchedTopologySource { + std::vector grids; // per-member device grid pointers (host-side vector) + std::vector leafCounts; // per-member leaf counts (host-side) + torch::Device device{torch::kCUDA}; +}; + +/// Result of one batched pass: a single device buffer holding all B output grids back-to-back, +/// plus the host-side layout needed to wrap it in a GridHandle or feed the next pass. +struct BatchedTopologyResult { + TorchDeviceBuffer buffer; + std::vector gridByteOffsets; // B+1 cumulative byte offsets into `buffer` + std::vector leafCounts; // per-member output leaf counts +}; + +// --------------------------------------------------------------------------------------------- +// Device helpers +// --------------------------------------------------------------------------------------------- + +/// Sentinel tile key marking an invalid emission slot (sorts to the end of its segment). +static constexpr uint64_t kInvalidTileKey = ~uint64_t(0); + +/// Root-tile sort key: the offset-shifted encoding used by PointsToGrid and RefineGrid::refineRoot +/// (NOT the encoding stored in Tile::key by RootData::CoordToKey). Grids built by every other CUDA +/// path order their root tiles by this key, and elementwise topology comparisons rely on that +/// canonical order. +__device__ inline uint64_t +tileSortKey(const nanovdb::Coord &ijk) { + static constexpr int64_t kOffset = int64_t(1) << 31; + return (uint64_t(uint32_t(int64_t(ijk[2]) + kOffset) >> 12)) | + (uint64_t(uint32_t(int64_t(ijk[1]) + kOffset) >> 12) << 21) | + (uint64_t(uint32_t(int64_t(ijk[0]) + kOffset) >> 12) << 42); +} + +/// In-tile node key ordering leaves by (upper child offset, lower child offset) -- x-major at each +/// level, matching the breadth-first node order every NanoVDB builder produces. +__device__ inline uint32_t +nodeSortKey(const nanovdb::Coord &leafOrigin) { + return (UpperT::CoordToOffset(leafOrigin) << 12) | LowerT::CoordToOffset(leafOrigin); +} + +/// Index of the segment containing element i: upper_bound(offsets, i) - 1 over numSegments+1 +/// monotone offsets. +template +__device__ inline int32_t +findSegment(const OffsetT *__restrict__ offsets, int32_t numSegments, OffsetT i) { + int32_t lo = 0, hi = numSegments - 1; // segment index range + while (lo < hi) { + const int32_t mid = (lo + hi + 1) >> 1; + if (offsets[mid] <= i) { + lo = mid; + } else { + hi = mid - 1; + } + } + return lo; +} + +/// Per-grid node layout within the shared output buffer; mirrors TopologyBuilder::getBuffer. +struct GridNodes { + GridT *grid; + TreeT *tree; + RootT *root; + UpperT *upper; // first upper node + LowerT *lower; // first lower node + LeafT *leaf; // first leaf node + uint32_t numUpper, numLower, numLeaf; + uint64_t size; // total grid size in bytes +}; + +struct BuildDeviceArrays { + uint8_t *dstBase; + const uint64_t *gridByteOffsets; // [B+1] + const uint32_t *upperStart; // [B+1] global upper-node index offsets per grid + const uint32_t *lowerStart; // [B+1] + const uint32_t *leafStart; // [B+1] +}; + +__device__ inline GridNodes +gridNodes(const BuildDeviceArrays &a, int32_t g) { + GridNodes n; + n.numUpper = a.upperStart[g + 1] - a.upperStart[g]; + n.numLower = a.lowerStart[g + 1] - a.lowerStart[g]; + n.numLeaf = a.leafStart[g + 1] - a.leafStart[g]; + n.size = a.gridByteOffsets[g + 1] - a.gridByteOffsets[g]; + uint8_t *base = a.dstBase + a.gridByteOffsets[g]; + uint64_t offset = 0; + n.grid = reinterpret_cast(base); + offset += GridT::memUsage(); + n.tree = reinterpret_cast(base + offset); + offset += TreeT::memUsage(); + n.root = reinterpret_cast(base + offset); + offset += RootT::memUsage(n.numUpper); + n.upper = reinterpret_cast(base + offset); + offset += UpperT::memUsage() * uint64_t(n.numUpper); + n.lower = reinterpret_cast(base + offset); + offset += LowerT::memUsage() * uint64_t(n.numLower); + n.leaf = reinterpret_cast(base + offset); + return n; +} + +/// Byte size of one grid given its node counts; must match gridNodes()'s layout. +inline uint64_t +gridByteSize(uint32_t numUpper, uint32_t numLower, uint32_t numLeaf) { + return GridT::memUsage() + TreeT::memUsage() + RootT::memUsage(numUpper) + + UpperT::memUsage() * uint64_t(numUpper) + LowerT::memUsage() * uint64_t(numLower) + + LeafT::DataType::memUsage() * uint64_t(numLeaf); +} + +// --------------------------------------------------------------------------------------------- +// Emission kernels: one slot per candidate output leaf, segmented per grid +// --------------------------------------------------------------------------------------------- + +struct EmissionArrays { + uint64_t *tileKey; // [N] root-tile sort key; kInvalidTileKey for dead slots + uint32_t *nodeKey; // [N] (upperChildOffset << 12) | lowerChildOffset + int32_t *origin; // [N*3] output leaf origin + uint64_t *mask; // [N*8] output leaf 512-bit activity-mask contribution + const int32_t *segOffsets; // [B+1] per-grid slot ranges + const GridT *const *srcGrids; // [B] device pointers to the source grids + int32_t numSegments; + int32_t numSlots; +}; + +/// Refine (factor 2): slot = (source leaf, octant). A source leaf spanning [o, o+7]^3 produces up +/// to 8 fine leaves tiling [2o, 2o+15]^3; the fine mask of octant (bi,bj,bk) is the bit-doubled +/// 4^3 sub-mask of the source mask. Exactly one producer per fine leaf, so no deduplication. +static __global__ void +emitRefinedLeaves(EmissionArrays em) { + using RefineOp = nanovdb::util::morphology::cuda::RefineLeafMasksFunctor; + for (int32_t s = blockIdx.x * blockDim.x + threadIdx.x; s < em.numSlots; + s += gridDim.x * blockDim.x) { + const int32_t g = findSegment(em.segOffsets, em.numSegments, s); + const int32_t localSlot = s - em.segOffsets[g]; + const int32_t leafLocal = localSlot >> 3; + const int32_t octant = localSlot & 7; + const int32_t bi = (octant >> 2) & 1, bj = (octant >> 1) & 1, bk = octant & 1; + + const LeafT &srcLeaf = em.srcGrids[g]->tree().template getFirstNode<0>()[leafLocal]; + const uint64_t *srcWords = srcLeaf.valueMask().words(); + + // Extract the (bi,bj,bk) 4^3 sub-block of the source mask. refineMask() only reads the + // low-nibble pattern, so the stray high bits from the shift are harmless there, but they + // must be masked out for the occupancy test. + nanovdb::Mask<3> fineMask; + uint64_t *w = fineMask.words(); + uint64_t occupied = 0; + for (int i = 0; i < 4; ++i) { + w[i] = srcWords[i + bi * 4] >> (4 * bk + 32 * bj); + occupied |= w[i] & 0x000000000f0f0f0fUL; + } + if (!occupied) { + em.tileKey[s] = kInvalidTileKey; + continue; + } + RefineOp::refineMask(fineMask); + + const nanovdb::Coord srcOrigin = srcLeaf.origin(); + const nanovdb::Coord fineOrigin( + srcOrigin[0] * 2 + 8 * bi, srcOrigin[1] * 2 + 8 * bj, srcOrigin[2] * 2 + 8 * bk); + em.tileKey[s] = tileSortKey(fineOrigin); + em.nodeKey[s] = nodeSortKey(fineOrigin); + em.origin[s * 3] = fineOrigin[0]; + em.origin[s * 3 + 1] = fineOrigin[1]; + em.origin[s * 3 + 2] = fineOrigin[2]; + const uint64_t *fw = fineMask.words(); + for (int i = 0; i < 8; ++i) { + em.mask[s * 8 + i] = fw[i]; + } + } +} + +/// Coarsen (factor 2): slot = source leaf. A source leaf collapses to a 4^3 block placed by the +/// parity of the coarsened origin inside one output leaf; up to 8 source leaves contribute to the +/// same output leaf and are deduplicated (mask-OR) downstream. +static __global__ void +emitCoarsenedLeaves(EmissionArrays em) { + using CoarsenOp = nanovdb::util::morphology::cuda::CoarsenLeafMasksFunctor; + for (int32_t s = blockIdx.x * blockDim.x + threadIdx.x; s < em.numSlots; + s += gridDim.x * blockDim.x) { + const int32_t g = findSegment(em.segOffsets, em.numSegments, s); + const int32_t leafLocal = s - em.segOffsets[g]; + + const LeafT &srcLeaf = em.srcGrids[g]->tree().template getFirstNode<0>()[leafLocal]; + if (srcLeaf.valueMask().isOff()) { // leaves always have active voxels; defensive + em.tileKey[s] = kInvalidTileKey; + continue; + } + + const nanovdb::Coord coarseOrigin = + nanovdb::util::morphology::coarsenCoord(srcLeaf.origin()); + // The 4^3 coarse block starts at a multiple of 4, so it lies in a single 8-aligned leaf. + const nanovdb::Coord dstLeafOrigin( + coarseOrigin[0] & ~7, coarseOrigin[1] & ~7, coarseOrigin[2] & ~7); + const int bi = (coarseOrigin[0] & 7) ? 1 : 0; + const int bj = (coarseOrigin[1] & 7) ? 1 : 0; + const int bk = (coarseOrigin[2] & 7) ? 1 : 0; + + nanovdb::Mask<3> coarseMask = srcLeaf.valueMask(); + CoarsenOp::coarsenMask(coarseMask); + + uint64_t contribution[8] = {}; + const uint64_t *cw = coarseMask.words(); + for (int wi = 0; wi < 4; ++wi) { + contribution[wi + 4 * bi] = cw[wi] << (4 * bk + 32 * bj); + } + + em.tileKey[s] = tileSortKey(dstLeafOrigin); + em.nodeKey[s] = nodeSortKey(dstLeafOrigin); + em.origin[s * 3] = dstLeafOrigin[0]; + em.origin[s * 3 + 1] = dstLeafOrigin[1]; + em.origin[s * 3 + 2] = dstLeafOrigin[2]; + for (int i = 0; i < 8; ++i) { + em.mask[s * 8 + i] = contribution[i]; + } + } +} + +/// Shifts a leaf-mask word along z by s in [-7,7], dropping bits that leave the leaf. +/// (Mask<3> word layout: word index = local x, byte within word = local y, bit within byte = z.) +__device__ inline uint64_t +shiftWordZ(uint64_t w, int s) { + if (s > 0) { + return (w & (0x0101010101010101UL * (0xffu >> s))) << s; + } + if (s < 0) { + return (w >> -s) & (0x0101010101010101UL * (0xffu >> -s)); + } + return w; +} + +/// Shifts a leaf-mask word along y by s in [-7,7]; byte granularity, natural truncation. +__device__ inline uint64_t +shiftWordY(uint64_t w, int s) { + return s > 0 ? w << (8 * s) : (s < 0 ? w >> (-8 * s) : w); +} + +/// BoxDilate: slot = (source leaf, target neighbor-leaf offset db). The contribution of a source +/// leaf S to the target leaf at S.origin + 8*db is U_{o in [boxLo,boxHi]} shift(S, o - 8*db), +/// which factorizes per axis (Minkowski sums of axis-aligned boxes compose); shifts with any +/// component of magnitude >= 8 vanish, so per axis: db==0 keeps all box offsets, db==+-1 keeps +/// only o==+-1 (shift of -+7). Up to 2^3 source leaves contribute to one target leaf per +/// one-sided pass (3^3 for the full box), deduplicated (mask-OR) downstream like coarsen. +static __global__ void +emitBoxDilatedLeaves(EmissionArrays em, nanovdb::Coord boxLo, nanovdb::Coord boxHi) { + // Per-axis target-offset ranges and fanout. + int dbLo[3], fan[3]; + for (int a = 0; a < 3; ++a) { + dbLo[a] = boxLo[a] < 0 ? -1 : 0; + fan[a] = (boxHi[a] > 0 ? 1 : 0) - dbLo[a] + 1; + } + const int32_t fanout = fan[0] * fan[1] * fan[2]; + + for (int32_t s = blockIdx.x * blockDim.x + threadIdx.x; s < em.numSlots; + s += gridDim.x * blockDim.x) { + const int32_t g = findSegment(em.segOffsets, em.numSegments, s); + const int32_t localSlot = s - em.segOffsets[g]; + const int32_t leafLocal = localSlot / fanout; + int32_t t = localSlot % fanout; + const int dbz = dbLo[2] + t % fan[2]; + t /= fan[2]; + const int dby = dbLo[1] + t % fan[1]; + const int dbx = dbLo[0] + t / fan[1]; + + const LeafT &srcLeaf = em.srcGrids[g]->tree().template getFirstNode<0>()[leafLocal]; + const uint64_t *srcWords = srcLeaf.valueMask().words(); + + // z-stage, then y-stage (per word), then x-stage (word permutation). + uint64_t wz[8], wy[8], out[8]; + for (int i = 0; i < 8; ++i) { + uint64_t acc = 0; + for (int oz = boxLo[2]; oz <= boxHi[2]; ++oz) { + const int sz = oz - 8 * dbz; + if (sz > -8 && sz < 8) { + acc |= shiftWordZ(srcWords[i], sz); + } + } + wz[i] = acc; + } + for (int i = 0; i < 8; ++i) { + uint64_t acc = 0; + for (int oy = boxLo[1]; oy <= boxHi[1]; ++oy) { + const int sy = oy - 8 * dby; + if (sy > -8 && sy < 8) { + acc |= shiftWordY(wz[i], sy); + } + } + wy[i] = acc; + } + uint64_t occupied = 0; + for (int x = 0; x < 8; ++x) { + uint64_t acc = 0; + for (int ox = boxLo[0]; ox <= boxHi[0]; ++ox) { + const int sx = ox - 8 * dbx; // out[x] gathers wy[x - sx] + const int src = x - sx; + if (src >= 0 && src < 8) { + acc |= wy[src]; + } + } + out[x] = acc; + occupied |= acc; + } + if (!occupied) { + em.tileKey[s] = kInvalidTileKey; + continue; + } + + const nanovdb::Coord srcOrigin = srcLeaf.origin(); + const nanovdb::Coord dstLeafOrigin( + srcOrigin[0] + 8 * dbx, srcOrigin[1] + 8 * dby, srcOrigin[2] + 8 * dbz); + em.tileKey[s] = tileSortKey(dstLeafOrigin); + em.nodeKey[s] = nodeSortKey(dstLeafOrigin); + em.origin[s * 3] = dstLeafOrigin[0]; + em.origin[s * 3 + 1] = dstLeafOrigin[1]; + em.origin[s * 3 + 2] = dstLeafOrigin[2]; + for (int i = 0; i < 8; ++i) { + em.mask[s * 8 + i] = out[i]; + } + } +} + +// --------------------------------------------------------------------------------------------- +// Sort / dedup kernels +// --------------------------------------------------------------------------------------------- + +static __global__ void +iotaKernel(uint32_t *out, int32_t n) { + for (int32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) { + out[i] = uint32_t(i); + } +} + +static __global__ void +gatherKeys64(const uint64_t *__restrict__ keys, + const uint32_t *__restrict__ perm, + uint64_t *__restrict__ out, + int32_t n) { + for (int32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) { + out[i] = keys[perm[i]]; + } +} + +struct SortedArrays { + const uint64_t *tileKey; // [N] canonical order, invalid slots at segment tails + uint32_t *nodeKey; // [N] + int32_t *origin; // [N*3] + uint64_t *mask; // [N*8] +}; + +static __global__ void +gatherSorted(EmissionArrays em, const uint32_t *__restrict__ perm, SortedArrays out) { + for (int32_t j = blockIdx.x * blockDim.x + threadIdx.x; j < em.numSlots; + j += gridDim.x * blockDim.x) { + const uint32_t s = perm[j]; + out.nodeKey[j] = em.nodeKey[s]; + out.origin[j * 3] = em.origin[s * 3]; + out.origin[j * 3 + 1] = em.origin[s * 3 + 1]; + out.origin[j * 3 + 2] = em.origin[s * 3 + 2]; + for (int i = 0; i < 8; ++i) { + out.mask[j * 8 + i] = em.mask[s * 8 + i]; + } + } +} + +/// Marks the first slot of every distinct leaf / lower / upper node within each grid's segment. +static __global__ void +computeHeadFlags(const uint64_t *__restrict__ tileKey, + const uint32_t *__restrict__ nodeKey, + const int32_t *__restrict__ segOffsets, + int32_t numSegments, + int32_t n, + uint32_t *__restrict__ leafFlag, + uint32_t *__restrict__ lowerFlag, + uint32_t *__restrict__ upperFlag) { + for (int32_t j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x) { + uint32_t lf = 0, wf = 0, uf = 0; + if (tileKey[j] != kInvalidTileKey) { + const int32_t g = findSegment(segOffsets, numSegments, j); + const bool first = (j == segOffsets[g]); + const bool newTile = first || tileKey[j] != tileKey[j - 1]; + const bool newLower = newTile || (nodeKey[j] >> 12) != (nodeKey[j - 1] >> 12); + const bool newLeaf = newLower || nodeKey[j] != nodeKey[j - 1]; + lf = newLeaf ? 1u : 0u; + wf = newLower ? 1u : 0u; + uf = newTile ? 1u : 0u; + } + leafFlag[j] = lf; + lowerFlag[j] = wf; + upperFlag[j] = uf; + } +} + +/// Per-grid global node-index offsets, read off the inclusive rank scans at segment boundaries. +static __global__ void +segmentNodeOffsets(const uint32_t *__restrict__ leafRank, + const uint32_t *__restrict__ lowerRank, + const uint32_t *__restrict__ upperRank, + const int32_t *__restrict__ segOffsets, + int32_t numSegments, + uint32_t *__restrict__ leafStart, + uint32_t *__restrict__ lowerStart, + uint32_t *__restrict__ upperStart) { + const int32_t g = blockIdx.x * blockDim.x + threadIdx.x; + if (g > numSegments) { + return; + } + const int32_t boundary = segOffsets[g]; // slots before this boundary belong to grids < g + leafStart[g] = boundary == 0 ? 0 : leafRank[boundary - 1]; + lowerStart[g] = boundary == 0 ? 0 : lowerRank[boundary - 1]; + upperStart[g] = boundary == 0 ? 0 : upperRank[boundary - 1]; +} + +/// Records, for every unique node, its head slot and its parent's global node index. +static __global__ void +scatterNodeTables(const uint32_t *__restrict__ leafFlag, + const uint32_t *__restrict__ lowerFlag, + const uint32_t *__restrict__ upperFlag, + const uint32_t *__restrict__ leafRank, + const uint32_t *__restrict__ lowerRank, + const uint32_t *__restrict__ upperRank, + int32_t n, + uint32_t *__restrict__ leafHeadSlot, + uint32_t *__restrict__ leafParent, + uint32_t *__restrict__ lowerHeadSlot, + uint32_t *__restrict__ lowerParent, + uint32_t *__restrict__ upperHeadSlot) { + for (int32_t j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x) { + if (leafFlag[j]) { + const uint32_t leafIdx = leafRank[j] - 1; + leafHeadSlot[leafIdx] = uint32_t(j); + leafParent[leafIdx] = lowerRank[j] - 1; + } + if (lowerFlag[j]) { + const uint32_t lowerIdx = lowerRank[j] - 1; + lowerHeadSlot[lowerIdx] = uint32_t(j); + lowerParent[lowerIdx] = upperRank[j] - 1; + } + if (upperFlag[j]) { + upperHeadSlot[upperRank[j] - 1] = uint32_t(j); + } + } +} + +/// Coarsen only: OR every duplicate slot's mask contribution into its unique leaf's head slot. +static __global__ void +combineDuplicateMasks(const uint64_t *__restrict__ tileKey, + const uint32_t *__restrict__ leafFlag, + const uint32_t *__restrict__ leafRank, + const uint32_t *__restrict__ leafHeadSlot, + int32_t n, + uint64_t *__restrict__ mask) { + for (int32_t j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x) { + if (tileKey[j] == kInvalidTileKey || leafFlag[j]) { + continue; + } + const uint32_t head = leafHeadSlot[leafRank[j] - 1]; + for (int i = 0; i < 8; ++i) { + const uint64_t w = mask[j * 8 + i]; + if (w) { + nanovdb::util::atomicOr(&mask[head * 8 + i], w); + } + } + } +} + +// --------------------------------------------------------------------------------------------- +// Build kernels (after the single sizing readback and output allocation) +// --------------------------------------------------------------------------------------------- + +/// Copies every source grid's GridData header to its output position (preserving grid name, map, +/// voxel size); the fields reset by initGridTreeRoot are rewritten afterwards. +static __global__ void +copyGridData(const GridT *const *__restrict__ srcGrids, BuildDeviceArrays a, int32_t numGrids) { + constexpr int32_t kWords = int32_t(sizeof(nanovdb::GridData) / sizeof(uint64_t)); + const int32_t totalWords = numGrids * kWords; + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < totalWords; + t += gridDim.x * blockDim.x) { + const int32_t g = t / kWords; + const int32_t w = t % kWords; + const uint64_t *src = reinterpret_cast(srcGrids[g]); + uint64_t *dst = reinterpret_cast(a.dstBase + a.gridByteOffsets[g]); + dst[w] = src[w]; + } +} + +/// Per-grid transcription of tools::cuda::topology::detail::BuildGridTreeRootFunctor with +/// mGridIndex/mGridCount/mGridSize set for a multi-grid buffer. +static __global__ void +initGridTreeRoot(BuildDeviceArrays a, int32_t numGrids) { + const int32_t g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= numGrids) { + return; + } + GridNodes n = gridNodes(a, g); + + auto &root = *n.root; + root.mTableSize = n.numUpper; + root.mBackground = RootT::ValueType(0); + root.mMinimum = root.mMaximum = RootT::ValueType(0); + root.mAverage = root.mStdDevi = RootT::FloatType(0); + root.mBBox = nanovdb::CoordBBox(); + + auto &tree = *n.tree; + tree.setRoot(&root); + if (n.numUpper) { + tree.setFirstNode(n.upper); + tree.setFirstNode(n.lower); + tree.setFirstNode(n.leaf); + } else { + tree.setFirstNode(nullptr); + tree.setFirstNode(nullptr); + tree.setFirstNode(nullptr); + } + tree.mNodeCount[2] = n.numUpper; + tree.mNodeCount[1] = n.numLower; + tree.mNodeCount[0] = n.numLeaf; + tree.mVoxelCount = 0; // set by finalizeGrids once leaf masks are in place + tree.mTileCount[2] = tree.mTileCount[1] = tree.mTileCount[0] = 0; + + auto &grid = *n.grid; + grid.mChecksum.disable(); + grid.mFlags.initMask({nanovdb::GridFlags::IsBreadthFirst}); + grid.mGridIndex = uint32_t(g); + grid.mGridCount = uint32_t(numGrids); + grid.mGridSize = n.size; + grid.mWorldBBox = nanovdb::Vec3dBBox(); + grid.mVoxelSize = grid.mMap.getVoxelSize(); + grid.mBlindMetadataOffset = n.size; + grid.mBlindMetadataCount = 0u; + grid.mData1 = 1u; +} + +/// One thread per unique upper node: writes its root tile (in canonical tile order) and preamble. +static __global__ void +buildUpperNodes(BuildDeviceArrays a, + SortedArrays sorted, + const uint32_t *__restrict__ upperHeadSlot, + int32_t numGrids, + int32_t totalUpper) { + for (int32_t u = blockIdx.x * blockDim.x + threadIdx.x; u < totalUpper; + u += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.upperStart, numGrids, uint32_t(u)); + const GridNodes n = gridNodes(a, g); + const uint32_t local = uint32_t(u) - a.upperStart[g]; + const uint32_t slot = upperHeadSlot[u]; + const nanovdb::Coord tileOrigin(sorted.origin[slot * 3] & ~4095, + sorted.origin[slot * 3 + 1] & ~4095, + sorted.origin[slot * 3 + 2] & ~4095); + UpperT &upper = n.upper[local]; + n.root->tile(local)->setChild(tileOrigin, &upper, n.root->data()); + upper.mBBox = nanovdb::CoordBBox(); + upper.mFlags = (uint64_t)nanovdb::GridFlags::HasBBox; + } +} + +/// One thread per unique lower node: links it under its upper node and writes its preamble. +static __global__ void +buildLowerNodes(BuildDeviceArrays a, + SortedArrays sorted, + const uint32_t *__restrict__ lowerHeadSlot, + const uint32_t *__restrict__ lowerParent, + int32_t numGrids, + int32_t totalLower) { + for (int32_t l = blockIdx.x * blockDim.x + threadIdx.x; l < totalLower; + l += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.lowerStart, numGrids, uint32_t(l)); + const GridNodes n = gridNodes(a, g); + const uint32_t local = uint32_t(l) - a.lowerStart[g]; + const uint32_t slot = lowerHeadSlot[l]; + const uint32_t upperOff = sorted.nodeKey[slot] >> 12; + UpperT &upper = n.upper[lowerParent[l] - a.upperStart[g]]; + LowerT &lower = n.lower[local]; + upper.mChildMask.setOnAtomic(upperOff); + upper.setChild(upperOff, &lower); + lower.mBBox = nanovdb::CoordBBox(); + lower.mFlags = (uint64_t)nanovdb::GridFlags::HasBBox; + } +} + +/// One thread per unique leaf: links it under its lower node, writes its activity mask, per-leaf +/// prefix sums, and voxel count. +static __global__ void +buildLeafNodes(BuildDeviceArrays a, + SortedArrays sorted, + const uint32_t *__restrict__ leafHeadSlot, + const uint32_t *__restrict__ leafParent, + int32_t numGrids, + int32_t totalLeaf, + uint64_t *__restrict__ voxelCounts) { // [totalLeaf+1]; element 0 stays 0 + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < totalLeaf; + t += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.leafStart, numGrids, uint32_t(t)); + const GridNodes n = gridNodes(a, g); + const uint32_t local = uint32_t(t) - a.leafStart[g]; + const uint32_t slot = leafHeadSlot[t]; + const uint32_t lowerOff = sorted.nodeKey[slot] & 0xFFFu; + LowerT &lower = n.lower[leafParent[t] - a.lowerStart[g]]; + LeafT &leaf = n.leaf[local]; + lower.mChildMask.setOnAtomic(lowerOff); + lower.setChild(lowerOff, &leaf); + leaf.mBBoxMin = nanovdb::Coord( + sorted.origin[slot * 3], sorted.origin[slot * 3 + 1], sorted.origin[slot * 3 + 2]); + leaf.mFlags = uint8_t(nanovdb::GridFlags::HasBBox); + + uint64_t *dstWords = leaf.mValueMask.words(); + for (int i = 0; i < 8; ++i) { + dstWords[i] = sorted.mask[slot * 8 + i]; + } + + // Per-leaf voxel count and the 9-bit encoded intra-leaf prefix sums (transcribed from + // TopologyBuilder's UpdateLeafVoxelCountsAndPrefixSumFunctor). + uint64_t prefixSum = 0, sum = nanovdb::util::countOn(dstWords[0]); + prefixSum = sum; + for (int wi = 1; wi < 7; ++wi) { + sum += nanovdb::util::countOn(dstWords[wi]); + prefixSum |= sum << (9 * wi); + } + sum += nanovdb::util::countOn(dstWords[7]); + voxelCounts[t + 1] = sum; + leaf.mPrefixSum = prefixSum; + } +} + +/// One thread per leaf: 1-based per-grid value offsets from the global voxel-count scan. +static __global__ void +setLeafOffsets(BuildDeviceArrays a, + const uint64_t *__restrict__ voxelScan, // [totalLeaf+1] inclusive scan, [0] == 0 + int32_t numGrids, + int32_t totalLeaf) { + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < totalLeaf; + t += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.leafStart, numGrids, uint32_t(t)); + const GridNodes n = gridNodes(a, g); + const uint32_t local = uint32_t(t) - a.leafStart[g]; + n.leaf[local].mOffset = voxelScan[t] - voxelScan[a.leafStart[g]] + 1; + } +} + +static __global__ void +propagateLeafBBox(BuildDeviceArrays a, + const uint32_t *__restrict__ leafParent, + int32_t numGrids, + int32_t totalLeaf) { + for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < totalLeaf; + t += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.leafStart, numGrids, uint32_t(t)); + const GridNodes n = gridNodes(a, g); + LeafT &leaf = n.leaf[uint32_t(t) - a.leafStart[g]]; + LowerT &lower = n.lower[leafParent[t] - a.lowerStart[g]]; + leaf.updateBBox(); + lower.mBBox.expandAtomic(leaf.bbox()); + } +} + +static __global__ void +propagateLowerBBox(BuildDeviceArrays a, + const uint32_t *__restrict__ lowerParent, + int32_t numGrids, + int32_t totalLower) { + for (int32_t l = blockIdx.x * blockDim.x + threadIdx.x; l < totalLower; + l += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.lowerStart, numGrids, uint32_t(l)); + const GridNodes n = gridNodes(a, g); + LowerT &lower = n.lower[uint32_t(l) - a.lowerStart[g]]; + UpperT &upper = n.upper[lowerParent[l] - a.upperStart[g]]; + upper.mBBox.expandAtomic(lower.bbox()); + } +} + +static __global__ void +propagateUpperBBox(BuildDeviceArrays a, int32_t numGrids, int32_t totalUpper) { + for (int32_t u = blockIdx.x * blockDim.x + threadIdx.x; u < totalUpper; + u += gridDim.x * blockDim.x) { + const int32_t g = findSegment(a.upperStart, numGrids, uint32_t(u)); + const GridNodes n = gridNodes(a, g); + n.root->mBBox.expandAtomic(n.upper[uint32_t(u) - a.upperStart[g]].bbox()); + } +} + +/// Per-grid epilogue: voxel counts, index-space -> world bbox, HasBBox flag. +static __global__ void +finalizeGrids(BuildDeviceArrays a, const uint64_t *__restrict__ voxelScan, int32_t numGrids) { + const int32_t g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= numGrids) { + return; + } + const GridNodes n = gridNodes(a, g); + if (n.numLeaf == 0) { // empty grid: keep the empty bbox and defaults + return; + } + const uint64_t voxelCount = voxelScan[a.leafStart[g + 1]] - voxelScan[a.leafStart[g]]; + n.tree->mVoxelCount = voxelCount; + n.grid->mData1 = voxelCount + 1; + + nanovdb::CoordBBox bbox = n.root->mBBox; + bbox.max() += nanovdb::Coord(1); + n.grid->mFlags.setMaskOn(nanovdb::GridFlags::HasBBox); + n.grid->mWorldBBox = bbox.transform(n.grid->data()->mMap); +} + +// --------------------------------------------------------------------------------------------- +// Host driver +// --------------------------------------------------------------------------------------------- + +inline BatchedTopologySource +sourceFromGridBatch(const GridBatchData &batch) { + BatchedTopologySource src; + src.device = batch.device(); + const int64_t batchSize = batch.batchSize(); + src.grids.reserve(batchSize); + src.leafCounts.reserve(batchSize); + for (int64_t i = 0; i < batchSize; ++i) { + src.grids.push_back(batch.deviceGridPtrAt(i)); + src.leafCounts.push_back(batch.numLeavesAt(i)); + } + return src; +} + +inline BatchedTopologySource +sourceFromResult(const BatchedTopologyResult &result, const torch::Device &device) { + BatchedTopologySource src; + src.device = device; + const size_t batchSize = result.leafCounts.size(); + src.grids.reserve(batchSize); + src.leafCounts = result.leafCounts; + const uint8_t *base = result.buffer.deviceData(); + for (size_t i = 0; i < batchSize; ++i) { + src.grids.push_back(reinterpret_cast(base + result.gridByteOffsets[i])); + } + return src; +} + +/// Runs one batched topology pass over all grids of `src`. +/// One cudaStreamSynchronize total (the node-count readback that sizes the output allocation). +inline BatchedTopologyResult +runBatchedTopologyPass(const BatchedTopologySource &src, + const TopologyPassSpec &pass, + cudaStream_t stream) { + TORCH_CHECK(src.device.is_cuda(), "batched topology passes require a CUDA device"); + const int32_t numGrids = int32_t(src.grids.size()); + TORCH_CHECK(numGrids > 0, "batched topology pass requires at least one grid"); + int32_t fanout = 1; + switch (pass.op) { + case TopologyPassSpec::Op::Refine: fanout = 8; break; + case TopologyPassSpec::Op::Coarsen: fanout = 1; break; + case TopologyPassSpec::Op::BoxDilate: + for (int a = 0; a < 3; ++a) { + TORCH_CHECK(pass.boxLo[a] >= -1 && pass.boxLo[a] <= 0 && pass.boxHi[a] >= 0 && + pass.boxHi[a] <= 1, + "BoxDilate pass components must be in {-1,0} / {0,1}"); + fanout *= (pass.boxHi[a] > 0 ? 1 : 0) - (pass.boxLo[a] < 0 ? -1 : 0) + 1; + } + break; + } + + // Per-grid emission slot ranges (host-known: leaf counts x fanout; no sync needed). + std::vector segOffsetsHost(numGrids + 1, 0); + int64_t total = 0; + for (int32_t g = 0; g < numGrids; ++g) { + segOffsetsHost[g] = int32_t(total); + total += src.leafCounts[g] * fanout; + } + TORCH_CHECK(total <= std::numeric_limits::max(), + "batched topology pass: emission slot count ", + total, + " exceeds int32 range"); + segOffsetsHost[numGrids] = int32_t(total); + const int32_t numSlots = int32_t(total); + const int64_t allocSlots = std::max(numSlots, 1); // torch::empty({0}) yields a null data_ptr + + const auto byteOpts = torch::TensorOptions().dtype(torch::kUInt8).device(src.device); + const auto i32Opts = byteOpts.dtype(torch::kInt32); + const auto i64Opts = byteOpts.dtype(torch::kInt64); + + constexpr int32_t kThreads = 256; + const auto numBlocks = [](int32_t n) { return std::max((n + kThreads - 1) / kThreads, 1); }; + const auto u32 = [](torch::Tensor &t) { + return reinterpret_cast(t.data_ptr()); + }; + const auto u64 = [](torch::Tensor &t) { + return reinterpret_cast(t.data_ptr()); + }; + + // Runs a cub function twice: once to query scratch size, once for real. + const auto callCub = [&](auto fn) { + size_t tempBytes = 0; + fn(nullptr, tempBytes); + torch::Tensor tempStorage = + torch::empty({std::max(int64_t(tempBytes), 1)}, byteOpts); + fn(tempStorage.data_ptr(), tempBytes); + }; + + // Small host-side staging uploaded once per pass: segment offsets + source grid pointers. + // (Pageable H2D copies are synchronous with the host but do not synchronize the stream.) + torch::Tensor segOffsetsDev = torch::empty({numGrids + 1}, i32Opts); + torch::Tensor srcGridsDev = torch::empty({numGrids}, i64Opts); + C10_CUDA_CHECK(cudaMemcpyAsync(segOffsetsDev.data_ptr(), + segOffsetsHost.data(), + sizeof(int32_t) * (numGrids + 1), + cudaMemcpyHostToDevice, + stream)); + static_assert(sizeof(const GridT *) == sizeof(int64_t)); + C10_CUDA_CHECK(cudaMemcpyAsync(srcGridsDev.data_ptr(), + src.grids.data(), + sizeof(const GridT *) * numGrids, + cudaMemcpyHostToDevice, + stream)); + + // --- Emit candidate output leaves. --- + torch::Tensor emTileKey = torch::empty({allocSlots}, i64Opts); + torch::Tensor emNodeKey = torch::empty({allocSlots}, i32Opts); + torch::Tensor emOrigin = torch::empty({allocSlots * 3}, i32Opts); + torch::Tensor emMask = torch::empty({allocSlots * 8}, i64Opts); + + EmissionArrays em; + em.tileKey = u64(emTileKey); + em.nodeKey = u32(emNodeKey); + em.origin = emOrigin.data_ptr(); + em.mask = u64(emMask); + em.segOffsets = segOffsetsDev.data_ptr(); + em.srcGrids = reinterpret_cast(srcGridsDev.data_ptr()); + em.numSegments = numGrids; + em.numSlots = numSlots; + + if (numSlots > 0) { + switch (pass.op) { + case TopologyPassSpec::Op::Refine: + emitRefinedLeaves<<>>(em); + break; + case TopologyPassSpec::Op::Coarsen: + emitCoarsenedLeaves<<>>(em); + break; + case TopologyPassSpec::Op::BoxDilate: + emitBoxDilatedLeaves<<>>( + em, pass.boxLo, pass.boxHi); + break; + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + // --- Canonical-order segmented sort: stable by node key, then stable by tile key. --- + torch::Tensor permA = torch::empty({allocSlots}, i32Opts); + torch::Tensor permB = torch::empty({allocSlots}, i32Opts); + torch::Tensor keysTmp = torch::empty({allocSlots}, i32Opts); + torch::Tensor keys64A = torch::empty({allocSlots}, i64Opts); + torch::Tensor keys64B = torch::empty({allocSlots}, i64Opts); + + if (numSlots > 0) { + iotaKernel<<>>(u32(permA), numSlots); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // Pass 1: sort by the 27-bit in-tile node key. + callCub([&](void *temp, size_t &bytes) { + C10_CUDA_CHECK(cub::DeviceSegmentedRadixSort::SortPairs(temp, + bytes, + em.nodeKey, + u32(keysTmp), + u32(permA), + u32(permB), + numSlots, + numGrids, + em.segOffsets, + em.segOffsets + 1, + 0, + 27, + stream)); + }); + + // Pass 2: stable sort by the 64-bit tile key (invalid slots carry ~0 and sort last). + gatherKeys64<<>>( + em.tileKey, u32(permB), u64(keys64A), numSlots); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + callCub([&](void *temp, size_t &bytes) { + C10_CUDA_CHECK(cub::DeviceSegmentedRadixSort::SortPairs(temp, + bytes, + u64(keys64A), + u64(keys64B), + u32(permB), + u32(permA), + numSlots, + numGrids, + em.segOffsets, + em.segOffsets + 1, + 0, + 64, + stream)); + }); + } + // From here on: permA = canonical-order permutation, keys64B = sorted tile keys. + + torch::Tensor sortedNodeKey = torch::empty({allocSlots}, i32Opts); + torch::Tensor sortedOrigin = torch::empty({allocSlots * 3}, i32Opts); + torch::Tensor sortedMask = torch::empty({allocSlots * 8}, i64Opts); + + SortedArrays sorted; + sorted.tileKey = u64(keys64B); + sorted.nodeKey = u32(sortedNodeKey); + sorted.origin = sortedOrigin.data_ptr(); + sorted.mask = u64(sortedMask); + + // --- Dedup: head flags, global node ranks, per-grid offsets, parent linkage. --- + torch::Tensor leafFlag = torch::empty({allocSlots}, i32Opts); + torch::Tensor lowerFlag = torch::empty({allocSlots}, i32Opts); + torch::Tensor upperFlag = torch::empty({allocSlots}, i32Opts); + torch::Tensor leafRank = torch::empty({allocSlots}, i32Opts); + torch::Tensor lowerRank = torch::empty({allocSlots}, i32Opts); + torch::Tensor upperRank = torch::empty({allocSlots}, i32Opts); + + torch::Tensor leafHeadSlot = torch::empty({allocSlots}, i32Opts); + torch::Tensor leafParent = torch::empty({allocSlots}, i32Opts); + torch::Tensor lowerHeadSlot = torch::empty({allocSlots}, i32Opts); + torch::Tensor lowerParent = torch::empty({allocSlots}, i32Opts); + torch::Tensor upperHeadSlot = torch::empty({allocSlots}, i32Opts); + + if (numSlots > 0) { + gatherSorted<<>>(em, u32(permA), sorted); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + computeHeadFlags<<>>(sorted.tileKey, + sorted.nodeKey, + em.segOffsets, + numGrids, + numSlots, + u32(leafFlag), + u32(lowerFlag), + u32(upperFlag)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const std::pair scans[3] = { + {&leafFlag, &leafRank}, {&lowerFlag, &lowerRank}, {&upperFlag, &upperRank}}; + for (const auto &[flags, ranks]: scans) { + callCub([&](void *temp, size_t &bytes) { + C10_CUDA_CHECK(cub::DeviceScan::InclusiveSum( + temp, bytes, u32(*flags), u32(*ranks), numSlots, stream)); + }); + } + + scatterNodeTables<<>>(u32(leafFlag), + u32(lowerFlag), + u32(upperFlag), + u32(leafRank), + u32(lowerRank), + u32(upperRank), + numSlots, + u32(leafHeadSlot), + u32(leafParent), + u32(lowerHeadSlot), + u32(lowerParent), + u32(upperHeadSlot)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (pass.op != TopologyPassSpec::Op::Refine) { // coarsen/dilate have duplicate producers + combineDuplicateMasks<<>>(sorted.tileKey, + u32(leafFlag), + u32(leafRank), + u32(leafHeadSlot), + numSlots, + sorted.mask); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + } + + torch::Tensor upperStart = torch::empty({numGrids + 1}, i32Opts); + torch::Tensor lowerStart = torch::empty({numGrids + 1}, i32Opts); + torch::Tensor leafStart = torch::empty({numGrids + 1}, i32Opts); + + segmentNodeOffsets<<>>(u32(leafRank), + u32(lowerRank), + u32(upperRank), + em.segOffsets, + numGrids, + u32(leafStart), + u32(lowerStart), + u32(upperStart)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // --- THE synchronization: read back per-grid node offsets, size and allocate the output. --- + std::vector leafStartHost(numGrids + 1), lowerStartHost(numGrids + 1), + upperStartHost(numGrids + 1); + C10_CUDA_CHECK(cudaMemcpyAsync(leafStartHost.data(), + leafStart.data_ptr(), + sizeof(uint32_t) * (numGrids + 1), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaMemcpyAsync(lowerStartHost.data(), + lowerStart.data_ptr(), + sizeof(uint32_t) * (numGrids + 1), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaMemcpyAsync(upperStartHost.data(), + upperStart.data_ptr(), + sizeof(uint32_t) * (numGrids + 1), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + + BatchedTopologyResult result; + result.gridByteOffsets.resize(numGrids + 1); + result.leafCounts.resize(numGrids); + uint64_t totalBytes = 0; + for (int32_t g = 0; g < numGrids; ++g) { + result.gridByteOffsets[g] = totalBytes; + const uint32_t numUpper = upperStartHost[g + 1] - upperStartHost[g]; + const uint32_t numLower = lowerStartHost[g + 1] - lowerStartHost[g]; + const uint32_t numLeaf = leafStartHost[g + 1] - leafStartHost[g]; + result.leafCounts[g] = numLeaf; + totalBytes += gridByteSize(numUpper, numLower, numLeaf); + } + result.gridByteOffsets[numGrids] = totalBytes; + + const int32_t totalLeaf = int32_t(leafStartHost[numGrids]); + const int32_t totalLower = int32_t(lowerStartHost[numGrids]); + const int32_t totalUpper = int32_t(upperStartHost[numGrids]); + + result.buffer = TorchDeviceBuffer(totalBytes, src.device); + C10_CUDA_CHECK(cudaMemsetAsync(result.buffer.deviceData(), 0, totalBytes, stream)); + + torch::Tensor gridOffDev = torch::empty({numGrids + 1}, i64Opts); + C10_CUDA_CHECK(cudaMemcpyAsync(gridOffDev.data_ptr(), + result.gridByteOffsets.data(), + sizeof(uint64_t) * (numGrids + 1), + cudaMemcpyHostToDevice, + stream)); + + BuildDeviceArrays build; + build.dstBase = result.buffer.deviceData(); + build.gridByteOffsets = u64(gridOffDev); + build.upperStart = u32(upperStart); + build.lowerStart = u32(lowerStart); + build.leafStart = u32(leafStart); + + // --- Build all grids. --- + constexpr int32_t kGridDataWords = int32_t(sizeof(nanovdb::GridData) / sizeof(uint64_t)); + copyGridData<<>>( + em.srcGrids, build, numGrids); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + initGridTreeRoot<<>>(build, numGrids); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + torch::Tensor voxelCounts = torch::zeros({int64_t(totalLeaf) + 1}, i64Opts); + + if (totalLeaf > 0) { + buildUpperNodes<<>>( + build, sorted, u32(upperHeadSlot), numGrids, totalUpper); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + buildLowerNodes<<>>( + build, sorted, u32(lowerHeadSlot), u32(lowerParent), numGrids, totalLower); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + buildLeafNodes<<>>(build, + sorted, + u32(leafHeadSlot), + u32(leafParent), + numGrids, + totalLeaf, + u64(voxelCounts)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // In-place inclusive scan over elements [1, totalLeaf]; element 0 stays 0. + callCub([&](void *temp, size_t &bytes) { + C10_CUDA_CHECK(cub::DeviceScan::InclusiveSum( + temp, bytes, u64(voxelCounts) + 1, u64(voxelCounts) + 1, totalLeaf, stream)); + }); + + setLeafOffsets<<>>( + build, u64(voxelCounts), numGrids, totalLeaf); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + propagateLeafBBox<<>>( + build, u32(leafParent), numGrids, totalLeaf); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + propagateLowerBBox<<>>( + build, u32(lowerParent), numGrids, totalLower); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + propagateUpperBBox<<>>( + build, numGrids, totalUpper); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + finalizeGrids<<>>(build, u64(voxelCounts), numGrids); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return result; +} + +/// Runs a chained sequence of batched passes over the batch and wraps the final buffer in a +/// GridHandle. `passes` must be non-empty. +inline nanovdb::GridHandle +batchedTopologyHandle(const GridBatchData &batch, + const std::vector &passes, + cudaStream_t stream) { + TORCH_CHECK(!passes.empty(), "batchedTopologyHandle requires at least one pass"); + BatchedTopologyResult result = + runBatchedTopologyPass(sourceFromGridBatch(batch), passes[0], stream); + for (size_t p = 1; p < passes.size(); ++p) { + result = + runBatchedTopologyPass(sourceFromResult(result, batch.device()), passes[p], stream); + } + return nanovdb::GridHandle(std::move(result.buffer)); +} + +} // namespace batched +} // namespace detail +} // namespace fvdb + +#endif // FVDB_DETAIL_UTILS_NANOVDB_BATCHEDTOPOLOGYBUILDER_CUH diff --git a/tests/unit/test_batched_topology_builder.py b/tests/unit/test_batched_topology_builder.py new file mode 100644 index 000000000..bf2c84748 --- /dev/null +++ b/tests/unit/test_batched_topology_builder.py @@ -0,0 +1,218 @@ +# Copyright Contributors to the OpenVDB Project +# SPDX-License-Identifier: Apache-2.0 +# +"""Equivalence tests for the batched leaf-mask topology builder (issue #755). + +On CUDA, ``refined_grid`` / ``coarsened_grid`` (and through them ``conv_grid`` / +``conv_transpose_grid`` for K == S) build all batch members in a single batched pass instead of +one NanoVDB build + merge per member. These tests pin the batched results against: + +- ``from_ijk`` (PointsToGrid) grids built from independently computed expected coordinates, + compared **elementwise** (``torch.equal`` on ``ijk.jdata``), which pins the canonical NanoVDB + node ordering (root tiles by offset-shifted key, then x-major upper/lower child offsets), and +- the CPU implementations of the same ops, compared as per-member coordinate sets. + +Emphasis is on the cases the batched pipeline must get right and a per-member loop got for free: +batches with empty members (first/middle/last/all), members of wildly unequal sizes, coordinates +straddling root-tile boundaries (+-4096) and negative octants (where the sort-key encoding differs +from the stored ``Tile::key`` encoding), and multi-pass factors (4 = two chained passes). +""" + +import unittest + +import torch +from parameterized import parameterized + +import fvdb +from fvdb import GridBatch, JaggedTensor + + +def _build(ijks, device): + jt = JaggedTensor([t.to(device=device, dtype=torch.int32) for t in ijks]) + return GridBatch.from_ijk(jt, voxel_sizes=1.0, origins=0.0) + + +def _expected_refine_ijk(ijk: torch.Tensor, factor: int) -> torch.Tensor: + """All fine coordinates of ``ijk`` subdivided by ``factor`` (unique by construction).""" + if ijk.numel() == 0: + return ijk.reshape(0, 3) + offsets = torch.stack( + torch.meshgrid(torch.arange(factor), torch.arange(factor), torch.arange(factor), indexing="ij"), + dim=-1, + ).reshape(-1, 3) + fine = ijk.to(torch.int64)[:, None, :] * factor + offsets[None, :, :].to(ijk.device) + return fine.reshape(-1, 3).to(torch.int32) + + +def _expected_coarsen_ijk(ijk: torch.Tensor, factor: int) -> torch.Tensor: + """Unique coarse coordinates floor(ijk / factor).""" + if ijk.numel() == 0: + return ijk.reshape(0, 3) + coarse = torch.div(ijk.to(torch.int64), factor, rounding_mode="floor") + return torch.unique(coarse, dim=0).to(torch.int32) + + +def _ijk_sets(grid: GridBatch): + return [set(map(tuple, t.cpu().to(torch.int64).tolist())) for t in grid.ijk.unbind()] + + +# Coordinate batches exercising the tricky regimes of the batched builder. +def _tricky_batches(): + torch.manual_seed(42) + return { + "mixed_sizes": [ + torch.randint(-8, 8, (200, 3), dtype=torch.int32), + torch.tensor([[0, 0, 0]], dtype=torch.int32), + torch.randint(20, 60, (500, 3), dtype=torch.int32), + ], + "empty_members": [ + torch.empty((0, 3), dtype=torch.int32), + torch.randint(-10, 10, (100, 3), dtype=torch.int32), + torch.empty((0, 3), dtype=torch.int32), + torch.tensor([[5, 5, 5], [5, 5, 6]], dtype=torch.int32), + torch.empty((0, 3), dtype=torch.int32), + ], + "all_empty": [ + torch.empty((0, 3), dtype=torch.int32), + torch.empty((0, 3), dtype=torch.int32), + ], + # Straddles the +-4096 root-tile boundaries and negative octants: multiple root tiles per + # grid and coordinates where the offset-shifted sort key and stored Tile::key encodings + # order tiles differently. + "tile_boundaries": [ + torch.tensor( + [ + [-4097, -4097, -4097], + [-4096, -4096, -4096], + [-2049, 0, 0], + [-1, -1, -1], + [0, 0, 0], + [4095, 4095, 4095], + [4096, 4096, 4096], + [4096, -4097, 0], + ], + dtype=torch.int32, + ), + torch.randint(-4200, -3900, (300, 3), dtype=torch.int32), + torch.randint(3900, 4200, (300, 3), dtype=torch.int32), + ], + "single_grid": [ + torch.randint(-16, 16, (400, 3), dtype=torch.int32), + ], + "larger_batch": [torch.randint(-32, 32, (50 + 37 * i, 3), dtype=torch.int32) for i in range(16)], + } + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required for the batched builder") +class TestBatchedTopologyBuilder(unittest.TestCase): + def _check_against_expected(self, result: GridBatch, expected_ijks, msg: str): + """Elementwise-pin `result` against a from_ijk build of the expected coordinates.""" + expected = _build(expected_ijks, result.device) + self.assertEqual(result.grid_count, expected.grid_count, msg) + self.assertTrue(torch.equal(result.num_voxels, expected.num_voxels), msg) + self.assertTrue( + torch.equal(result.ijk.jdata, expected.ijk.jdata), + f"{msg}: voxel enumeration (canonical node order) differs from PointsToGrid", + ) + for b in range(result.grid_count): + self.assertTrue( + torch.equal(result.bbox_at(b).cpu(), expected.bbox_at(b).cpu()), + f"{msg}: bbox of member {b}", + ) + + def _check_against_cpu(self, result: GridBatch, cpu_result: GridBatch, msg: str): + self.assertEqual(result.grid_count, cpu_result.grid_count, msg) + self.assertTrue(torch.equal(result.num_voxels.cpu(), cpu_result.num_voxels), msg) + self.assertEqual(_ijk_sets(result), _ijk_sets(cpu_result), msg) + + @parameterized.expand([(name,) for name in _tricky_batches().keys()]) + def test_refined_grid_matches_expected_and_cpu(self, name): + coords = _tricky_batches()[name] + for factor in (2, 4): + grid = _build(coords, "cuda") + result = grid.refined_grid(factor) + expected = [_expected_refine_ijk(c, factor) for c in coords] + self._check_against_expected(result, expected, f"{name} refine x{factor}") + cpu_result = _build(coords, "cpu").refined_grid(factor) + self._check_against_cpu(result, cpu_result, f"{name} refine x{factor} vs CPU") + + @parameterized.expand([(name,) for name in _tricky_batches().keys()]) + def test_coarsened_grid_matches_expected_and_cpu(self, name): + coords = _tricky_batches()[name] + for factor in (2, 4): + grid = _build(coords, "cuda") + result = grid.coarsened_grid(factor) + expected = [_expected_coarsen_ijk(c, factor) for c in coords] + self._check_against_expected(result, expected, f"{name} coarsen x{factor}") + cpu_result = _build(coords, "cpu").coarsened_grid(factor) + self._check_against_cpu(result, cpu_result, f"{name} coarsen x{factor} vs CPU") + + def test_conv_grid_k2s2_multi_grid_matches_per_member(self): + coords = _tricky_batches()["mixed_sizes"] + full = _build(coords, "cuda") + conv = full.conv_grid(kernel_size=2, stride=2) + convt = full.conv_transpose_grid(kernel_size=2, stride=2) + for b, c in enumerate(coords): + single = _build([c], "cuda") + conv_single = single.conv_grid(kernel_size=2, stride=2) + convt_single = single.conv_transpose_grid(kernel_size=2, stride=2) + self.assertTrue( + torch.equal(conv.ijk.unbind()[b], conv_single.ijk.jdata), + f"conv_grid k2s2 member {b}", + ) + self.assertTrue( + torch.equal(convt.ijk.unbind()[b], convt_single.ijk.jdata), + f"conv_transpose_grid k2s2 member {b}", + ) + + @parameterized.expand([("mixed_sizes",), ("empty_members",), ("tile_boundaries",)]) + def test_conv_stride1_and_k3s2_match_cpu(self, name): + # Stride-1 uniform K routes through batched box-dilate passes (odd K: [-1,1]^3 per pass; + # even K: one-sided {-1,0}^3 / {0,1}^3 passes), and k3s2 transpose through refine + one + # negative pad pass. Pin against the CPU implementation (coordinate sets) and against a + # from_ijk build of the CPU coordinates (elementwise: canonical node order). + coords = _tricky_batches()[name] + cuda_grid = _build(coords, "cuda") + cpu_grid = _build(coords, "cpu") + cases = ( + [("conv_grid", k, 1) for k in (2, 3, 4, 5)] + + [("conv_transpose_grid", k, 1) for k in (2, 3, 4, 5)] + + [("conv_transpose_grid", 3, 2)] + ) + for op, k, s in cases: + msg = f"{name} {op} k{k}s{s}" + result = getattr(cuda_grid, op)(kernel_size=k, stride=s) + reference = getattr(cpu_grid, op)(kernel_size=k, stride=s) + self._check_against_cpu(result, reference, msg) + expected = _build(list(reference.ijk.unbind()), "cuda") + self.assertTrue( + torch.equal(result.ijk.jdata, expected.ijk.jdata), + f"{msg}: voxel enumeration (canonical node order) differs from PointsToGrid", + ) + + def test_masked_refine_multi_grid(self): + # Masked subdivision routes through pruneGrid then the batched refine. + coords = _tricky_batches()["mixed_sizes"] + grid = _build(coords, "cuda") + mask = JaggedTensor([((t.sum(-1) % 2) == 0) for t in grid.ijk.unbind()]) + result = grid.refined_grid(2, mask=mask) + cpu_grid = _build(coords, "cpu") + cpu_mask = JaggedTensor([m.cpu() for m in mask.unbind()]) + cpu_result = cpu_grid.refined_grid(2, mask=cpu_mask) + self._check_against_cpu(result, cpu_result, "masked refine x2 vs CPU") + + def test_refine_coarsen_roundtrip_grid_count(self): + # Serialization-visible metadata: grid_count/address arithmetic must be consistent for a + # multi-grid handle produced by the batched builder. + coords = _tricky_batches()["empty_members"] + grid = _build(coords, "cuda") + fine = grid.refined_grid(2) + back = fine.coarsened_grid(2) + self.assertEqual(fine.grid_count, grid.grid_count) + self.assertEqual(back.grid_count, grid.grid_count) + self.assertEqual(_ijk_sets(back), _ijk_sets(grid)) + self.assertTrue(torch.equal(back.num_voxels, grid.num_voxels)) + + +if __name__ == "__main__": + unittest.main()