[Kernel] Add submanifold sparse 3D convolution (bf16 implicit GEMM, gfx950)

[jiacao-amd]

Summary

Adds a self-contained submanifold sparse 3D convolution (spconv-style): raw voxel
coordinates in, dense output features out, kernel map built internally. Matches the
input convention of spconv / MinkowskiEngine / TorchSparse so it drops into their call
sites. Four GPU kernels, one public entry point.

coords --[map]--> LUT --[compact]--> pair list --[gemm]--> output

Draft — see Open questions at the bottom before merging.

Design, and the measurements behind it

Sized against measured production shapes rather than synthetic ones: nuScenes
CenterPoint over all 6019 val frames is median 91,090 active voxels (p95 114,827, max
135,306), and PTv3 captures give ScanNet ~148k and SemanticKITTI ~120k. Earlier tuning
had used a 10k-point toy input, which understates both N and neighbour density by enough
to invert several conclusions.

Two mappers, both required. build_lut_auto routes by bounding-box volume:

grid	dense	Z-Delta
195³	41 µs	160 µs
595×595×245	93 µs	180 µs
2095²×695	cannot run (3.05e9 cells, 12.2 GB, int32 index overflow)	203 µs

The dense presence grid is 2-4x faster on compact grids but is O(volume) in memory and
fill; Z-Delta (Spira, arXiv 2511.20834 §5.2) sorts
packed keys once and does K² binary searches instead of K³, flat in grid volume. The
last row is PTv3 on SemanticKITTI, so the second mapper is a functional requirement
outdoors, not a tuning option.

Compaction is the largest single win (2.4-3.3x). An indexed [tile, kv, row] layout
runs all 16 rows of a tile when only 2-3 have a neighbour for that tap — measured
5.4-6.2x wasted MACs, i.e. 16-21% of multiplies useful, while executed throughput was
already 45-53% of peak. The hardware was saturated computing zeros.

No LDS, deliberately. With one 64-lane wave and a 16-wide tile every operand element
is read by exactly one lane — reuse is 1.00x in fp32 and again in bf16 after
vectorisation. Isolated by compiling for a padded C_IN so useful work is unchanged: LDS
alone made it 17x slower by collapsing occupancy to 5 blocks/CU at C_IN=512. An
LDS-transposed epilogue was also built and measured 8x slower, with padding making no
difference — bounded by the round-trip, not by bank conflicts.

Weight layout is [kv, t, k_outer, u, k_inner]. The obvious K-major
[kv, t, u, ci] gives each lane its 8 consecutive k but puts adjacent lanes C_IN*2 bytes
apart, so a wave reads 16 half-used cache lines: 50% coalescing, and that was 84% of the
kernel's requested traffic. Interleaving k_inner innermost makes it 8 fully-used lines —
1.8x, a pure packing change.

dtype is bf16 because production sparse conv is 16-bit: spconv ships fp32/fp16/bf16,
its own benchmark tables are F16/TF32, and NVIDIA's shipped CenterPoint logs are
perf-float16-*.log and perf-int8.log with no fp32 variant. An fp32 GEMM was built
first and measured 1.6-2.9x slower at every shape.

Performance

End to end through the public API on production shapes (warm cache, includes mapping):

case	N	C	µs
nuScenes	32371	128	47
nuScenes	26509	256	101
nuScenes	15052	512	172
SemKITTI	94517	128	178

A 6-layer stage sharing one indice_key runs in 0.596 ms.

API style

Commits 3-5 port the kernels to the current fx surface per the kernel-code-cleanup
skill (fx.Index → fx.Int32, arith.* → operators, vector.* → fx.Vector), −34
lines net. Verified bit-identical against a golden capture and perf-neutral
(median-of-5: 1.00x / 1.02x / 1.02x).

Three sites are left legacy on purpose, each documented in-code:

• scf.IfOp everywhere. _collect_assigned_vars treats the receiver of a method
  call inside a branch as branch state, so any lut_.store(...) puts a GTensor into
  result_names and scf_if_dispatch rejects it. Three formulations were tried cold;
  all fail identically. Fixing this means changing kernels/common/tensor_shim.py,
  shared by four other kernels.
• The Z-Delta kernel. Its keys are KEY_T (i32 or i64 by grid extent) and
  GTensor.load returns a raw value typed by that dtype, so operator compares don't
  apply and fx.Int32 offsets take a different buffer_load path.
• The GEMM epilogue. Migrating it is op-count-identical but measured 185 µs vs 170 µs
  at C_OUT=512 (median-of-7, range 184.8-186.4) — the same instruction-scheduling class
  of exception Port GEMM/MoE/conv kernels to the layout API #913 documents for mxfp_moe gemm1 and fp8_4wave.

Also note this kernel does not use buffer_ops/create_buffer_resource directly, so the
layout-API port in #913 does not apply to it; it goes through the shared GTensor shim.

Test plan

FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest tests/kernels/test_sparse_conv3d_implicit.py -v

39 tests: generic shapes (K=3 and K=5, non-divisible channel counts, degenerate C=1),
[N,4] batched coords, explicit vs derived spatial_shape, dense-vs-Z-Delta bit
equality, the >2^31-cell grid, 32- vs 64-bit key selection, weight-cache invalidation and
eviction, and full-tap coverage of the compacted pair list.

Open questions before this leaves draft

1. bf16 accuracy is not validated against a pretrained model. Per-layer rounding is
   ~2.5e-3; across a 22-layer stack with normalisation and residuals it compounds to
   ~1.3e-2 (cosine similarity 0.99994, argmax agreement 98.6%) — but that is synthetic
   weights, not a real mIoU comparison. This should gate any accuracy-sensitive use.
2. BIG grids are only exercised synthetically. The 3.05e9-cell path is tested with
   generated coordinates, not a real SemanticKITTI frame.
3. The host side uses torch for allocation, torch.sort and the compaction scan; the GPU
   kernels are pure FlyDSL. Replacing those host ops with hand-written kernels was tried
   and measured 2-8x slower (torch's nonzero/cumsum are parallel scans; the
   replacement serialised on atomics). Flagging in case the repo prefers a different
   dependency posture.

🤖 Generated with Claude Code