fx.gemm uses more VGPR than raw MFMA, and fails to compile with BufferCopy* atoms for fp8 M=64/N=16/K=128

[fsx950223]

Summary

While optimizing kernels/pa_decode_tile.py's QK matmul (raw mfma_f32_16x16x32_fp8_fp8 intrinsics, M=64/N=16/K=128 fp8), I compared it against the equivalent computation expressed via the high-level fx.gemm + make_tiled_mma + make_tiled_copy_A/B/C API. Two issues showed up:

1. fx.gemm uses more VGPR than the equivalent raw MFMA sequence for the same M/N/K/dtype (confirmed, modest but real: +2 VGPR / +8% in an isolated microbenchmark).
2. fx.gemm's fragment/copy path fails to compile at all when paired with fx.rocdl.BufferCopy{32,64,128}b (the AMD buffer-descriptor copy atoms recommended in the docs for global memory access) for this fp8 M=64/N=16/K=128 shape — only fx.UniversalCopy32b (plain-pointer copy) worked.

Repro

Attached: gemm_vgpr_compare.py (self-contained, no other repo files needed beyond flydsl/torch).

FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 gemm_vgpr_compare.py raw    # raw MFMA version
FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 gemm_vgpr_compare.py gemm  # fx.gemm version

Both kernels do the same M=64,N=16,K=128 fp8 mfma_f32_16x16x32_fp8_fp8 GEMM with a 4-warp (NWARP,1,1) tiled_mma split (matching pa_decode_tile.py's own QK tiling).

Issue 1 — VGPR overhead

.amdhsa_next_free_vgpr from the dumped final ISA (FLYDSL_DUMP_IR=1):

Version	VGPR
raw MFMA intrinsics	24
fx.gemm (with fx.UniversalCopy32b, since BufferCopy* doesn't compile — see below)	26

Reproducible across repeated runs. The gap is small in this toy kernel (no other loop-carried state), but in the real pa_decode_tile.py context, VGPR is the kernel's binding occupancy constraint (waves/SIMD = 512 // vgpr), so even a modest per-operand overhead compounds across a kernel with many such GEMM calls and loop-carried fragments — worth knowing about when choosing between the two APIs for register-pressure-sensitive kernels.

Issue 2 — BufferCopy* + fx.gemm fragments fail to compile

Swapping fx.UniversalCopy32b() for fx.rocdl.BufferCopy128b() in the gemm variant (same script, same shapes) hits a C++ assertion:

python3: include/flydsl/Dialect/Fly/Utils/LayoutUtils.h:791: std::pair<_FIter, _FIter> mlir::fly::detail::complementImpl(...): Assertion `(R == 1 || filteredStride.isStatic()) && "stride must be static for complement"' failed.

Switching to fx.rocdl.BufferCopy64b() (matching the fragment's 8-element/lane granularity) avoids the assertion but hits an MLIR legalization failure instead:

error: failed to legalize operation 'fly.copy_atom_call_ssa' that was explicitly marked illegal:
  %80 = "fly.copy_atom_call_ssa"(%55, %78) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>}> :
    (!fly.copy_atom<!fly_rocdl.cdna3.buffer_copy<64>, 8>, !fly.memref<f8E4M3FNUZ, global, 8:1>) -> vector<8xi8>

fx.rocdl.BufferCopy32b() produces the same class of legalization failure (just with vector<4xi8> / buffer_copy<32>).

Only fx.UniversalCopy32b() (plain llvm.load-based copy, no buffer-resource descriptor) compiles successfully for this fp8 shape. Given the docs (docs/cute_layout_algebra_guide.md §6.2) present BufferCopy* as the recommended AMD-specific path for efficient global memory access, and examples/03-tiledMma.py's own working example uses BufferCopy32b successfully for fp32 — this looks like a gap specific to fp8 element types (1-byte) combined with make_tiled_copy_A/B + thr_mma.make_fragment_A/B for this particular MNK/warp-tiling combination, rather than BufferCopy* being broadly broken.

Environment

• gfx942 (MI300X-class)
• Repro'd via FLYDSL_DUMP_IR=1 FLYDSL_RUNTIME_ENABLE_CACHE=0

Context

Found while working on pa_decode_tile.py (a hand-tuned reference PA-decode kernel that intentionally uses raw MFMA intrinsics instead of fx.gemm for its QK/PV matmuls, for exactly this kind of register-pressure reason — this issue is filed to check whether the raw-intrinsic-over-fx.gemm choice there is also masking an underlying fx.gemm/BufferCopy compile gap worth fixing independently).

[fsx950223]

Repro script (gemm_vgpr_compare.py) referenced above:

# SPDX-License-Identifier: Apache-2.0
"""Isolated microbenchmark: does fx.gemm cost more VGPR than raw MFMA
intrinsics for an equivalent fp8 16x16x32 GEMM (M=64,N=16,K=128, matching
pa_decode_tile's QK per-chunk shape)?
"""
import os

os.environ["FLYDSL_DUMP_IR"] = "1"
os.environ["FLYDSL_DUMP_DIR"] = "/tmp/dump_gemm_vgpr"
os.environ["FLYDSL_RUNTIME_ENABLE_CACHE"] = "0"

import torch  # noqa: E402

import flydsl.compiler as flyc  # noqa: E402
import flydsl.expr as fx  # noqa: E402
from flydsl.expr.typing import T  # noqa: E402

MFMA_MNK = 16
MFMA_K = 32
WAVE = 64
NWARP = 4
M = 64  # tokens
N = 16  # qheads
K = 128  # head_dim
FP8 = fx.Float8E4M3FNUZ


@flyc.kernel(known_block_size=(NWARP * WAVE, 1, 1))
def raw_mfma_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor):
    tid = fx.Int32(fx.gpu.thread_id("x"))
    warp = tid // WAVE
    lane = tid - warp * WAVE
    lane16 = lane - (lane // 16) * 16
    rgroup = lane // 16

    A_buf = A
    B_buf = B

    def load16(buf, byte_off):
        copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), FP8)
        reg = fx.make_rmem_tensor(fx.make_layout(16, 1), FP8)
        flat = fx.Tensor(fx.make_view(fx.get_iter(buf), fx.make_layout(1 << 30, 1)))
        div = fx.logical_divide(flat, fx.make_layout(1, 1))
        fx.copy(copy_atom, fx.slice(div, (None, byte_off)), reg)
        return fx.Vector(fx.memref_load_vec(reg)).bitcast(fx.Int64)

    N_SUBCHUNKS = K // 32  # 4
    acc = fx.arith.constant_vector(0.0, T.f32x4)
    for qkhe in range(N_SUBCHUNKS // 2):
        a_off = (warp * 16 + lane16) * K + (qkhe * 4 + rgroup) * 16
        a_w = load16(A_buf, a_off)
        b_off = lane16 * K + (qkhe * 4 + rgroup) * 16
        b_w = load16(B_buf, b_off)
        acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [a_w[0], b_w[0], acc, 0, 0, 0])
        acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [a_w[1], b_w[1], acc, 0, 0, 0])

    out = fx.Vector(acc)
    row = warp * 16 + lane16
    for r in range(4):
        col = rgroup * 4 + r
        C[row, col] = out[r]


@flyc.kernel(known_block_size=(NWARP * WAVE, 1, 1))
def gemm_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor):
    tid = fx.Int32(fx.gpu.thread_id("x"))
    bid = fx.Int32(fx.gpu.block_id("x"))

    bA = fx.zipped_divide(A, (M, K))
    bB = fx.zipped_divide(B, (N, K))
    bC = fx.zipped_divide(C, (M, N))
    bA = fx.slice(bA, (None, bid))
    bB = fx.slice(bB, (None, bid))
    bC = fx.slice(bC, (None, bid))

    mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(MFMA_MNK, MFMA_MNK, MFMA_K, FP8))
    tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((NWARP, 1, 1), (1, 0, 0)))
    thr_mma = tiled_mma.get_slice(tid)

    copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), FP8)
    tiled_copy_A = fx.make_tiled_copy_A(copy_atom, tiled_mma)
    tiled_copy_B = fx.make_tiled_copy_B(copy_atom, tiled_mma)
    tiled_copy_C = fx.make_tiled_copy_C(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), tiled_mma)

    thr_copy_A = tiled_copy_A.get_slice(tid)
    thr_copy_B = tiled_copy_B.get_slice(tid)
    thr_copy_C = tiled_copy_C.get_slice(tid)

    frag_A = thr_mma.make_fragment_A(bA)
    frag_B = thr_mma.make_fragment_B(bB)
    frag_C = thr_mma.make_fragment_C(bC)

    fx.copy(copy_atom, thr_copy_A.partition_S(bA), thr_copy_A.retile(frag_A), pred=None)
    fx.copy(copy_atom, thr_copy_B.partition_S(bB), thr_copy_B.retile(frag_B), pred=None)

    frag_C.fill(0.0)
    fx.gemm(mma_atom, frag_C, frag_A, frag_B, frag_C)

    fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), thr_copy_C.retile(frag_C), thr_copy_C.partition_D(bC), pred=None)


@flyc.jit
def launch_raw(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
    raw_mfma_kernel(A, B, C).launch(grid=(1, 1, 1), block=(NWARP * WAVE, 1, 1), stream=stream)


@flyc.jit
def launch_gemm(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
    gemm_kernel(A, B, C).launch(grid=(1, 1, 1), block=(NWARP * WAVE, 1, 1), stream=stream)


def main():
    dev = "cuda"
    A = torch.zeros(M, K, dtype=torch.float8_e4m3fnuz, device=dev)
    B = torch.zeros(N, K, dtype=torch.float8_e4m3fnuz, device=dev)
    C = torch.zeros(M, N, dtype=torch.float32, device=dev)

    import sys

    which = sys.argv[1] if len(sys.argv) > 1 else "gemm"
    if which == "raw":
        launch_raw(A, B, C)
    else:
        launch_gemm(A, B, C)
    torch.cuda.synchronize()


if __name__ == "__main__":
    main()

[fsx950223]

Re-verified on latest main (41b8f90a, includes #824) — built FlyDSL fresh against this exact commit (not reusing a stale binary) to make sure the repro is current.

Issue 1 (VGPR overhead) — reproduces identically:

• raw MFMA intrinsics: .amdhsa_next_free_vgpr = 24
• fx.gemm (UniversalCopy32b): .amdhsa_next_free_vgpr = 26

Issue 2 (BufferCopy* + fx.gemm fragments) — core bug still present, but BufferCopy128b's failure mode has changed:

• BufferCopy64b and BufferCopy32b still fail exactly as originally reported, same MLIR legalization error:

  error: failed to legalize operation 'fly.copy_atom_call_ssa' that was explicitly marked illegal:
    ... (!fly.copy_atom<!fly_rocdl.cdna3.buffer_copy<64>, 8>, !fly.memref<f8E4M3FNUZ, global, 8:1>) -> vector<8xi8>
  

  (and the analogous buffer_copy<32> / vector<4xi8> variant for BufferCopy32b).
• BufferCopy128b no longer hits the LayoutUtils.h:791 C++ assertion (stride must be static for complement) that crashed the process. On latest main it now fails the same way as 64b/32b — a clean, catchable DSLCompileError from the fly.copy_atom_call_ssa legalization pass instead of an assertion abort.

This looks consistent with include/flydsl/Dialect/Fly/Utils/LayoutUtils.h having been touched since this issue was filed by #784 ("Add coalesceX to keep layout boundary"), which changed the coalesceImpl/complementImpl logic right around the code the old assertion fired in. So the crash-on-128b regression is gone, but the underlying incompatibility between BufferCopy* and fx.gemm's fragment/copy path for this fp8 M=64/N=16/K=128 shape is not fixed — all three BufferCopy{32,64,128}b widths still fail to compile, just now uniformly via the legalization pass rather than a mix of assertion-crash (128b) and legalize-failure (64b/32b).