Batched leaf-mask topology construction for generated conv grids

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,

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()

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<double, 2>();
     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<double, 2>();
     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;
 }

src/fvdb/detail/GridBatchDataFactory.cu

@@ -178,15 +178,18 @@ makeGridBatchData(nanovdb::GridHandle<TorchDeviceBuffer> &&gridHdl,
     TORCH_CHECK(listIndices.numel() == 0 || listIndices.size(0) == (batchOffsets.size(0) - 1),
                 "Invalid list indices when building grid");
 
-    std::vector<torch::Tensor> 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<int64_t, 1>();
     for (int64_t i = 0; i < batchSize; i += 1) {
-        leafBatchIdxs.push_back(
-            torch::full({hostMeta[i].mNumLeaves},
-                        static_cast<fvdb::JIdxType>(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<nanovdb::GridHandle<TorchDeviceBuffer>>(std::move(gridHdl));
 

src/fvdb/detail/ops/BuildCoarseGridFromFine.cu

@@ -11,12 +11,11 @@
 #include <fvdb/detail/utils/AccessorHelpers.cuh>
 #include <fvdb/detail/utils/Utils.h>
 #include <fvdb/detail/utils/VoxelSizeUtils.h>
-#include <fvdb/detail/utils/nanovdb/CreateEmptyGridHandle.h>
+#include <fvdb/detail/utils/nanovdb/BatchedTopologyBuilder.cuh>
 
 #include <nanovdb/NanoVDB.h>
 #include <nanovdb/tools/CreateNanoGrid.h>
 #include <nanovdb/tools/GridBuilder.h>
-#include <nanovdb/tools/cuda/CoarsenGrid.cuh>
 
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -53,10 +52,11 @@ dispatchBuildCoarseGridFromFine(const GridBatchData &fineGridBatch,
 nanovdb::GridHandle<TorchDeviceBuffer>
 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,39 +65,19 @@ 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
         // (possibly sliced) selected grids into a fresh contiguous handle.
         return ops::contiguousGridHandle(fineGridBatch);
     }
 
-    std::vector<nanovdb::GridHandle<TorchDeviceBuffer>> 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<TorchDeviceBuffer> handle;
-        for (int p = 0; p < nPasses; p += 1) {
-            nanovdb::tools::cuda::CoarsenGrid<nanovdb::ValueOnIndex, BuilderResource> op(
-                grid, stream.stream());
-            op.setChecksum(nanovdb::CheckMode::Default);
-            op.setVerbose(0);
-            handle = op.getHandle(guide);
-            C10_CUDA_KERNEL_LAUNCH_CHECK();
-            grid = handle.deviceGrid<nanovdb::ValueOnIndex>();
-        }
-        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<batched::TopologyPassSpec> passes(nPasses,
+                                                        batched::TopologyPassSpec::coarsen());
+    return batched::batchedTopologyHandle(fineGridBatch, passes, stream.stream());
 }
 
 template <>