Generated-topology grid construction dominates per-iteration ConvolutionPlan cost (serial per-grid build + merge)

[swahtz]

Summary

When training generative sparse networks (topology changes every iteration, so ConvolutionPlans must be rebuilt per step), plan construction dominates wall-clock and leaves the GPU underutilized (~35% GPU, one CPU core pegged in the shape-VAE example from #753). Profiling shows the cost is not the kernel-map build — it is the construction of the generated topology grid itself.

Measurements

Shape-VAE example (#753) on the 254-shoe GSO dataset, batch 16 (~110k voxels), resolution 64, 4-level encoder/decoder:

• One training iteration: 165 ms, of which 69.6 ms is ConvolutionPlan construction (21 plans) and 22 ms per-minibatch conv_grid pyramid rebuild.
• cProfile of a single ConvolutionPlan.from_grid_batch_transposed(2, 2, grid_batch16): 24 of 26 ms is _fvdb_cpp.conv_transpose_grid (generated target-grid construction). The transpose kernel map (gs_build_transpose_topology) is only 0.5 ms; a k3 s1 same-topology kernel map (gs_build_topology) is 0.8 ms.
• Plan-construction time scales ~linearly with batch voxel count (65 / 126 / 182 ms per iteration at batch 16 / 32 / 48), so batching the training step does not amortize it.

Root cause

dispatchBuildGridForConvTranspose (and the analogous conv_grid path) constructs the generated grid one batch member at a time and then merges:

• src/fvdb/detail/ops/BuildGridForConvTranspose.cu: even on the fast leaf-mask subdivision path (K=S=2, which this workload hits), perItemGridHandle loops over the batch building one NanoVDB grid handle per member, each with its own stream work, then nanovdb::cuda::mergeGridHandles. At batch 16 that is ~1.5 ms of fixed per-grid overhead x 16, serialized.

So per-iteration cost is dominated by per-grid fixed overhead in generated-topology construction, not by topology size or kernel-map work.

Suggested direction

Batched generated-grid construction: build the strided/subdivided support for all batch members in one pass (single staging + single grid-batch build) instead of per-member handle construction + merge. This would directly speed up:

• GridBatch.conv_grid / conv_transpose_grid
• ConvolutionPlan.from_grid_batch(_transposed) with generated targets (target_grid=None)
• any per-iteration generative training loop (the Generative Transposed Convolution #741 / Add fvdb.nn.Prune and generative shape completion / VAE examples #753 use case; MinkowskiEngine and spconv build their equivalents fully on-device, which is why they sustain higher GPU utilization on the same workload)

Secondary (much smaller): from_grid_batch(1, 1, g, g) already short-circuits execution to _MatmulBackend, but construction still pays ~1.2 ms of Python-side validation/coverage setup per call; a fast path for the identity-plan case would tidy up per-iteration classifier heads.

Repro

import torch, cProfile
from fvdb import ConvolutionPlan
from fvdb.utils.examples import load_gso_shoes
import fvdb

meshes = load_gso_shoes(limit=16)
v = fvdb.JaggedTensor([(m[0] - m[0].amin(0)) / m[0].amax() * 0.96 + 0.02 for m in meshes])
f = fvdb.JaggedTensor([m[1].int() for m in meshes])
g = fvdb.GridBatch.from_mesh(v, f, voxel_sizes=1/64, origins=0.0)
cProfile.run('ConvolutionPlan.from_grid_batch_transposed(2, 2, g)', sort='cumulative')
# -> ~24ms in _fvdb_cpp.conv_transpose_grid, ~0.5ms in gs_build_transpose_topology