[Kernel] Fix the RDNA3 GEMM grid swizzle and migrate it to the tiled copy/MMA API

[vlluvia]

The part that matters most

rdna3_f16_gemm computes a wrong C, or faults the GPU, at ordinary shapes on today's main.

The grid swizzle derives bid_m from a fixed group width of 8 that need not divide grid_m, so the trailing partial group addresses tiles past the end of the grid. At the default 128x128x32 tile on gfx1100, bf16 in and out:

M = N	grid_m	8 divides it	Result on main	Max abs error
1024	8	yes	correct	0.0039
1152	9	no	wrong C	1.3295
1280	10	no	wrong C	1.5940
1536	12	no	GPU page fault	—
1664	13	no	wrong C	1.6725
2048	16	yes	correct	0.0039
2560	20	no	GPU page fault	—
3072	24	yes	correct	0.0039

K is 1024 throughout; each shape ran in its own process so a fault would not take the sweep with it. A correct bf16 result lands at 0.0039, so the failures are ~400x the rounding floor rather than a precision artifact. The damage is always confined to output tile rows at index 8 and above — the remainder group — and never touches the first eight, which is the swizzle's signature. Whether a given shape produces a wrong answer or a hard fault depends on whether the out-of-bounds address happens to be mapped: 1536 faulted in one harness and returned wrong numbers in another. The silent wrong answer is the common case.

_group_width now snaps the width down to the largest divisor of grid_m within the cap, which restores the bijection. Every shape above is correct after this PR.

The rest: migration to the layout API

Replaces the hand-rolled index arithmetic, which is how the swizzle bug surfaced.

Epilogue. C is now partitioned through a tiled_mma and make_tiled_copy_C instead of an explicit g_row = base + 2*si + klane store loop. The tiled_mma carries a permutation that reproduces the kernel's wave banding — the default stamping interleaves the repeats instead, which measured 62% slower at 3072x3072x1024, so the permutation is load-bearing rather than cosmetic.

GMEM to LDS. Now goes through flat_divide plus a tiled copy with BufferCopy128b, replacing the precomputed offset tables and the flat _v8_store helper. The thread geometry is the same assignment the tables computed.

Grid swizzle. Lifted out of the kernel body into _group_width and _swizzle_tile_id, which is what makes it unit-testable without a GPU.

Performance: unchanged, by construction

The block tile stays a parameter defaulting to 128x128x32, and the generated ISA is byte-identical for rn and rs output across 128x128x32, 64x64x64 and 128x64x32 on gfx1100.

Measured anyway, against the unmodified kernel over 16 shapes where main is correct enough to time: geomean 1.002x, spread 0.978x to 1.034x. Since both sides run identical machine code, that spread is the harness noise floor rather than an effect.

Tests

• test_rdna_gemm — device-level regression at grid_m of 3, 9 and 20, with an assert grid_m % 8 guard so the shapes cannot silently stop covering the bug.
• test_rdna3_grid_swizzle — the bijection, over the grids that faulted (12, 20) and the grids that returned wrong answers (9, 10, 13). No GPU needed.
• test_rdna3_wmma_atom — the gfx11 tiled copy atoms, which have to handle lanes 16-31 mirroring lanes 0-15 under the v16 operand ABI.

Full suite on gfx1100 (Radeon Pro W7900): 1609 passed, 0 failed.

[vlluvia]

Follow-up #980 stacks on this branch and adds shape-driven tile selection on top. It targets main only because the base branch has to live in this repo, so its diff currently replays this commit; it will collapse to its own 5 files once this merges.

[coderfeli]

Looks good. But style check failed. @vlluvia

[vlluvia]

Looks good. But style check failed. @vlluvia

Sorry, that’s on me. I’ll get the style check fixed.