[MFMA] Add 16x16x16 bf16/f16 support with fly-fix-bitcast-width pass

[RichardChamberlain1]

Summary

• Fix MFMA 16x16x16 bf16/f16 lowering crash when BufferCopy128b loads 128-bit values that feed 64-bit MFMA operands
• Add fly-fix-bitcast-width compiler pass to prevent MLIR canonicalizer from folding extract chains into width-mismatched bitcasts
• Add matchWidth helper in CDNA3/CDNA4 emitAtomCallSSA for SSA-path width narrowing
• Add narrowToMmaWidth in ConvertAtomCallToSSAForm for early operand narrowing
• Add FileCheck tests for 16x16x16 bf16/f16 (same-width and wider-than-MFMA-operand inputs)

Bug

When BufferCopy128b loads i128 and the 16x16x16 MFMA expects vector<4xi16> (64 bits), the MLIR canonicalizer folds:

buffer_load(i128) -> bitcast(vector<8xbf16>) -> extract_strided_slice([0:4]) -> bitcast(vector<4xi16>)

into:

bitcast(i128 -> vector<4xi16>)

which LLVM rejects with "Invalid cast!" (128 != 64 bits).

Fix

Three layers of defense:

1. emitAtomCallSSA matchWidth (CDNA3/MmaAtom.cpp, CDNA4/MmaAtom.cpp): When the incoming SSA value is a wider vector than the MFMA operand, bitcast to the target element type at full width then vector.extract_strided_slice to narrow. Handles the SSA-only path (FileCheck tests).

2. ConvertAtomCallToSSAForm narrowToMmaWidth (ConvertAtomCallToSSAForm.cpp): After loading the register value via PtrLoadOp, compares the loaded type's width against the MMA atom's expected operand width (from getThrValLayoutA/B() cosize). Inserts ExtractStridedSliceOp if wider.

3. fly-fix-bitcast-width pass (FixBitcastWidth.cpp): Runs before the canonicalizer in the ROCm pipeline. Walks all vector.extract_strided_slice ops whose source traces back through bitcasts to a wider integer type (e.g. i128 from rocdl.raw.ptr.buffer.load). Inserts llvm.freeze on the extract result, which blocks the canonicalizer from folding the extract chain back to the wider source. freeze is an identity on well-defined values so correctness is preserved.

Files changed

File	Change
lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp	matchWidth lambda replaces direct bitcast
lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp	Same matchWidth fix
lib/Dialect/Fly/Transforms/FixBitcastWidth.cpp	New — fly-fix-bitcast-width pass
lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp	narrowToMmaWidth helper
lib/Dialect/Fly/CMakeLists.txt	Register new pass source
include/flydsl/Dialect/Fly/Transforms/Passes.td	Pass tablegen definition
python/flydsl/compiler/backends/rocm.py	Add pass to pipeline (before canonicalize)
tests/mlir/Conversion/mma_atom.mlir	4 new FileCheck tests (bf16/f16 x same-width/wide)
tests/mlir/Conversion/mma_atom_16x16x16.mlir	New — standalone 16x16x16 FileCheck tests

Verification

All three configurations produce correct results with matching LLVM/FlyDSL build:

Config	max_err	cos	Status
mma_k=32, block_n=32	0.004	1.0	PASS (no regression)
mma_k=16, block_n=32	0.004	1.0	PASS
mma_k=16, block_n=64	0.0002	1.0	PASS

Test plan

• bash scripts/build.sh -j64 compiles cleanly
• FileCheck: fly-opt tests/mlir/Conversion/mma_atom.mlir and mma_atom_16x16x16.mlir pass
• bash scripts/run_tests.sh — no regressions
• Device: mma_k=32 block_n=32 produces correct attention output
• Device: mma_k=16 block_n=32 produces correct attention output
• Device: mma_k=16 block_n=64 produces correct attention output