Generalize leaf-mask morphology to arbitrary factors: non-power-of-two coarsen/subdivide and general conv strides still pay the coordinate-list + radix-sort path

[swahtz]

Summary

The leaf-mask topology work (#710, #712 / issue #711) converted the hot CUDA topology-builder paths to nanovdb::tools::cuda::TopologyBuilder morphology, but the fast paths are gated on uniform power-of-two factors because NanoVDB's CoarsenGrid/RefineGrid are hard-wired 2×. Everything else — non-power-of-two factors, anisotropic factors (even anisotropic powers of two, e.g. (2, 2, 1)), and general conv strides — falls back to the expanded coordinate-list + PointsToGrid radix-sort path, with all the costs catalogued in #711 (12–20 B/candidate torch tensors incl. written-then-discarded jidx arrays, ~40 B/candidate raw-cudaMalloc sort scratch outside the torch caching allocator, per-batch-item serial sorts, a joffsets().cpu() host sync, and silent int32 truncation above 2³¹ candidates).

Investigation result (details below): NanoVDB's CoarsenGrid is only two factor-2-specific functors away from supporting arbitrary factors. All the heavy machinery — node counting, buffer allocation, tree assembly, canonical-order value indexing — is the factor-agnostic TopologyBuilder core. A generalized "coarsen by (fx, fy, fz)" tool eliminates the remaining coordinate-list fallbacks for coarsen/subdivide, and a trivial variant (subsample instead of union) covers general conv strides.

Plan of record: implement the generalized tool upstream in NanoVDB (parameterizing / extending CoarsenGrid/RefineGrid, keeping the 2× bit-trick as a specialization), then bump the get_nanovdb.cmake pin and plumb it into fvdb's builders.

Line references: fvdb files on branch issue-711-leaf-mask-topology (PR #712); NanoVDB files at the pinned commit b7fc4fc7 (src/cmake/get_nanovdb.cmake).

Current fallback sites

Site	Trigger	Cost (per the #711 shared-costs analysis)
BuildCoarseGridFromFine.cu:61-62, :112-113	uniformPowerOfTwoLog2(factor) < 0 — any non-uniform or non-pow2 factor (3, 5, 6, (2,2,1), (2,3,4), …)	K=1: ~20 B/fine-voxel torch + ~40 B raw scratch, full sort
BuildFineGridFromCoarse.cu:391, :457	subdivUniformPowerOfTwoLog2(factor) < 0	K=f³ expansion: e.g. factor 3 → 27 candidates/coarse voxel ≈ 540 B torch + ~1080 B raw per coarse voxel
BuildGridForConv.cu:253-255	non-uniform kernel, or stride ∉ {1, kernelSize}	K=kernelVol + a full torch::zeros mask + compaction-while-expanded-live
BuildGridForConv.cu:200-202	stride == kernelSize routes to the coarsen builder — which itself falls back for non-pow2 stride, so e.g. a 3³ stride-3 conv pays the coord path	as coarsen row
BuildGridForConvTranspose.cu:280-282 (and :216 → subdivide fallback for non-pow2 stride)	same shape as conv	K=kernelVol, no compaction

Note the anisotropic-power-of-two case: uniformPowerOfTwoLog2 requires f[0]==f[1]==f[2] (BuildCoarseGridFromFine.cu:32-34), so (2, 2, 1) — a perfectly reasonable factor for 2.5D/BEV-style data — takes the sort path today even though each axis is a power of two. The generalized driver removes the uniformity restriction as a side effect.

Anatomy of CoarsenGrid — what is actually factor-2-specific

nanovdb::tools::cuda::CoarsenGrid<BuildT>::getHandle() (CoarsenGrid.cuh:95-150) runs:

1. coarsenRoot() (:155) — host side: D2H-copies the source root topology, remaps each root-tile origin with coarsenCoord (= per-component floor(n/2), MorphologyHelpers.h:85-99), dedupes in a std::map, uploads the new root. Factor-2 assumption: the ÷2 image of a 4096-aligned tile always lands in exactly one coarse tile.
2. CoarsenInternalNodesFunctor (Morphology.cuh:449) — one thread per source leaf: coarsenedOrigin = coarsenCoord(leaf.origin()), then atomically sets one upper-mask bit and one lower-mask bit under the matching root tile. Factor-2 assumption: a source leaf's ÷2 image (span 4) always lands in exactly one coarse leaf → a single origin suffices.
3. Generic TopologyBuilder steps — allocateInternalMaskBuffers, countNodes (prefix sums over the masks), buffer allocation, processGridTreeRoot/UpperNodes/LowerNodes. All factor-agnostic; this is the same machinery PadGrid.cuh drives.
4. CoarsenLeafMasksFunctor (Morphology.cuh:876) — one thread per source leaf: folds the 512-bit value mask 2× per axis with word-level bit tricks (coarsenMask, :879-905), locates the destination leaf via root().probeLeaf(coarsenedOrigin), and atomicOrs four words with the appropriate sub-octant shift. Factor-2 assumptions: the bit-fold trick, and the single-destination-leaf property.
5. Generic finalization — processLeafOffsets (canonical-order ValueOnIndex numbering), processBBox, postProcessGridTree. Factor-agnostic.

So the generalization touches exactly steps 1, 2, and 4; steps 3 and 5 are reused verbatim.

Proposed design: coarsen by arbitrary (fx, fy, fz)

Mapping: coarse = (floor(i/fx), floor(j/fy), floor(k/fz)) (floor division toward −∞, matching fvdb semantics and coarsenComponent's negative-rounding).

Geometric bound that keeps the scatter small: a source leaf spans 8 fine voxels per axis, so its coarse image spans ≤ 8 coarse voxels per axis (f=1 worst case), which straddles at most 2 coarse leaves per axis → ≤ 2³ = 8 destination leaves per source leaf. Likewise a source root tile (span 4096) maps to ≤ 2 coarse root tiles per axis → ≤ 8 candidates per source tile.

Passes (mirroring the CoarsenGrid skeleton):

1. Root pass (host, same map-dedupe): emit up to 8 candidate coarse tiles per source tile instead of 1. To avoid empty root tiles in the output, either filter candidates by the source tile's actual leaf bounding box, or accept conservative tiles (a root tile whose upper mask stays empty contributes zero nodes in countNodes — verify TopologyBuilder tolerates this; CoarsenGrid never produces one today).
2. Internal pass (device, per source leaf): from leaf.valueMask(), determine which of the ≤ 8 candidate coarse leaves actually receive at least one active bit, and atomically set upper/lower mask bits only for those. Exactness here matters: conservative (bbox-based) marking would allocate all-off leaves, which downstream code assumes don't exist (and which PruneGrid exists to remove).
3. Generic: countNodes → allocate → processGridTreeRoot/UpperNodes/LowerNodes (unchanged).
4. Leaf-mask pass (device, per source leaf): no bit-fold trick exists for general f, and none is needed — iterate the source mask's on-bits, computing for each the destination (leaf, bit); accumulate into up to 8 local Mask<3> fragments (512 B — local/shared memory, not registers), then atomicOr whole words into the destination leaves via probeLeaf (≤ 8 leaves × 8 words = 64 atomics per source leaf, vs. per-voxel atomics). Since f³ fine voxels collapse onto each coarse bit, the local reduction absorbs almost all the write traffic.
5. Generic: processLeafOffsets, processBBox, postProcessGridTree (unchanged) — which also preserves the canonical-order invariant (feature row i ↔ getValue() == i+1) that all differentiable consumers rely on.

Transient memory: only the upper/lower mask scratch TopologyBuilder already allocates — O(root-tile count), identical to the existing dilate/pad/coarsen ops. No per-candidate arrays, no sort scratch, no joffsets().cpu() sync, no 2³¹ overflow surface.

Variant A — arbitrary-factor refinement (subdivide)

Inverse mapping: coarse voxel c → fine box c*f + [0, f)³. A coarse leaf (span 8) images to a fine span of 8·f per axis → ≤ f+1 fine leaves per axis, so the fan-out is (f+1)³ destination leaves per source leaf — larger than coarsening but still sort-free and proportional to output size (which the result must materialize anyway; today's fallback pays K=f³ sorted candidates on top). Masked subdivision composes exactly as it does on the pow-2 path already shipped: PruneGrid(coarse, mask) → refine (BuildFineGridFromCoarse.cu:383-386).

Variant B — subsample mode (general conv strides)

The conv fallback's semantics (convIJKForGrid) are: Minkowski-sum by the kernel window, keep only coordinates ≡ 0 (mod stride), divide by stride. That second step is subsampling (decimation), not union-coarsening — the coarse bit is on iff the single fine voxel at coarse*f is active. That is strictly simpler than union coarsening (one probe per output bit, no union), and composes with the existing dilate/pad fast paths to cover BuildGridForConv/BuildGridForConvTranspose for any stride: dilate/pad by window → subsample by stride. A mode flag on the same driver suffices.

Implementation plan (upstream NanoVDB + fvdb plumbing)

Phase 1 — NanoVDB (openvdb repo): add the arbitrary-factor tool alongside CoarsenGrid.cuh/RefineGrid.cuh in nanovdb/tools/cuda/ — either a factor parameter on the existing classes or a sibling tool (e.g. ResampleGrid) sharing their TopologyBuilder-driven skeleton. The 2× bit-fold (coarsenMask) stays as a fast-path specialization; the generalized functors (root candidates ≤ 8/tile, internal-pass ≤ 8 leaves/leaf, leaf-pass local Mask<3>-fragment scatter) live in util/cuda/Morphology.cuh next to their 2× counterparts. The union-vs-subsample reduction mode (variant B) is a template/enum parameter on the same functors. Estimated small: steps 3/5 of the skeleton are unchanged library calls, and CoarsenGrid.cuh itself is ~250 lines.

Phase 2 — fvdb: bump the get_nanovdb.cmake pin, then replace the fallback branches:

• BuildCoarseGridFromFine.cu — drop uniformPowerOfTwoLog2 gating; route every factor to the new tool (the repeated-CoarsenGrid-pass loop for pow-2 factors can be retired too, see bonus note below).
• BuildFineGridFromCoarse.cu — same for subdivision (masked variant keeps the existing PruneGrid-then-refine composition).
• BuildGridForConv.cu / BuildGridForConvTranspose.cu — replace the general-stride coordinate fallbacks with dilate/pad + subsample-mode resample.
• Delete the then-unreferenced CUDA coordinate-list legs (CoarseIjkForFineGrid.{cu,h}, fineIJKForCoarseGrid, convIJKForGrid) — CPU and PrivateUse1 keep their paths.

PadGrid.cuh from #710 remains the in-repo reference for driving TopologyBuilder if any glue is needed while the upstream PR is in flight, but the tool itself belongs in NanoVDB where CoarsenGrid/RefineGrid live.

What it unlocks

1. coarsened_grid with any factor — including anisotropic pow-2 ((2,2,1)) which is arguably a latent performance bug today.
2. refined_grid with any factor (masked included, via prune-first).
3. conv_grid / conv_transpose_grid for all kernel/stride combinations — closing the last coordinate-list fallbacks in the DL hot path (BuildGridForConv.cu:253, BuildGridForConvTranspose.cu:280) and retiring their K=kernelVol expansions and the 3³-conv 79.5M-voxel silent-overflow threshold noted in Grid topology ops build coordinate lists where leaf-mask operations would be far cheaper #711.
4. Deletion candidates once no callers remain: CoarseIjkForFineGrid.{cu,h}, fineIJKForCoarseGrid, convIJKForGrid (CUDA legs; CPU and PrivateUse1 keep their paths).

Verification plan (mirrors the #711/#712 test template)

• CUDA↔CPU parity (ordered ijk per batch item — the canonical-order invariant) for factors {1, 2, 3, 5, 6, 8, (2,2,1), (2,3,4)} on random, empty, single-voxel, and multi-item batches.
• Sliced/non-contiguous views (tail/gap/reversed/single selections) vs. contiguous reference, per tests/unit/test_sliced_batch.py.
• Boundary crossings: leaf-local coord 7, root-tile boundary 4095/4096, and negative coordinates (floor-division rounding is the classic off-by-one source; coarsenComponent's negative handling is the model).
• Consumer round-trip: max/avg pool and SparseConv3d forward+backward at a non-pow2 factor/stride against the ground-truth tests.
• Peak-memory regression (torch.cuda.max_memory_allocated) showing the drop vs. the coordinate path, per the test_dual.py::test_dual_grid_peak_memory pattern.
• Equivalence with the pow-2 fast path: for factors 2/4/8 the new driver must produce bit-identical grids to the repeated-CoarsenGrid path (then the repeated-pass loop can optionally be retired in favor of one generalized pass — a factor-8 coarsen today runs 3 full build cycles).

That last point is a bonus: even the existing pow-2 path does log2(f) full build-allocate-finalize cycles with intermediate handles; the generalized driver does one.