Batched leaf-mask topology construction for generated conv grids

[swahtz]

Batched leaf-mask topology construction for generated conv grids

Fixes #755.

Summary

conv_grid / conv_transpose_grid with generated targets (and refined_grid / coarsened_grid / the stride-1 dilate-pad paths generally) built their output one batch member at a time: each member paid a full NanoVDB RefineGrid/CoarsenGrid/DilateGrid/PadGrid build (3+ stream synchronizations each, a host-side speculative root refinement with a pageable D2H readback, a GridHandle constructor with 2 blocking memcpys + raw cudaMalloc/cudaFree), plus a host-built proxy grid + H2D per empty member, and finally nanovdb::cuda::mergeGridHandles added another synchronization per grid. That is ~1.5 ms of fixed overhead x B members x 21 plans per training iteration in generative workloads where topology changes every step (#741 / #753) — construction cost scaled linearly in batch size and left the GPU at ~35% utilization.

This PR batches the leaf-mask topology machinery across the whole grid batch (extending #712's leaf-mask direction — no return to coordinate staging), in four self-contained commits:

1. Batched topology builder + refine/coarsen passes

New src/fvdb/detail/utils/nanovdb/BatchedTopologyBuilder.cuh: one batched pass builds ALL grids at once. Per pass:

• one emission kernel over all source leaves of all members produces candidate output leaves as (root-tile sort key, in-tile node key, origin, 512-bit mask) slots, segmented per grid — mask bit math is NanoVDB's own (RefineLeafMasksFunctor::refineMask, CoarsenLeafMasksFunctor::coarsenMask);
• two stable segmented radix sorts put each grid's slots in canonical NanoVDB node order (root tiles by the PointsToGrid offset-shifted key, then x-major upper/lower child offsets), so produced grids enumerate voxels identically to every other build path (pinned by elementwise torch.equal tests);
• head-flag + scan passes dedup nodes and derive per-grid counts and parent linkage; root tiles are derived on device from the unique upper keys (the host-side speculative refineRoot and its readback disappear);
• one cudaStreamSynchronize reads back per-grid node counts, one buffer is allocated, and batched kernels (transcribed from tools::cuda::TopologyBuilder's functors with (gridIndex, localIndex) indexing) write every grid's headers (mGridIndex=g, mGridCount=B), nodes, leaf mOffset/mPrefixSum, and bboxes. Empty members become valid empty grids inline.

fineGridHandleFromCoarseCUDA / coarseGridHandleFromFineCUDA route through it for all batch sizes; multi-pass factors (4, 8) chain passes, mirroring the previous per-pass semantics. Checksums are disabled on the output (matching ops::contiguousGridHandle and mergeGridHandles). CPU / PrivateUse1 paths and non-power-of-two coordinate fallbacks unchanged; masked subdivision still prunes upstream, then batched-refines.

2. Batched box-dilate passes (stride-1 K>1, k3s2)

A BoxDilate pass computes the Minkowski sum with an axis-aligned unit box via pure per-axis bit shifts (scatter formulation, up to 27 target leaves per source leaf, deduplicated by the same back-end). This covers DilateGrid's 26-neighbor dilation ([-1,1]^3) and both PadGrid octants ({-1,0}^3, {0,1}^3), so the stride-1 uniform-K conv/conv-transpose paths and the k3s2 transpose (refine + negative pad) become batched pass sequences. The per-member perItemGridHandle drivers are deleted from both conv builders.

3. Identity-plan fast path (issue's secondary item)

ConvolutionPlan.from_grid_batch[_transposed](1, 1, g[, g]) short-circuits to the matmul backend without building tensor-valued transform diagnostics (which are trivially exact when source and target share GridBatchData), preserving channel-pair validation, backend-name rejection, and the general path for distinct-but-equal-looking grids: ~1.2 ms -> 0.08 ms per call.

4. Grid-construction cheap wins

• makeGridBatchData: leafBatchIndices via one repeat_interleave instead of B x torch::full + torch::cat (B+1 dispatches on every grid construction).
• voxelSizesTensor / voxelOriginsTensor: accessor fill instead of per-element ATen indexing (6 dispatches per grid per call, on every plan construction's transform validation).

Measurements

RTX PRO 6000 Blackwell, synthetic shell batches (~7.8k voxels/member, resolution 64), median of 20 CUDA-event-timed iterations (src/benchmarks/convolution/benchmark_conv_grid_build.py):

op	before B=16	after B=16	before B=48	after B=48
conv_transpose_grid k2s2	9.0 ms	0.80 ms	25.6 ms	0.94 ms
conv_grid k2s2	8.7 ms	0.80 ms	23.4 ms	0.94 ms
conv_grid k3s1	9.8 ms	0.83 ms	25.1 ms	1.00 ms
plan from_grid_batch(2,2,g)	8.9 ms	1.27 ms	24.8 ms	1.62 ms
plan from_grid_batch_transposed(2,2,g)	10.2 ms	1.94 ms	29.4 ms	3.27 ms
4-level plan-pyramid rebuild (8 plans + 3 conv_grids)	114 ms	17.4 ms	309 ms	27.0 ms
identity plan from_grid_batch(1,1,g,g)	~1.2 ms	0.08 ms	—	—

Construction cost is now near-flat in batch size (B=1: within noise of the old single-grid path). The issue's per-iteration plan-construction share (~70 ms of a 165 ms shape-VAE iteration at B=16) drops to single-digit milliseconds.

Follow-ups (out of scope, same back-end): an ijk emission front-end to resolve BuildGridFromIjk.cu's per-member FIXME (from_ijk/from_points/shifted-geometry fallbacks), and dilated_grid/BuildPaddedGrid's standalone per-member loops.

Test plan

• New tests/unit/test_batched_topology_builder.py (18 tests): elementwise (torch.equal on ijk.jdata, num_voxels, per-member bboxes) equivalence against from_ijk-built expected topologies — pinning canonical node order — plus per-member coordinate-set equality against the CPU paths, across: mixed member sizes, empty members (first/middle/last/all), coordinates straddling +-4096 root-tile boundaries and negative octants (where the sort-key and stored Tile::key encodings order differently), single-grid and 16-grid batches, factors 2 and 4, masked refine, conv_grid/conv_transpose_grid K in {2,3,4,5} at stride 1, k3s2 transpose, k2s2 vs per-member, and a refine->coarsen round trip.
• Existing suites, all green (724 passed, 1 skipped): pytest unit/test_conv_semantics_integration.py unit/test_conv_default.py unit/test_conv_transpose_default.py unit/test_batching.py unit/test_basic_ops.py unit/test_sliced_batch.py unit/test_conv_semantics.py unit/test_conv_ground_truth.py unit/test_nn_modules.py — includes the elementwise ijk order pins, sliced-view coverage, resource-stats path pinning, and the matmul/identity plan contract tests.
• Benchmark: python src/benchmarks/convolution/benchmark_conv_grid_build.py (numbers above; --gso runs the issue's verbatim GSO repro).

🤖 Generated with Claude Code