[FlyDSL] [Perf] [gfx950] Swap A/B in the mxfp8 8-wave epilogue for 64-bit stores

aiter/ops/flydsl/gemm_a8w8_mxfp8_8wave.py

@@ -13,11 +13,13 @@
 
 from aiter.jit.utils.chip_info import get_lds_capacity_bytes
 
-# Fixed by the kernel: MFMA_Scale(16, 16, 128) over a 128-deep K tile, one
-# 256x256 block tile, and scales staged a 256-K pair at a time.
+# Fixed by the kernel: MFMA_Scale(16, 16, 128) over a 128-deep K tile and one
+# 256x256 block tile. Row-major scales are staged a chunk of K steps at a time
+# and need only whole K tiles; the preshuffled layout is addressed by K-pair,
+# so it needs 256 (checked separately, below).
 BLOCK_K = 128
 BLOCK_M = BLOCK_N = 256
-K_ALIGN = 256
+K_ALIGN = BLOCK_K
 SCALE_GROUP = 32
 
 # A/B stages: 4 buffers each of (BLOCK/2) x 128 fp8; the row layout adds two
@@ -66,9 +68,10 @@ def _validate(
     AssertionError escaping here would kill the tuner worker process rather than
     just failing the candidate.
     """
-    if K <= 0 or K % K_ALIGN != 0:
+    if K <= 0 or K % K_ALIGN != 0 or (scale_preshuffled and K % 256):
         raise ValueError(
-            f"[FlyDSL mxfp8 8wave] K must be a positive multiple of {K_ALIGN}, got {K}"
+            f"[FlyDSL mxfp8 8wave] K must be a positive multiple of "
+            f"{256 if scale_preshuffled else K_ALIGN}, got {K}"
         )
     if int(waves_per_eu) < 1:
         raise ValueError(

aiter/ops/flydsl/kernels/mxfp8_gemm_8wave.py

@@ -13,6 +13,8 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm as _llvm_d
 from flydsl.expr import const_expr, range_constexpr, rocdl
 from flydsl.expr.typing import Vector as Vec
 
@@ -189,7 +191,7 @@ def __init__(
             ),
         ]
 
-    def _soffset(self, op, chunk):
+    def _chunk_off(self, op, chunk):
         if const_expr(op == 0 and self.a_col):
             return fx.Int32(chunk * self.CHUNK) * self.c_m
         if const_expr(op == 1 and self.b_blk):
@@ -198,7 +200,13 @@ def _soffset(self, op, chunk):
         return fx.Int32(chunk * 16)
 
     def dma(self, chunk):
-        """Stage chunk ``chunk``; the last chunk may run past K (those bytes are never read)."""
+        """Stage chunk ``chunk``; past the tensor's end it reads zero, never consumed.
+
+        The chunk offset goes in the per-lane offset: the raw-buffer range check
+        covers the VGPR and immediate offsets but not ``soffset``, so through
+        ``soffset`` the last chunk would read beyond the tensor. Not the
+        immediate either: on an LDS-direct load it also moves the destination.
+        """
         base_i32 = fx.Int32(fx.ptrtoint(self.lds.ptr))
         for op in range_constexpr(2):
             src, offs = self.srcs[op]
@@ -214,17 +222,17 @@ def dma(self, chunk):
                     rocdl.buffer_load_to_lds(
                         src,
                         fx.to_llvm_ptr(lds_dst),
-                        fx.Int32(offs[step]).ir_value(),
+                        (fx.Int32(offs[step]) + self._chunk_off(op, chunk)).ir_value(),
                         size_bytes=1,
-                        soffset=self._soffset(op, chunk).ir_value(),
                     )
                     continue
                 dst = fx.make_view(lds_dst, fx.make_layout(1, 1))
                 fx.copy(
                     self.atom,
-                    fx.slice(src, (None, fx.Int32(offs[step]))),
+                    fx.slice(
+                        src, (None, fx.Int32(offs[step]) + self._chunk_off(op, chunk))
+                    ),
                     dst,
-                    soffset=self._soffset(op, chunk),
                 )
 
     def _byte(self, op, off):
@@ -317,8 +325,79 @@ def call(self, a, b, c, sa, sb, *, k_pack, set_prio=True):
         return [cf.load().ir_value() for cf in c_frags]
 
 
-class StoreBf16:
-    """Converts the already-scaled fp32 accumulators to bf16 and stores them."""
+class SwappedMxMfma:
+    """``MxMfma`` with A/B -- and their scales and opsels -- exchanged.
+
+    ``B^T A^T = (A B)^T``, so every accumulator comes out transposed and a lane
+    ends up holding four consecutive *columns* instead of four rows. That is
+    what lets the epilogue store 64 bits at a time; see ``StoreBf16Swapped``.
+
+    The inner atom is built with the tile counts swapped, since its "A" is our
+    B -- its own asserts and ``idx`` then line up without further adjustment.
+    ``idx`` forwards swapped so callers keep indexing in (M-tile, N-tile) order.
+    """
+
+    def __init__(self, n_tiles_a, n_tiles_b, paired_tiles=True):
+        self._inner = MxMfma(n_tiles_b, n_tiles_a, paired_tiles=paired_tiles)
+        self.zero_value = self._inner.zero_value
+        self.n_tiles_a = n_tiles_a
+        self.n_tiles_b = n_tiles_b
+
+    def idx(self, i, j):
+        return self._inner.idx(j, i)
+
+    def call(self, a, b, c, sa, sb, *, k_pack, set_prio=True):
+        return self._inner.call(b, a, c, sb, sa, k_pack=k_pack, set_prio=set_prio)
+
+
+def _raw(v):
+    """Unwrap a DSL value to its raw ``ir.Value``, for ops that take one."""
+    return v.ir_value() if hasattr(v, "ir_value") else v
+
+
+def _permlane16_swap(x, y):
+    """gfx950 ``v_permlane16_swap_b32``: exchange 16-lane rows between two VGPRs.
+
+    Verified on hardware rather than assumed::
+
+        vdst_new = [X.r0, Y.r0, X.r2, Y.r2]
+        vsrc_new = [X.r1, Y.r1, X.r3, Y.r3]
+
+    Position *within* a row is preserved, so a lane keeps its ``lane % 16`` and
+    therefore its C row.
+    """
+    st = ir.Type.parse("!llvm.struct<(i32, i32)>")
+    res = fx.rocdl.permlane16_swap(st, _raw(x), _raw(y), False, True)
+    i32 = ir.IntegerType.get_signless(32)
+    return (
+        fx.Int32(
+            _llvm_d.ExtractValueOp(i32, res, ir.DenseI64ArrayAttr.get([0])).result
+        ),
+        fx.Int32(
+            _llvm_d.ExtractValueOp(i32, res, ir.DenseI64ArrayAttr.get([1])).result
+        ),
+    )
+
+
+def _ds_bpermute(value, src_lane):
+    """Pull ``value`` from ``src_lane``. Byte-addressed, hence the <<2."""
+    i32 = ir.IntegerType.get_signless(32)
+    return fx.Int32(
+        fx.rocdl.ds_bpermute(i32, _raw(fx.Int32(src_lane) * fx.Int32(4)), _raw(value))
+    )
+
+
+class StoreBf16Swapped:
+    """C store for an A/B-swapped MFMA: 64 bits per lane instead of 16.
+
+    After the swap lane ``l`` holds ``D[l%16][4*(l/16)+k]`` -- four consecutive
+    columns, i.e. eight contiguous bytes of a row-major C. The unswapped layout
+    gives a lane four consecutive *rows* of one column, which can only be
+    written one bf16 at a time.
+
+    The accumulators are already scaled by the MFMA, so the epilogue is a plain
+    fp32 -> bf16 convert.
+    """
 
     def __init__(self, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b):
         self.c_rows = c_rows
@@ -331,31 +410,105 @@ def __init__(self, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b):
             C, max_size=False, num_records_bytes=c_rows * c_cols * 2
         )
         self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1))
-        self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), fx.BFloat16)
-        self.reg_bf16_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.BFloat16)
+        self.out_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.BFloat16)
+        self.reg_bf16_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.BFloat16)
+        self.out_atom_8 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.BFloat16)
+        self.reg_bf16_8 = fx.make_rmem_tensor(fx.make_layout(8, 1), fx.BFloat16)
 
     def store(self, c_frag, base_row, base_col):
+        lane = self.lane_id
         for ti in range_constexpr(self.n_tiles_a):
-            row = base_row + ti * 16 + (self.lane_id // 16) * 4
+            row = base_row + ti * 16 + lane % 16
             for tj in range_constexpr(self.n_tiles_b):
-                col = base_col + tj * 16 + self.lane_id % 16
-                col_valid = col < self.c_cols
+                col = base_col + tj * 16 + (lane // 16) * 4
+                # A lane's four columns are consecutive, so one predicate covers
+                # them; out-of-range lands on the descriptor's OOB index.
                 oob = fx.Int32(self.c_rows * self.c_cols)
                 vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)])
-                for i in range_constexpr(4):
-                    fx.memref_store_vec(
-                        fx.full(1, vec_f32[i].to(fx.BFloat16), fx.BFloat16),
-                        self.reg_bf16_1,
-                    )
-                    c_index = (row + i) * self.c_cols + col
-                    fx.copy(
-                        self.out_atom_1,
-                        self.reg_bf16_1,
-                        fx.slice(
-                            self.c_div,
-                            (None, fx.Int32(col_valid.select(c_index, oob))),
-                        ),
-                    )
+                vals = [vec_f32[k].to(fx.BFloat16) for k in range_constexpr(4)]
+                idx = (col + 3 < self.c_cols).select(row * self.c_cols + col, oob)
+                fx.memref_store_vec(
+                    Vec.from_elements(vals, fx.BFloat16), self.reg_bf16_4
+                )
+                fx.copy(
+                    self.out_atom_4,
+                    self.reg_bf16_4,
+                    fx.slice(self.c_div, (None, fx.Int32(idx))),
+                )
+
+
+class StoreBf16Permlane(StoreBf16Swapped):
+    """One cross-lane step on top, giving 16B per lane and 64B rows.
+
+    After the swap a lane owns four consecutive columns of one N-tile, i.e.
+    eight bytes, and the two N-tiles it holds are 16 columns apart -- so a row
+    still only gets 32 bytes per instruction. Two ``permlane16_swap`` fix that
+    exactly, with no ``ds_bpermute`` and no LDS::
+
+        (A, B) = permlane16_swap(tile0.d0, tile1.d0)
+        (C, D) = permlane16_swap(tile0.d1, tile1.d1)
+        lane group g stores (A, C, B, D)
+
+    which lands g = 0,1,2,3 on columns 0-7, 16-23, 8-15, 24-31. The four groups
+    together cover columns 0..31 contiguously: 64 bytes per row, from 16 bytes
+    per lane. The column permutation is absorbed into the address, so it is
+    free.
+    """
+
+    _lane_transpose = False
+
+    def store(self, c_frag, base_row, base_col):
+        assert self.n_tiles_b == 2, (
+            "the permlane mapping pairs exactly two N-tiles (BLOCK_N == 256); "
+            f"got n_tiles_b={self.n_tiles_b}"
+        )
+        lane = self.lane_id
+        grp = lane // 16
+        for ti in range_constexpr(self.n_tiles_a):
+            row = base_row + ti * 16 + lane % 16
+            dwords = []
+            for tj in range_constexpr(self.n_tiles_b):
+                vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)])
+                packed = Vec.from_elements(
+                    [vec_f32[k].to(fx.BFloat16) for k in range_constexpr(4)],
+                    fx.BFloat16,
+                ).bitcast(fx.Int32)
+                dwords.append((packed[0], packed[1]))
+            a, b = _permlane16_swap(dwords[0][0], dwords[1][0])
+            c, d = _permlane16_swap(dwords[0][1], dwords[1][1])
+            if const_expr(self._lane_transpose):
+                # The permlane stage leaves a row's four 8-column chunks in
+                # lanes 16 apart, so a 16-lane group touches 16 rows at 16 bytes
+                # each. Transposing the lane index -- lane l' = 4r'+q' pulls
+                # from the lane holding (row r', chunk q') -- puts adjacent
+                # lanes on one row, so lanes 0-3 write 64 contiguous bytes and a
+                # group covers 4 rows instead of 16.
+                #
+                # chunk q' lives at column q'*8, and the permlane stage put
+                # column (g%2)*16 + (g//2)*8 in group g, so g is the swap of
+                # q's two bits: 0,1,2,3 -> 0,2,1,3.
+                q = lane % 4
+                src_lane = ((q % 2) * 2 + q // 2) * 16 + lane // 4
+                a, c, b, d = (_ds_bpermute(v, src_lane) for v in (a, c, b, d))
+                row = base_row + ti * 16 + lane // 4
+                col = base_col + q * 8
+            else:
+                col = base_col + (grp % 2) * 16 + (grp // 2) * 8
+            out8 = Vec.from_elements([a, c, b, d], fx.Int32).bitcast(fx.BFloat16)
+            oob = fx.Int32(self.c_rows * self.c_cols)
+            idx = (col + 7 < self.c_cols).select(row * self.c_cols + col, oob)
+            fx.memref_store_vec(out8, self.reg_bf16_8)
+            fx.copy(
+                self.out_atom_8,
+                self.reg_bf16_8,
+                fx.slice(self.c_div, (None, fx.Int32(idx))),
+            )
+
+
+class StoreBf16LaneTranspose(StoreBf16Permlane):
+    """``StoreBf16Permlane`` plus the ds_bpermute lane transpose."""
+
+    _lane_transpose = True
 
 
 def compile_mxfp8_gemm_8w(
@@ -397,14 +550,18 @@ def compile_mxfp8_gemm_8w(
     assert (
         BLOCK_M == 256 and BLOCK_N == 256
     ), "MXFP8 8-wave is specialized to a 256x256 block tile"
-    assert (
-        K % 256 == 0
-    ), f"K must be a multiple of 256 (MX scale chunk staging), got {K}"
+    assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}"
 
     K_ITERS = K // BLOCK_K
-    # Scale words are addressed by K-pair, so K must contain a whole number of
-    # them (K % 256 above already guarantees it).
-    assert K_ITERS % 2 == 0
+    # Only the preshuffled reader addresses scale words by K-pair (k1 = k // 2),
+    # so only it needs a whole number of pairs. ScaleRawLds stages CHUNK steps
+    # at a time and its dma() already tolerates a partial last chunk.
+    if scale_preshuffled:
+        assert K % 256 == 0, (
+            f"preshuffled scales are addressed by K-pair, so K must be a "
+            f"multiple of 256, got {K}"
+        )
+        assert K_ITERS % 2 == 0
 
     N_TILES_A = BLOCK_M // 64
     N_TILES_B = BLOCK_N // 128
@@ -506,13 +663,13 @@ def kernel_gemm(
             lane_id, wave_id, K, N_LDS_ROUNDS, preshuffled=b_preshuffled
         )
 
-        mfma = MxMfma(N_TILES_A, N_TILES_B, paired_tiles=scale_preshuffled)
+        mfma = SwappedMxMfma(N_TILES_A, N_TILES_B, paired_tiles=scale_preshuffled)
 
         a_g2s = G2SLoader(a_div, gl_off_a, N_LDS_STEPS_A, F8_IR_t, wave_id)
         b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, F8_IR_t, wave_id)
         a_s2r = S2RLoader(wave_m, N_TILES_A)
         b_s2r = S2RLoader(wave_n, N_TILES_B)
-        store_c = StoreBf16(C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B)
+        store_c = StoreBf16LaneTranspose(C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B)
 
         if const_expr(scale_preshuffled):
             a_sc = ScalePreshuffledS2R(A_scale, c_m, K, N_TILES_A)