[Kernel][Dialect] Support RDNA 3 / 3.5 WMMA (PoC): fix atom + add baseline GEMM kernel

[KerwinTsaiii]

Motivation

This PR adds initial PoC-level support for AMD RDNA 3 / 3.5 wave32 WMMA on FlyDSL's atom path, validated on gfx1151 (Strix Halo iGPU, Radeon 8060S).

The goals are:

1. Make the MmaOpRDNA3_WMMAType atom functionally and numerically correct on gfx110x / gfx115x — the previous registration had a C/D layout mismatch with the actual hardware ordering and a suspect bf16-acc lowering, both papered over at runtime by ~28 ds_bpermute instructions per WMMA.
2. Demonstrate the atom path works end-to-end on production-shape GEMM (≤ 4096³) by adding a multi-wave + LDS double-buffered kernel, so other RDNA 3 / 3.5 kernels can be built on the same primitives.
3. Provide a rocBLAS-comparison benchmark so future perf work has a tracked baseline (currently ≈ 56% of rocBLAS hgemm at 4096³).

This is a PoC baseline, not a competitive flagship kernel. For peak GEMM perf on RDNA 3 / 3.5 today users should fall back to PyTorch /rocBLAS hgemm. Closing the gap to rocBLAS is left for follow-up PRs (see the Out of scope list below).

Technical Details

Atom-level fixes (MmaOpRDNA3_WMMAType)

Three issues found while reviewing the existing atom registration against the gfx11 ISA + empirically validating on gfx1151 hardware:

• C/D layout mismatch. The f32/i32 accumulator's per-lane natural M ordering is interleaved between lane groups (lane n, val v → M = 2·v + n/16), not contiguous halves as getThrValLayoutCD declared. The mismatch was being papered over by calling reorderAccLaneValues four times per WMMA call (3× on the C input as a hand-rolled inverse permutation, 1× on the D output), costing ~32 ds_bpermute per call. Updating getThrValLayoutCD to match the real layout and dropping the runtime swizzles takes the per-WMMA cross-lane op count from ~28 → 8 (the irreducible A/B expansion for the WMMA256b 16-wide K input).

• bf16-acc path. The native bf16-acc WMMA variant uses op_sel-packed 16-wide C/D lanes which conflicted with FlyDSL's vec8 fragment representation. Tests passed because verify_output uses cosine-similarity with atol=0.1 and only exercised 16×16×16 — almost certainly silently wrong on larger shapes. Now lowers through the f32 WMMA op + fptrunc, matching the f16-acc fallback. (BF16 has the same exponent range as f32, so the promote/truncate round-trip is overflow-safe.)

• Dead code. Remove the now-unused duplicateTo16WideSimple and reorderAccLaneValues helpers (~110 LOC of MmaAtom.cpp).

The MLIR FileCheck test (tests/mlir/Conversion/wmma_rdna3.mlir) is updated for the new bf16 → f32 lowering.

New PoC baseline GEMM kernel — kernels/rdna3_gemm.py

A multi-wave atom-based GEMM exercising MmaOpRDNA3_WMMAType on production-shape sizes. Reaches ≈ 56% of rocBLAS hgemm at 4096³ on gfx1151 (≈ 30% of WMMA F16 peak). Explicitly framed as a PoC baseline, not a peak-perf kernel — the docstring lists the optimizations needed to close the gap to rocBLAS (LDS XOR swizzle, CShuffle epilogue, sched intrinsics, occupancy tuning, tail-tile predication) for follow-up PRs.

Atom-API optimizations applied:

• 2×2 wave layout (4 × wave32 = 128 threads/workgroup), per-warp reg_m × reg_n = 4×4 WMMA tile, BLOCK = 128 × 128 × 32.
• LDS staging for both A and B with a 2-stage ping-pong buffer; LDS layout is row-major (K-fastest) to match the row-major GMEM source, layout-lowering pass handles the reshape into the column-major WMMA fragment convention.
• Software pipelining: each loop iteration prefetches the next K-tile from GMEM into LDS[(k+1)&1] while computing from LDS[k&1], overlapping GMEM latency with WMMA compute.
• Outer loop is a runtime SCF for-op iterating over K-tile pairs (so the IR doesn't blow up at K=4096 from unrolling 128 iterations); the pair body unrolls stage 0 and stage 1 explicitly so the LDS-stage index stays constexpr.
• Implementation note: make_tiled_mma's third argument is the K-direction permutation, not a per-warp reg-tile knob. reg_m / reg_n / reg_k are implicit from the ratio BLOCK_M / (waves_m * WMMA_M).

Tests + benchmark infrastructure

• Atom validation (tests/kernels/test_rdna3_wmma_gemm.py): expanded 17 → 25 cases — bf16-acc multi-K-tile (was only 16×16×16) and multi-block sizes; f16-acc multi-K-tile; iu8 → i32 multi-block (was only single-tile contract test).

• Production-shape correctness (tests/kernels/test_rdna3_gemm.py, new): 19 tests covering f16/bf16 inputs, f32/f16/bf16 accumulators, multi-block sizes (128³ → 1024³), and a wave / reg-config sweep.

• rocBLAS-comparison benchmark (tests/bench_rdna3_gemm.py, new): reports TFLOPS, %rocBLAS, and %WMMA-peak. Includes a chip table for back-calculating theoretical peak from per-CU per-cycle WMMA throughput (511 ops/cycle/CU on RDNA 3, derived from RX 7900 XTX's published 122.6 TFLOPS / 96 CUs / 2.5 GHz).

Test Plan

Validated on gfx1151 (AMD Ryzen AI Max+ 395, Radeon 8060S iGPU, ROCm 7.12).

# Atom validation (25 tests covering all 6 dtype combos + multi-tile)
python -m pytest tests/kernels/test_rdna3_wmma_gemm.py -v

# Production-shape GEMM correctness (19 tests covering dtypes,
# accumulators, multi-block sizes, and a wave/reg-config sweep)
python -m pytest tests/kernels/test_rdna3_gemm.py -v

# Benchmark 
bash scripts/run_benchmark.sh 
[run_benchmark] GPU arch: gfx1151 (CDNA=false, RDNA4=false)
========================================================================
Benchmarks (logs under /tmp/flydsl_bench)
========================================================================

op                     shape                              dtype            TB/s     TFLOPS
---------------------- ---------------------------------- ---------- ---------- ----------
softmax                32768x8192                         bf16            0.209          -
layernorm              32768x8192                         bf16            0.226          -
rmsnorm                32768x8192                         bf16            0.225          -

========================================================================
Benchmark Summary
========================================================================
Total: 3 tests
Success: 3
Failed: 0
Logs: /tmp/flydsl_bench

All benchmarks passed! 

[sjfeng1999]

I am not familiar with WMMA in RDNA3, but this part looks so strange. Generating so many ops inside AtomCall is not a good idea, since it may replicate multiple times when gemmOp is unpacked into atom calls.

If this wrapping step is the common way for RDNA3 WMMA, I would suggest giving FlyROCDL_MmaOpRDNA3_WMMA a wrapper-specific suffix/name, and reserving the instruction-matching name for an unwrapped version which AtomCall basically is a rocdl.wmma call. This would leave room for future that users could perform these cross-lane operations explicitly at the Python level if they want.

[KerwinTsaiii]

Thanks for point out this

With the lane mapping using matrix_calculator.py, and RDNA3 is the odd one here: in wave32, A/B need replication between lane i and i+16. For example:

./matrix_calculator.py --architecture rdna3 --instruction v_wmma_f16_16x16x16_f16 --A-matrix --get-register --I-coordinate 3 --K-coordinate 2
Architecture: RDNA3
Instruction: V_WMMA_F16_16X16X16_F16
A[3][2] = v1{3}.[15:0]
A[3][2] = v1{19}.[15:0]

Same query on RDNA4 and CDNA3 returns only one lane:

./matrix_calculator.py --architecture rdna4 --instruction v_wmma_f16_16x16x16_f16 --A-matrix --get-register --I-coordinate 3 --K-coordinate 2
Architecture: RDNA4
Instruction: V_WMMA_F16_16X16X16_F16
A[3][2] = v1{3}.[15:0]

./matrix_calculator.py --architecture cdna3 --instruction v_mfma_f32_16x16x16_f16 --A-matrix --get-register --I-coordinate 3 --K-coordinate 2
Architecture: CDNA3
Instruction: V_MFMA_F32_16X16X16_F16
A[3][2] = v1{3}.[15:0]

The reason I implemented it this way is that I wanted to keep the current FlyDSL compact fragment model (vec8 split across lane groups) unchanged and compatible with the existing tiling/copy pipeline. Given that model, reconstructing the operand cross-lane inside emitAtomCallSSA was the most direct way to satisfy RDNA3 WMMA operand ABI requirements.

But I also agree this weakens the design contract of MmaOpRDNA3_WMMA, because it is no longer a clean instruction-matching atom and now embeds cross-lane wrapper behavior in the atom call path.

Given this affects atom semantics (instruction-matching vs wrapper behavior), I don’t think I should make that call solo in this PR. I’m going to close this for now, align on direction first (raw-only / wrapper-only / dual-mode), then continue in a follow-up.