diff --git a/.claude/skills/4wave-mfma-coverage-analysis/SKILL.md b/.claude/skills/4wave-mfma-coverage-analysis/SKILL.md new file mode 100644 index 000000000..03976271b --- /dev/null +++ b/.claude/skills/4wave-mfma-coverage-analysis/SKILL.md @@ -0,0 +1,67 @@ +--- +name: 4wave-mfma-coverage-analysis +description: From an ATT trace, find which hot-loop instructions are NOT hidden behind MFMA execution -- i.e. what is keeping cyc/mfma above the MFMA execute floor (16 for fp4 16x16x128, 32 for fp8 16x16x128). Models the matrix unit as a pipeline (next_free) for the EXPOSED cycles, then tiles those idle cycles by OCCUPYING instruction (issue_dur + stall, intersected with the idle gaps) so the ranking says which non-MFMA ops steal issue bandwidth between MFMAs. Reports %exp and %all. Use when a GEMM/attention kernel is MFMA-bound and is 256 threads/4waves and you want to push cyc/mfma toward the floor, or which non-MFMA op is the biggest exposed stall. +--- + +# MFMA Coverage Analysis + +## The model + +**Step 1 — EXPOSED via next_free (matrix-unit pipeline).** Each MFMA occupies ONE EXEC-cycle execute slot, but consecutive MFMAs can **issue < EXEC apart** and still both be hidden (dense fp4 MFMAs issue (4~8) cyc apart yet each execute-slot is 16). +Track `next_free` = the cycle the unit becomes free: +- MFMA issues at `t <= next_free` → hidden; `next_free += EXEC` +- MFMA issues at `t > next_free` → unit IDLE for `t - next_free` → **EXPOSED** + +EXPOSED = sum of idle gaps; shrink toward 0 so cyc/mfma → EXEC(This replaces an older union-of-`[issue,issue+EXEC)` model that over-reported exposure.) + +**Step 2 — attribute EXPOSED by OCCUPANCY.** For every idle cycle, credit it to whichever instruction was on the issue port then. Each non-MFMA instruction occupies `[its issue, next instruction's issue)` = its issue_dur PLUS any stall (r[2]) — we do NOT split issue vs stall, and do NOT ask "who blocks". The question is simply: while the matrix unit sat idle, which non-MFMA ops were stealing issue bandwidth? Intersect each op's occupancy with the idle gaps and sum. This tiles the whole EXPOSED window (every idle cycle has an owner), so the ranking directly says which ops must be REMOVED / SHRUNK from between the MFMAs to push cyc/mfma toward the floor. + +Two percent columns: **%exp** (of the EXPOSED cycles) and **%all** (of the whole window incl. MFMA — the real wall-clock lever, since MFMA execute is ~85%+). + +Example (fp4_gemm_4wave, cyc/mfma 18.73, EXPOSED 6344 = 14.4% of all): `buffer_load_dwordx4 30%exp/4.3%all · v_readfirstlane 17%/2.5% · s_barrier 16%/2.3% · s_waitcnt(lgkmcnt) 8%/1.2% · s_add_i32 6%/0.8% · ds_read_b128 6%/0.8% · …`. +Read it as: the exposed time is scale/g2s **gather + readfirstlane + address arithmetic** stealing issue slots between MFMAs; cut their count/occupancy. +Row schema: `[cycle, issue_dur, stall, total_dur, code_id]`; s_waitcnt is split vmcnt / lgkmcnt so you see whether VMEM or LDS waits dominate. + +EXEC is the MFMA **execute** latency, NOT issue latency: +- fp4 `mfma_scale_f32_16x16x128_f8f6f4` → **16** +- fp8 `mfma..16x16x128` → **32** (its issue latency) +Pass the right `--exec`. + +## How to run + +```bash +# 1. capture ATT on an idle GPU +# 2. run, first without --range to print the cycle span: +python3 .claude/skills/mfma-coverage-analysis/scripts/mfma_coverage.py +# 3. pick a mid steady slice spanning ~10 outer-loop iters, re-run: +python3 .claude/skills/mfma-coverage-analysis/scripts/mfma_coverage.py \ + --range 219000,245000 --exec 16 +``` + +Output: MFMA-covered %, EXPOSED % (and % of all), cyc/mfma vs floor, then the exposed +cycles tiled by OCCUPYING instruction with two columns (%exp, %all). + +## Reading it + +- **Attack the top occupancy ops** = the non-MFMA instructions eating the most idle + cycles between MFMAs. Reduce their COUNT or per-op occupancy so more MFMAs pack in: + fewer/cheaper address ops, precompute wave-uniform values once (kill readfirstlane), + merge/spread loads, deepen prefetch so a value is on the port earlier. +- **buffer_load / ds_read high** → too many gathers/reads issue between MFMAs; merge + (wider load), move to a different MFMA's shadow, or cut the count. +- **v_readfirstlane high** → v→s serialization; precompute the wave-uniform SGPR once + outside the loop instead of per-use. +- **s_add / s_and / v_cmp / s_lshl** (address & loop arithmetic) → recompute less; hoist loop-invariants, use s_add increment chains, fold offsets into instruction immediates. +- **s_barrier high** → per-iter cross-wave sync (4 waves finish the iter at slightly different times). Structural for a fixed wave count; not cuttable without changing sync granularity / wave layout. +- **s_waitcnt(lgkmcnt/vmcnt)** → LDS/VMEM waits; deepen the corresponding prefetch or relax the count. Usually small once loads are hidden. + +## Caveats + +- Uses one wave file (se0_sm0_sl0_wv0.json by default). Different waves are + equivalent repeats; pass `--wave` to check another. +- Occupancy tiles the EXPOSED window with no gaps (every idle cycle has an owner), so + %exp sums to ~100. It counts issue_dur + stall together on purpose — the goal is + "what steals issue bandwidth between MFMAs", not "who stalls". +- Row schema: `[cycle, issue_dur, stall, total_dur, code_id]`; + `code.json["code"][cid][0]` is the asm text. (See att-hotloop-benchmark for the + complementary barrier-window cyc/mfma summary.) diff --git a/.claude/skills/4wave-mfma-coverage-analysis/scripts/mfma_coverage.py b/.claude/skills/4wave-mfma-coverage-analysis/scripts/mfma_coverage.py new file mode 100644 index 000000000..574a1b54e --- /dev/null +++ b/.claude/skills/4wave-mfma-coverage-analysis/scripts/mfma_coverage.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Find which instructions are NOT hidden behind MFMA in a GEMM hot loop. + +Model: each MFMA occupies a fixed execute window of EXEC cycles from its issue +cycle. Back-to-back MFMAs tile [c, c+EXEC) windows; while the matrix unit is busy +any co-issued scalar/VMEM/LDS op is "free". Cycles OUTSIDE the union of those +windows are EXPOSED -- the matrix unit idles there, so cyc/mfma > EXEC. We attribute +each exposed gap to the (non-MFMA) instruction at its start = what blocked the next +MFMA from issuing on time. + +Input: an ATT rocprofv3 UI dispatch dir (has code.json + se*_wv*.json). Pick the +steady-state cycle window with --range (avoid prologue/tail). + +Usage: + python3 mfma_coverage.py [--wave se0_sm0_sl0_wv0.json] + [--range LO,HI] [--exec 16] + + # find a steady window first (cycles): the script prints the wave cycle span; + # pick a mid slice spanning ~10 outer loop iterations. + +Notes: +- EXEC is the MFMA execute latency, NOT the issue latency. fp4 + mfma_scale_f32_16x16x128 ~ 16; fp8 16x16x128 ~ 32. Pass --exec accordingly. +- "idle" gaps = pure latency stalls (waitcnt drain / dependency) with no issuing + instruction; real wins come from cutting the named-instruction gaps. +""" +import argparse +import collections +import glob +import json +import os +import sys + + +def load(dispatch_dir, wave): + code = json.load(open(os.path.join(dispatch_dir, "code.json")))["code"] + if wave: + wpath = os.path.join(dispatch_dir, wave) + else: + cands = sorted(glob.glob(os.path.join(dispatch_dir, "se*_wv0.json"))) + wpath = cands[0] + wj = json.load(open(wpath)) + return code, wj["wave"]["instructions"], os.path.basename(wpath) + + +def op_of(code, cid): + a = code[cid][0].strip().split() if cid < len(code) else ["?"] + return a[0] if a else "?" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("dispatch_dir") + ap.add_argument("--wave", default=None) + ap.add_argument("--range", default=None, help="LO,HI cycle window (steady state)") + ap.add_argument("--exec", type=int, default=16, dest="exec_cyc", help="MFMA execute latency") + args = ap.parse_args() + + code, insts, wname = load(args.dispatch_dir, args.wave) + cyc = [r[0] for r in insts] + print(f"wave={wname} inst cycle span {min(cyc)}..{max(cyc)} n={len(insts)}") + if not args.range: + print("pass --range LO,HI (a mid steady slice, ~10 outer iters). e.g. " + f"--range {min(cyc) + (max(cyc)-min(cyc))//4},{min(cyc) + (max(cyc)-min(cyc))//2}") + return + lo, hi = (int(x) for x in args.range.split(",")) + seg = sorted((r for r in insts if lo <= r[0] < hi), key=lambda r: r[0]) + span = hi - lo + + E = args.exec_cyc + mfma = sorted(r[0] for r in seg if op_of(code, r[4]).startswith("v_mfma")) + if not mfma: + print("no MFMA in window") + return + + # next_free model (matrix-unit pipeline): each MFMA occupies ONE E-cycle execute + # slot, but slots pipeline -- consecutive MFMAs can ISSUE < E apart and still both + # be hidden (the unit stays busy). Track next_free = when the matrix unit frees. + # - issue t <= next_free : hidden (co-issued in the shadow); slot advances +E + # - issue t > next_free : the unit was IDLE for (t - next_free) -> EXPOSED + # This fixes the older union-of-[issue,issue+E) model, which capped overlapping + # windows and so mislabeled shadow-hidden loads (dense 8-cyc-apart MFMAs) as + # exposed. Blame each exposed gap on the first non-MFMA op issuing inside it. + next_free = mfma[0] + gaps = [] + for t in mfma: + if t > next_free: + gaps.append((next_free, t)) + next_free = t + E + else: + next_free = next_free + E + exp = sum(b - a for a, b in gaps) + cov = span - exp + print(f"\nsegment [{lo},{hi}) span={span} mfma={len(mfma)} exec={E}") + print(f"MFMA-covered: {cov} ({cov*100//span}%) EXPOSED: {exp} ({exp*100//span}%)") + print(f"cyc/mfma = {span/len(mfma):.2f} (floor = {E})") + + # OCCUPANCY attribution: for every EXPOSED cycle (matrix unit idle), credit it to + # whatever instruction was occupying the issue port then. Each non-MFMA instruction + # occupies [its issue, next instruction's issue) -- i.e. its issue_dur PLUS any + # stall (r[2]); we do NOT separate issue vs stall, we just ask "while the matrix + # unit sat idle, which op was on the port?". Intersect each op's occupancy span + # with the next_free idle gaps and sum. This tiles the whole EXPOSED window with no + # gaps (every idle cycle has an owner), so the ranking directly says: to push + # cyc/mfma toward the floor, which non-MFMA ops must be REMOVED / SHRUNK from + # between the MFMAs. (Stall-vs-issue and who-"blocks" views were dropped -- what + # matters for MFMA-bound speedup is total occupancy stealing issue bandwidth.) + # Row = [cycle, issue_dur, stall, total_dur, code_id]. + def label(cid): + t = code[cid][0].strip() if cid < len(code) else "?" + o = op_of(code, cid) + if o == "s_waitcnt": + v, l = "vmcnt" in t, "lgkmcnt" in t + return "s_waitcnt(vm+lgkm)" if v and l else "s_waitcnt(vmcnt)" if v else \ + "s_waitcnt(lgkmcnt)" if l else "s_waitcnt" + return o + + def gap_overlap(a, b): + tot = 0 + for g0, g1 in gaps: + if g1 <= a or g0 >= b: + continue + tot += min(g1, b) - max(g0, a) + return tot + + ni = [r for r in seg if not op_of(code, r[4]).startswith("v_mfma")] + ni.sort(key=lambda r: r[0]) + occ = collections.Counter() + for i, r in enumerate(ni): + nxt = ni[i + 1][0] if i + 1 < len(ni) else hi + c = gap_overlap(r[0], nxt) + if c: + occ[label(r[4])] += c + total = exp or 1 + print(f"\n== EXPOSED {exp} cyc ({exp*100/span:.1f}% of全局) by occupying instruction ==") + print(f" {'cyc':>6} {'%exp':>5} {'%all':>5} op") + for o, c in occ.most_common(18): + print(f" {c:>6} {c*100//total:>4d}% {c*100/span:>4.1f}% {o}") + + +if __name__ == "__main__": + main() diff --git a/kernels/fp4_gemm_4wave.py b/kernels/fp4_gemm_4wave.py index ef0fd93ef..e016a4613 100644 --- a/kernels/fp4_gemm_4wave.py +++ b/kernels/fp4_gemm_4wave.py @@ -27,20 +27,149 @@ import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl._mlir import ir as _ir from flydsl._mlir.dialects import llvm as _llvm from flydsl.expr import arith, const_expr, range_constexpr from flydsl.expr import buffer_ops as _buffer_ops +from flydsl.expr import rocdl as _rocdl +from flydsl.expr.typing import T as _T from flydsl.expr.typing import Vector as Vec from kernels.fp8_gemm_utils import ( - G2SLoader, ceildiv, compute_global_swizzle, divmod, - make_fp8_buffer_tensor, swizzle_128, wait_barrier, ) +_N_WAVES = 4 # block is always 256 threads -> 4 waves (compile-time constant) + + +class _Buf: + """LDS sub-buffer handle over the ONE contiguous LDS array. + + Holds the shared base pointer plus this buffer's compile-time byte offset + SEPARATELY (not pre-added). Keeping the buffer offset as an int lets the ds_read + path build the address as ``(base + dynamic_swizzle) + const(buffer_off+tile)`` + so the constant outer GEP folds into the ds_read 16-bit offset: field. (Pre-adding + buffer_off into .ptr made the dynamic swizzle the outermost term -> ptrtoint + + int-add -> unfoldable, leaving ~24 address VGPRs materialized.) + + ``.ptr`` (base + buffer_off) is still provided for G2SLoaderAsm, which needs the + actual per-buffer LDS base for m0. + """ + + def __init__(self, base_ptr, byte_off): + self.base_ptr = base_ptr + self.byte_off = byte_off + + @property + def ptr(self): + return fx.add_offset(self.base_ptr, self.byte_off) + + +_gep = _buffer_ops.get_element_ptr + + +def _lds_ptr_t(): + return _ir.Type.parse("!llvm.ptr<3>") + + +def _asm_void(operands, asm_string, constraints): + """Side-effecting void inline asm (LLVM sees no memory op -> no waitcnt added).""" + _llvm.inline_asm(None, operands, asm_string, constraints, has_side_effects=True) + + +def _uniform_i32(value): + """Cast to i32 and force a wave-uniform SGPR value for scalar inline-asm operands.""" + raw = arith._to_raw(value) if not isinstance(value, _ir.Value) else value + if raw.type != _T.i32: + raw = arith._to_raw(fx.Int32(raw)) + return _rocdl.readfirstlane(_T.i32, raw) + + +class G2SLoaderAsm: + """Global->LDS DMA via INLINE-ASM ``buffer_load_dwordx4 ... lds`` instead of the + BufferCopyLDS128b copy atom. + + Mirrors the mla_fwd_decode trick: with the 8 LDS buffers merged into ONE + symbol (so ds_read cross-buffer offsets fold into the 16-bit imm), LLVM can no + longer prove the g2s LDS writes don't alias the ds_reads, and would insert a + spurious ``s_waitcnt vmcnt(0)`` before every ds_read. Emitting the load as + opaque inline asm means LLVM sees no LDS write at all, so it adds no drain -- + vmcnt ordering is managed entirely by our explicit ``wait_barrier`` (which is + itself inline-asm ``s_waitcnt vmcnt(N); s_barrier``). The inline-asm + buffer_load IS still counted toward vmcnt by hardware, so the manual counts in + wait_barrier stay correct. + + Each ``buffer_load_dwordx4 vN, rsrc, soffset offen lds`` writes 16 bytes to the + LDS address in m0; m0 holds the per-step LDS byte base (wave-uniform via + readfirstlane), the per-lane swizzle is the voffset VGPR (loop-invariant), and + the K-step offset is the scalar soffset operand (hardware-free add). + """ + + def __init__(self, rsrc, gl_offsets, n_load_steps, wave_id): + self.rsrc = arith._to_raw(rsrc) + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + @property + def _step_stride(self): + # m0 (LDS byte) advance between consecutive steps of one load. Must be a + # Python int (baked into the asm string), so use the compile-time wave count + # (block is always 256 -> 4 waves) rather than the runtime block_dim value. + return _N_WAVES * 1024 + + def set_wave_base(self, base_ptr): + # Precompute the wave-uniform LDS base (ptrtoint(base) + wave_id*1024) into an + # SGPR ONCE. _lds_base_sgpr then adds the per-buffer compile-time byte_off with + # scalar arithmetic, so the step-0 m0 needs NO per-load readfirstlane (was 8/ + # iter: the m0 base came from wave_id, a VGPR, forcing readfirstlane). + wb = fx.Int32(fx.ptrtoint(base_ptr)) + fx.Int32(self.wave_id * 1024) + self._wave_base_s = _rocdl.readfirstlane(_T.i32, arith._to_raw(wb)) + + def _lds_base_sgpr(self, lds_dst): + # Step-0 LDS byte base (wave-uniform) = precomputed wave base (SGPR) + this + # buffer's compile-time byte_off. Scalar -> no readfirstlane. Later steps add + # _step_stride to m0. + m0 = fx.Int32(self._wave_base_s) + fx.Int32(lds_dst.byte_off) + return arith._to_raw(m0) + + def _voffset(self, step): + # Per-lane global byte offset = swizzle only (loop-invariant). The K-step + # offset is NOT added here -- it goes in the buffer instruction's scalar + # soffset field (hardware-free add), so voffset is a constant VGPR reused + # across all K iterations instead of an `v_add k_offset` every step. + return arith._to_raw(fx.Int32(self.gl_offsets[step])) + + def _emit(self, lds_dst, k_offset, step): + # m0 idiom (gcnasm async_copy): set m0 once for step 0, then advance it with + # s_add for later steps instead of recomputing readfirstlane+s_mov per step. + # The N_TILES steps of one load are issued back-to-back (interleaved into one + # MFMA cluster, in order), so the m0 add-chain stays coherent; s_add m0 is + # volatile inline-asm so the compiler can't reorder across it. Cuts the + # per-step v_readfirstlane (32->8/main-loop) and turns s_mov into s_add, + # freeing scalar-issue slots so the MFMAs pack tighter (toward cyc/mfma 16). + voff = self._voffset(step) + soff = _uniform_i32(k_offset) # scalar soffset (K-step), folded by hardware + stride = self._step_stride + if step == 0: + m0 = self._lds_base_sgpr(lds_dst) + asm = "s_mov_b32 m0, $0\nbuffer_load_dwordx4 $1, $2, $3 offen lds" + _asm_void([m0, voff, self.rsrc, soff], asm, "s,v,s,s") + else: + asm = f"s_add_u32 m0, {stride}, m0\nbuffer_load_dwordx4 $0, $1, $2 offen lds" + _asm_void([voff, self.rsrc, soff], asm, "v,s,s") + + def load(self, lds_dst, k_offset): + for step in range_constexpr(self.n_load_steps): + self._emit(lds_dst, k_offset, step) + + def load_one(self, lds_dst, k_offset, step): + self._emit(lds_dst, k_offset, step) + class S2RLoaderFp4: """fp4 S2R LDS->reg loader. Unlike the fp8 loader it does NOT pack the two @@ -56,15 +185,41 @@ def __init__(self, wave_idx, n_tiles): self.wave_idx = wave_idx self.n_tiles = n_tiles - def _vec_load_16xf8(self, lds_src, offset): - off_tup = fx.make_int_tuple(offset) - ptr_off = fx.add_offset(lds_src.ptr, off_tup) - i8_iter = fx.recast_iter(fx.Uint8, ptr_off) - view = fx.make_view(i8_iter, fx.make_layout(16, 1)) - return view.load() - - def _offset(self, i, step, preshuffled): - row = self.wave_idx * (self.n_tiles * 16) + i * 16 + self.lane_id % 16 + def _vec_load_16xf8(self, lds_src, dyn_offset, const_offset): + # Plain LLVM ds_read (NOT inline-asm). The single-symbol LDS layout would + # normally make the compiler insert an s_waitcnt vmcnt(0) drain before each + # ds_read (it can't prove the g2s global->LDS DMA writes don't alias) -- but + # because g2s is now inline-asm (G2SLoaderAsm), the LDS WRITE is invisible to + # the compiler, so no drain is emitted. And being a REAL LDS load, the + # compiler tracks it with fine-grained lgkmcnt (matching the 8-symbol + # baseline's sync) instead of the conservative per-MFMA VMEM vmcnt it would + # insert for an opaque inline-asm block. + # + # Address folding: the per-lane vaddr VGPR carries only the DYNAMIC part + # (symbol base + 64KB-window half + lane swizzle); the constant (buffer + + # tile) byte offset folds into the ds_read 16-bit offset: immediate via the + # inttoptr + GEP(static) below. ds offset is 16-bit (max 0xFFFF=64KB) but the + # 8 buffers span 0x1d800 (>64KB), so split into two windows: buffers 0-3 + # (base+0) and 4-7 (base+0x10000). imm = buffer_off - window_base + tile*2048 + # stays <= 0xc000+0x1800 = 0xd800 < 0xFFFF. The window base (0 or 0x10000) is + # the only buffer-dependent term left in the dynamic vaddr, so all 8 buffers + # collapse to 2 base VGPRs per (operand, step) -- 4 total in the hot loop. + total_off = lds_src.byte_off + const_offset + window_base = (total_off // 0x10000) * 0x10000 + imm = total_off - window_base + assert 0 <= imm <= 0xFFFF + vaddr = fx.Int32(fx.ptrtoint(lds_src.base_ptr)) + fx.Int32(window_base + dyn_offset) + lds_ptr = _llvm.inttoptr(_lds_ptr_t(), arith._to_raw(vaddr)) + if imm != 0: + lds_ptr = _gep(lds_ptr, static_byte_offset=imm) + vec4_i32 = _ir.VectorType.get([4], fx.Int32.ir_type) + raw = _llvm.LoadOp(vec4_i32, lds_ptr, alignment=16).result + return Vec(raw) + + def _dyn_offset(self, step, preshuffled): + # Lane-dependent part (tile i=0); the per-tile i*2048 is added as a constant. + # Verified: _offset(i,step) - _offset(0,step) == i*2048 on both paths. + row = self.wave_idx * (self.n_tiles * 16) + self.lane_id % 16 col = (self.lane_id // 16) * 16 + step * 64 if const_expr(preshuffled): return (row // 8) * 1024 + (row % 8) * 16 + (col // 16) * 128 @@ -76,14 +231,16 @@ def load(self, lds_src, preshuffled=False): for i in range_constexpr(self.n_tiles): halves = [] for step in range_constexpr(2): - v = self._vec_load_16xf8(lds_src, self._offset(i, step, preshuffled)) + dyn = self._dyn_offset(step, preshuffled) + v = self._vec_load_16xf8(lds_src, dyn, i * 2048) halves.append(v.bitcast(fx.Int32)) # i32x4, the K=128 MFMA operand frag.append(halves) # [ksub0_i32x4, ksub1_i32x4] return frag def load_one(self, lds_src, i, ksub, preshuffled=False): """One i32x4 (tile i, K=128 sub-block ksub) -- the interleave granularity.""" - v = self._vec_load_16xf8(lds_src, self._offset(i, ksub, preshuffled)) + dyn = self._dyn_offset(ksub, preshuffled) + v = self._vec_load_16xf8(lds_src, dyn, i * 2048) return v.bitcast(fx.Int32) @@ -187,7 +344,7 @@ def __init__(self, n_tiles_a, n_tiles_b): def idx(self, i, j): return i * self.n_tiles_b + j - def call(self, a, b, c, sa, sb, interleave=None): + def call(self, a, b, c, sa, sb, interleave=None, interleave_stride=1): """``sa`` / ``sb`` are lists (len n_groups) of packed-E8M0 i32 scales (4 sub-fields each, one full K=256 step for a 32-row pack-group). @@ -208,6 +365,7 @@ def call(self, a, b, c, sa, sb, interleave=None): # ds_read/buffer_load fits free between MFMAs). Mirrors fp8 _interleaved_cluster. thunks = list(interleave) if interleave else [] nth = [0] # python-level counter (compile-time), not loop-carried + mth = [0] # MFMA counter, for spacing thunks every `interleave_stride` MFMAs for ksub in range_constexpr(_FP4_PACK): for i in range_constexpr(self.n_tiles_a): a_op = a[i][ksub] @@ -218,9 +376,13 @@ def call(self, a, b, c, sa, sb, interleave=None): sb_v = sb[j // _FP4_PACK] jb = j % _FP4_PACK c[self.idx(i, j)] = self._mfma_agpr(a_op, b_op, c[self.idx(i, j)], sa_v, sb_v, ksub, ia, jb) - if nth[0] < len(thunks): + # Issue a thunk only every `interleave_stride` MFMAs so a load gets + # >1 MFMA execute-shadow to hide behind (vs bunched after the first + # few MFMAs). Spreads buffer_load/ds_read across the quad. + if nth[0] < len(thunks) and (mth[0] % interleave_stride) == 0: thunks[nth[0]]() nth[0] += 1 + mth[0] += 1 while nth[0] < len(thunks): thunks[nth[0]]() nth[0] += 1 @@ -246,43 +408,108 @@ def _mfma_agpr(self, a_op, b_op, acc, sa_v, sb_v, ksub, ia, jb): ) -class ScaleLoader: - """Loads ``shuffle_scale_w4``-PRESHUFFLED per-1x32 E8M0 scales. - - The shuffled layout (gate_up=False) packs the e8m0 as - ``[N1, K1, K_Lane, N_Lane, K_Pack, N_Pack]`` (K_Lane=4, N_Lane=16, - K_Pack=N_Pack=2), so the 4 e8m0 selected by a lane's opsel values 0..3 are - exactly the 4 contiguous bytes of one i32. Per lane the element offset is:: - - group * (K1 * 64) + kstep * 64 + lane_div_16 * 16 + lane_mod_16 (i32) - - where ``group = base_tile // 32`` is the N1 index (one pack-group = the 2 - tiles n_tiles=2), ``K1 = K // 256``, and the i32 holds [K_Pack, N_Pack]. - The MFMA opsel ``ksub*2 + tile_in_pair`` selects the byte at runtime-free. - One buffer_load per K-step feeds all (tile, ksub) MFMAs of the group. +# Scale-LDS geometry. Each wave needs 4 scale blocks (256 B each) per operand: +# blocks {R0g0, R0g1, R1g0, R1g1} for A, {C0g0, C0g1, C1g0, C1g1} for B. One +# ``buffer_load_dwordx4 ... lds`` (64 lanes x 16 B = 1024 B/wave) gathers all 4 +# blocks of one operand. A region (1024 B) + B region (1024 B) = 2048 B/wave; +# x4 waves = 8192 B per pipeline slot. Slots are indexed kstep%_SCALE_SLOTS. +# With depth-2 prefetch AND an unroll-by-2 main loop, up to 4 K-steps are live at +# once (read[kk], read[kk+1], gather[kk+2], gather[kk+3]), so 4 slots are needed +# to avoid one being overwritten before its read (3 slots -> nan). 4*8192 = 32 KB. +_SCALE_WAVE_BYTES = 2048 +_SCALE_SLOT_BYTES = _N_WAVES * _SCALE_WAVE_BYTES # 8192 +_SCALE_SLOTS = 4 +_SCALE_LDS_BYTES = _SCALE_SLOTS * _SCALE_SLOT_BYTES # 32 KB +_SCALE_A_REGION = 0 +_SCALE_B_REGION = 1024 + + +class ScaleLoaderLDS: + """Loads ``shuffle_scale_w4``-PRESHUFFLED per-1x32 E8M0 scales via a single + ``buffer_load_dwordx4 ... lds`` per operand per K-step, into a scale LDS + region, then per-lane ``ds_read_b32`` -- replacing the 8 per-step + ``buffer_load_dword`` (each with a dur-6 voffset v_add, the #1 exposed + hot-loop cost). + + Layout (gate_up=False): per (N1 group, K-step) the e8m0 form a 64-i32 + (256 B) block ``[K_Lane(4), N_Lane(16)]`` in which lane L's MFMA scale is + element L. A wave's 4 blocks are groups ``{G, G+1, G+4, G+5}`` (G+4 == the + second M/N half's group, since LDS_BLOCK/32 == 4). The dwordx4 gather has + lane g fetch 4 contiguous i32 ``(g%16)*4..+3`` of block ``g//16`` and write + them to ``m0 + g*16``; that lands i32 j of block blk at LDS byte + ``blk*256 + j*4`` (natural order), so the read is ``region + blk*256 + L*4``. """ - def __init__(self, scale_arg, n_tiles, K, lane_id): + def __init__(self, scale_arg, n_tiles, K, lane_id, wave_id, lds_base_ptr, region_off): assert n_tiles % _FP4_PACK == 0 - self.n_tiles = n_tiles - self.n_groups = n_tiles // _FP4_PACK # pack-groups of 2 tiles (32 rows) + self.n_groups = n_tiles // _FP4_PACK # pack-groups per M/N half (=2) self.K1 = K // 256 - self.row_stride = self.K1 * 64 # i32 elems per N1 group - self.lane_off = (lane_id // 16) * 16 + (lane_id % 16) - self.rsrc = _buffer_ops.create_buffer_resource(scale_arg, max_size=True) - - def load_step(self, kstep, base_tile): - """list[n_groups] of packed i32 (K=256 step for each 32-row pack-group). - All addends cast to Int32 so a runtime fx ``kstep`` (scf.for loop var) - doesn't trip arith.addi's same-type requirement (base_tile is Index).""" - base_group = fx.Int32(base_tile // 32) - kterm = fx.Int32(kstep) * fx.Int32(64) - lane = fx.Int32(self.lane_off) - out = [] - for g in range_constexpr(self.n_groups): - i32_off = (base_group + fx.Int32(g)) * fx.Int32(self.row_stride) + kterm + lane - out.append(_buffer_ops.buffer_load(self.rsrc, i32_off, vec_width=1, dtype=fx.Int32)) - return out + self.row_i32 = self.K1 * 64 # i32 per N1 group + self.lane_id = lane_id + self.wave_id = wave_id + self.region_off = region_off + self.rsrc = arith._to_raw(_buffer_ops.create_buffer_resource(scale_arg, max_size=True)) + # Gather per-lane block / within-block index (loop-invariant). + self._blk = lane_id // 16 # 0..3 -> which of the 4 blocks + self._in16 = lane_id % 16 # 0..15 -> which 4-i32 chunk within the block + self._lds_base = fx.Int32(fx.ptrtoint(lds_base_ptr)) + + def _slot_wave_byte(self, slot): + return ( + self._lds_base + + fx.Int32(slot) * fx.Int32(_SCALE_SLOT_BYTES) + + fx.Int32(self.wave_id * _SCALE_WAVE_BYTES + self.region_off) + ) + + def set_wave_base(self): + """Precompute the wave-uniform LDS base (lds_base + wave_id*2048 + region) + into an SGPR ONCE via a single readfirstlane. gather() then computes m0 = + wave_base_s + slot*8192 with scalar arithmetic, so the "s" constraint on m0 + needs NO per-gather readfirstlane (was 4/iter: one per A/B gather per step).""" + wave_base = self._lds_base + fx.Int32(self.wave_id * _SCALE_WAVE_BYTES + self.region_off) + self._wave_base_s = _rocdl.readfirstlane(_T.i32, arith._to_raw(wave_base)) + # soffset=0 as a wave-uniform SGPR (readfirstlane'd once, reused every gather). + self._soff0 = _uniform_i32(fx.Int32(0)) + + def gather(self, kstep, slot, base_tile): + """Issue ONE buffer_load_dwordx4...lds gathering all 4 blocks of this + operand into scale-LDS ``slot``. ``base_tile`` = the M/N-half-0 base row + (half 1 is reached via group +4). Inline-asm so LLVM emits no LDS-write + drain; vmcnt accounting is owned by wait_barrier (hardware counts it).""" + # LDS write (layout A): lane L writes its 16 B to m0 + L*16. So lane L must + # READ from global the data destined for LDS slot L*16 = block (L//16), + # i32 chunk (L%16)*4..+3. block g//16 -> group G + (blk//2)*4 + (blk%2). + G = fx.Int32(base_tile // 32) + grp = G + (self._blk // 2) * fx.Int32(4) + (self._blk % 2) + i32_off = grp * fx.Int32(self.row_i32) + fx.Int32(kstep) * fx.Int32(64) + self._in16 * fx.Int32(4) + voff = arith._to_raw(i32_off * fx.Int32(4)) # bytes + # m0 = precomputed wave base (SGPR) + slot*8192 (scalar): no readfirstlane. + m0 = arith._to_raw(fx.Int32(self._wave_base_s) + fx.Int32(slot) * fx.Int32(_SCALE_SLOT_BYTES)) + asm = "s_mov_b32 m0, $0\nbuffer_load_dwordx4 $1, $2, $3 offen lds" + _asm_void([m0, voff, self.rsrc, self._soff0], asm, "s,v,s,s") + + def read_half(self, slot, half): + """Per-lane ds_read of ONE half (2 blocks) -> list[n_groups] of i32. + Split from read() so half-1 can be issued as a thunk in an MFMA shadow + (half-0 feeds c00/c01, half-1 feeds c10/c11 -- two MFMA clusters later).""" + # lane-contiguous write: lane g's 4 i32 at m0 + g*16. block-elem e of block + # blk was written by lane (blk*16 + e//4) as its (e%4)-th i32 -> LDS byte + # blk*256 + (e//4)*16 + (e%4)*4. MFMA lane L wants block-elem L. + L = self.lane_id + base = self._slot_wave_byte(slot) + fx.Int32((L // 4) * 16 + (L % 4) * 4) + grp_list = [] + for gi in range_constexpr(self.n_groups): + blk = half * 2 + gi + vaddr = base + fx.Int32(blk * 256) + lds_ptr = _llvm.inttoptr(_lds_ptr_t(), arith._to_raw(vaddr)) + raw = _llvm.LoadOp(fx.Int32.ir_type, lds_ptr, alignment=4).result + grp_list.append(fx.Int32(raw)) + return grp_list + + def read(self, slot): + """Per-lane ds_read of the 4 blocks -> (half0, half1), each list[n_groups] + of i32 (the same shape the MFMA consumes: sa[i//2] / sb[j//2]).""" + return self.read_half(slot, 0), self.read_half(slot, 1) class StoreCFp4: @@ -360,32 +587,42 @@ def compile_fp4_gemm_4w( a_lds_size = LDS_BLOCK_M * BLOCK_K_BYTES b_lds_size = LDS_BLOCK_N * BLOCK_K_BYTES + # One contiguous LDS array (single static `make_ptr` -> single `@__shared_alloc` + # symbol), NOT 8 independent leaf fields. With one base symbol every cross-buffer + # offset is a compile-time constant against the same pointer, so the compiler + # folds it into the 16-bit ds_read offset: field instead of materializing ~24 + # per-(buffer,tile) address VGPRs (252 -> ~228 arch VGPR). The catch: a single + # LDS symbol breaks the compiler's alias analysis between the g2s global->LDS DMA + # writes and the ds_reads, so it would insert a spurious `s_waitcnt vmcnt(0)` + # before every ds_read -- which is why g2s uses inline-asm buffer_load (see + # G2SLoaderAsm) so LLVM sees no LDS write and emits no drain. + assert a_lds_size == b_lds_size + _lds_buf = a_lds_size # 16KB; 8 buffers laid out contiguously # 128KB + @fx.struct class SharedStorage: - A_lds_cur_0: fx.Array[fx.Int8, a_lds_size, 16] - A_lds_cur_1: fx.Array[fx.Int8, a_lds_size, 16] - A_lds_next_0: fx.Array[fx.Int8, a_lds_size, 16] - A_lds_next_1: fx.Array[fx.Int8, a_lds_size, 16] - B_lds_cur_0: fx.Array[fx.Int8, b_lds_size, 16] - B_lds_cur_1: fx.Array[fx.Int8, b_lds_size, 16] - B_lds_next_0: fx.Array[fx.Int8, b_lds_size, 16] - B_lds_next_1: fx.Array[fx.Int8, b_lds_size, 16] + all_lds: fx.Array[fx.Int8, 8 * _lds_buf, 16] + scale_lds: fx.Array[fx.Int8, _SCALE_LDS_BYTES, 16] @flyc.kernel def kernel_gemm( A: fx.Tensor, B_T: fx.Tensor, C: fx.Tensor, A_scale: fx.Tensor, B_scale: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 ): - I8_IR_t = fx.Int8.ir_type - lds = fx.SharedAllocator().allocate(SharedStorage).peek() - a_cur0 = lds.A_lds_cur_0 - a_cur1 = lds.A_lds_cur_1 - a_next0 = lds.A_lds_next_0 - a_next1 = lds.A_lds_next_1 - b_cur0 = lds.B_lds_cur_0 - b_cur1 = lds.B_lds_cur_1 - b_next0 = lds.B_lds_next_0 - b_next1 = lds.B_lds_next_1 + # 8 buffers as compile-time offsets off ONE base pointer (single symbol). + _base_ptr = lds.all_lds.ptr + + def _buf(idx): + return _Buf(_base_ptr, idx * _lds_buf) + + a_cur0 = _buf(0) + a_cur1 = _buf(1) + a_next0 = _buf(2) + a_next1 = _buf(3) + b_cur0 = _buf(4) + b_cur1 = _buf(5) + b_next0 = _buf(6) + b_next1 = _buf(7) lane_id = fx.thread_idx.x % 64 wave_id = fx.thread_idx.x // 64 @@ -410,16 +647,18 @@ def kernel_gemm( # K-step (same constant fp8_gemm_4wave uses for b_preshuffled). B_K_STEP = 2 * 1024 - gA = make_fp8_buffer_tensor(A, I8_IR_t) - gB = make_fp8_buffer_tensor(B_T, I8_IR_t) - ga_div = fx.logical_divide(gA, fx.make_layout(1, 1)) - gb_div = fx.logical_divide(gB, fx.make_layout(1, 1)) - mfma = Mfma16x16x128Fp4(N_TILES_A, N_TILES_B) - # One i32 scale per K-step (256 fp4) per M/N-pair; K-step index = k. - a_scale_ld = ScaleLoader(A_scale, N_TILES_A, K, lane_id) - b_scale_ld = ScaleLoader(B_scale, N_TILES_B, K, lane_id) + # Scale via dwordx4...lds gather + ds_read (see ScaleLoaderLDS). One gather + # per operand per K-step into a triple-buffered scale-LDS slot; read[kc] / + # gather[kc+2] use slots kc%3 / (kc+2)%3 (3 distinct -> race-free). + _scale_base_ptr = lds.scale_lds.ptr + a_scale_ld = ScaleLoaderLDS(A_scale, N_TILES_A, K, lane_id, wave_id, _scale_base_ptr, _SCALE_A_REGION) + b_scale_ld = ScaleLoaderLDS(B_scale, N_TILES_B, K, lane_id, wave_id, _scale_base_ptr, _SCALE_B_REGION) + # Precompute each loader's wave-uniform LDS base into an SGPR once, so the + # per-gather m0 needs no readfirstlane (was 4/iter -> exposed ~19% in ATT). + a_scale_ld.set_wave_base() + b_scale_ld.set_wave_base() base_row = tile_i * BLOCK_M + wave_i * (N_TILES_A * 16) base_col = tile_j * BLOCK_N + wave_j * (N_TILES_B * 16) @@ -428,24 +667,24 @@ def kernel_gemm( sb_C0 = base_col sb_C1 = base_col + LDS_BLOCK_N - def _a_sc(k, base): - return a_scale_ld.load_step(k, base) - - def _b_sc(k, base): - return b_scale_ld.load_step(k, base) - - def _load_scales(k): - """All four packed-i32 scales for K-step ``k`` (prefetchable). ``k`` may - be a compile-time int OR a runtime fx index (load_step's kstep*64 offset - arithmetic and buffer_load both accept a runtime value).""" - return ( - _a_sc(k, sa_R0), - _a_sc(k, sa_R1), - _b_sc(k, sb_C0), - _b_sc(k, sb_C1), - ) - - _load_scales_rt = _load_scales # runtime-k alias used inside the scf.for body + def _slot(k): + return fx.Int32(k) % fx.Int32(_SCALE_SLOTS) + + def _gather_scales(k, slot): + """Issue the 2 dwordx4...lds gathers (A, B) for K-step ``k`` into LDS + ``slot``. base_row/base_col are the M/N half-0 bases; half-1 is reached + via group +4 inside gather().""" + a_scale_ld.gather(k, slot, base_row) + b_scale_ld.gather(k, slot, base_col) + + def _gather_scale_thunks(k, slot): + """gather (A,B) as thunks so they co-issue in the MFMA execute shadow + instead of two back-to-back buffer_load after all MFMAs (ATT showed the + end-of-step gather pair exposed, not hidden).""" + return [ + lambda: a_scale_ld.gather(k, slot, base_row), + lambda: b_scale_ld.gather(k, slot, base_col), + ] # Accumulators: 2x2 64x64 quadrants per wave. c00_frag = [mfma.zero_value] * N_ACCUMS @@ -456,15 +695,32 @@ def _load_scales(k): gl_off_a = compute_global_swizzle(lane_id, wave_id, K_BYTES, N_LDS_ROUNDS, preshuffled=False) gl_off_b = compute_global_swizzle(lane_id, wave_id, K_BYTES, N_LDS_ROUNDS, preshuffled=True) - a_g2s = G2SLoader(ga_div, gl_off_a, N_TILES_A, I8_IR_t, wave_id) - b_g2s = G2SLoader(gb_div, gl_off_b, N_TILES_B, I8_IR_t, wave_id) + # Inline-asm g2s (see G2SLoaderAsm): needs the raw buffer resource. Build it + # once from the i8 buffer tensor (max_size OOB check; all addresses in-bounds). + a_rsrc = _buffer_ops.create_buffer_resource(A, max_size=True) # why max size here... + b_rsrc = _buffer_ops.create_buffer_resource(B_T, max_size=True) # why??? + a_g2s = G2SLoaderAsm(a_rsrc, gl_off_a, N_TILES_A, wave_id) + b_g2s = G2SLoaderAsm(b_rsrc, gl_off_b, N_TILES_B, wave_id) + # Precompute the g2s wave-uniform LDS base into SGPR once (all 8 buffers share + # _base_ptr; per-buffer byte_off is compile-time) -> no per-load readfirstlane. + a_g2s.set_wave_base(_base_ptr) + b_g2s.set_wave_base(_base_ptr) a_s2r = S2RLoaderFp4(wave_i, N_TILES_A) b_s2r = S2RLoaderFp4(wave_j, N_TILES_B) store_c = StoreCFp4(C, c_m, c_n, mfma.idx, N_TILES_A, N_TILES_B, mn_aligned=mn_aligned) - # Prologue. - a_g2s.load(a_cur0, A0_gl_offset + 0 * A_K_STEP) - b_g2s.load(b_cur0, B0_gl_offset + 0 * B_K_STEP) + # Prologue. Scale gathers for step 0/1/2 go FIRST (before the 32 g2s) so they + # are the OLDEST outstanding VMEM -- the main loop's wait_barrier(17) then + # drains them naturally (vs issuing them last, where vmcnt(17) can't reach + # them past the 32 newer g2s -> step-0 read got un-landed LDS -> big-K nan). + # DEPTH-3: scale[0] feeds the initial VGPR carry, scale[1]/scale[2] are read + # (into carry) during step 0/1; step kc then gathers scale[kc+3]. + _gather_scales(0, _slot(0)) + _gather_scales(1, _slot(1)) + _gather_scales(2, _slot(2)) + + a_g2s.load(a_cur0, A0_gl_offset + 0 * A_K_STEP) # 4个load. + b_g2s.load(b_cur0, B0_gl_offset + 0 * B_K_STEP) # 4个... b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) a_g2s.load(a_cur1, A1_gl_offset + 0 * A_K_STEP) @@ -481,8 +737,19 @@ def _do_quad(a_frag, b_frag, c_frag, sa_ksub, sb_ksub): wait_barrier((3 * N_TILES_A) + (3 * N_TILES_B)) b0_frag = b_s2r.load(b_cur0, preshuffled=True) - # Step-0 scale (consumed in iter 0); each iter prefetches step k+1. - saR0, saR1, sbC0, sbC1 = _load_scales(0) + # Initial VGPR scale carry = scale[0] (gathered first in prologue, landed by + # the barriers above). The loop carries scale[kc] in VGPR; each step reads + # scale[kc+1] in the MFMA shadow. saR0/saR1/sbC0/sbC1 each list[n_groups] i32. + sc0_saR0, sc0_saR1 = a_scale_ld.read(_slot(0)) + sc0_sbC0, sc0_sbC1 = b_scale_ld.read(_slot(0)) + sc0 = (sc0_saR0, sc0_saR1, sc0_sbC0, sc0_sbC1) + + # Main-loop wait_barrier vmcnt. g2s and scale gathers are inline-asm so the + # compiler uses our literal vmcnt verbatim. With depth-2 end-of-step scale, + # the scale[kc] gather (issued 2 steps ago) is the oldest outstanding VMEM at + # the step top; vmcnt(17) drains exactly it while keeping the g2s window (16) + # in flight, so the ds_read of scale[kc] sees landed LDS. + _MAIN_VMCNT = 17 # ---- Main K-loop as scf.for, unrolled by 2 ------------------------------ # Why scf.for (not range_constexpr full unroll): fully unrolling all 30 main @@ -496,21 +763,40 @@ def _do_quad(a_frag, b_frag, c_frag, sa_ksub, sb_ksub): # ``buf`` arg names below are fixed (cur0/cur1/next0/next1); a single step # mutates which physical buffer is "cur" via the pointer-pair swap, so the # step body is parameterized by the current pointer set passed in. + def _read_scale_thunks(kc_idx, holder): + """4 thunks, each doing one half-read of scale[kc_idx] into holder + (holder = [saR0, saR1, sbC0, sbC1]). Co-issued in an MFMA shadow so the + read of NEXT step's scale is fully hidden (no exposed lgkmcnt at loop top; + scale lives in VGPR carry).""" + s = _slot(kc_idx) + + def _r(dst, ld, half, _s=s): + holder[dst] = ld.read_half(_s, half) + + return [ + lambda: _r(0, a_scale_ld, 0), + lambda: _r(1, a_scale_ld, 1), + lambda: _r(2, b_scale_ld, 0), + lambda: _r(3, b_scale_ld, 1), + ] + def _one_step(kc, a0f, b0f, sc, accs, bufs): # bufs = (a_cur0, a_cur1, a_next0, a_next1, b_cur0, b_cur1, b_next0, b_next1) ac0, ac1, an0, an1, bc0, bc1, bn0, bn1 = bufs + # DEPTH-3 scale + VGPR carry: ``sc`` = scale[kc] ALREADY in VGPR (read in + # the prior step's MFMA shadow), so there is NO scale ds_read / lgkmcnt + # wait at the top of this step. This step (a) reads scale[kc+1] into the + # next carry inside the MFMA shadow, and (b) gathers scale[kc+3] (depth-3, + # so scale[kc+1] -- gathered at step kc-2 -- is landed at the barrier here). saR0, saR1, sbC0, sbC1 = sc c00f, c01f, c10f, c11f = accs - kc_i = fx.Int32(kc) # Int32 for scale load_step (matches lane_off Int32) - saR0_n, saR1_n, sbC0_n, sbC1_n = _load_scales_rt(kc_i + 1) + kc_i = fx.Int32(kc) _b1 = [None] * N_TILES_B _a1 = [None] * N_TILES_A _a0n = [None] * N_TILES_A _b0n = [None] * N_TILES_B - # This step prefetches K-step (kc+2). g2s offsets fully in Int32 - # (A*_gl_offset is Index; kc_i is the Int32 loop var), so arith.addi - # operands match. a*_off = base + (kc+2)*A_K_STEP. + # This step prefetches K-step (kc+2) for g2s. a*_off = base + (kc+2)*STEP. ak = (kc_i + fx.Int32(2)) * fx.Int32(A_K_STEP) bk = (kc_i + fx.Int32(2)) * fx.Int32(B_K_STEP) a0_off = fx.Int32(A0_gl_offset) + ak @@ -518,41 +804,63 @@ def _one_step(kc, a0f, b0f, sc, accs, bufs): b0_off = fx.Int32(B0_gl_offset) + bk b1_off = fx.Int32(B1_gl_offset) + bk - wait_barrier((2 * N_TILES_A) + (2 * N_TILES_B)) - il = _g2s_thunks(a_g2s, ac0, a0_off, N_TILES_A) + _s2r_thunks(b_s2r, bc1, _b1, N_TILES_B, True) - c00f = mfma.call(a0f, b0f, c00f, saR0, sbC0, interleave=il) + # Next-step scale carry: read scale[kc+1] in the MFMA shadow (4 thunks). + _scn = [None, None, None, None] # saR0, saR1, sbC0, sbC1 for kc+1 + _rd_scn = _read_scale_thunks(kc_i + 1, _scn) + # Scale[kc+3] gather (depth-3), co-issued in the MFMA shadow (2 thunks). + # Clamp to K_ITERS-1: the last loop step would gather scale[K_ITERS] (OOB + # -> memory fault); the clamped extra gather re-reads scale[K_ITERS-1] + # into the same slot (idempotent, its result is never consumed). + _gk = _min(kc_i + fx.Int32(3), fx.Int32(K_ITERS - 1)) + _sc_gather = _gather_scale_thunks(_gk, _slot(_gk)) + + wait_barrier(_MAIN_VMCNT) + il = ( + _g2s_thunks(a_g2s, ac0, a0_off, N_TILES_A) + _s2r_thunks(b_s2r, bc1, _b1, N_TILES_B, True) + _rd_scn[:2] + ) + c00f = mfma.call(a0f, b0f, c00f, saR0, sbC0, interleave=il, interleave_stride=2) b1f = _b1 - il = _g2s_thunks(b_g2s, bc0, b0_off, N_TILES_A) + _s2r_thunks(a_s2r, ac1, _a1, N_TILES_A, False) - c01f = mfma.call(a0f, b1f, c01f, saR0, sbC1, interleave=il) + il = ( + _g2s_thunks(b_g2s, bc0, b0_off, N_TILES_A) + + _s2r_thunks(a_s2r, ac1, _a1, N_TILES_A, False) + + _rd_scn[2:] + ) + c01f = mfma.call(a0f, b1f, c01f, saR0, sbC1, interleave=il, interleave_stride=2) a1f = _a1 - wait_barrier((2 * N_TILES_A) + (2 * N_TILES_B)) - il = _g2s_thunks(b_g2s, bc1, b1_off, N_TILES_A) + _s2r_thunks(a_s2r, an0, _a0n, N_TILES_A, False) - c10f = mfma.call(a1f, b0f, c10f, saR1, sbC0, interleave=il) + wait_barrier(_MAIN_VMCNT) + il = ( + _g2s_thunks(b_g2s, bc1, b1_off, N_TILES_A) + + _s2r_thunks(a_s2r, an0, _a0n, N_TILES_A, False) + + _sc_gather[:1] + ) + c10f = mfma.call(a1f, b0f, c10f, saR1, sbC0, interleave=il, interleave_stride=2) a0nf = _a0n - il = _g2s_thunks(a_g2s, ac1, a1_off, N_TILES_A) + _s2r_thunks(b_s2r, bn0, _b0n, N_TILES_B, True) - c11f = mfma.call(a1f, b1f, c11f, saR1, sbC1, interleave=il) + il = ( + _g2s_thunks(a_g2s, ac1, a1_off, N_TILES_A) + + _s2r_thunks(b_s2r, bn0, _b0n, N_TILES_B, True) + + _sc_gather[1:] + ) + c11f = mfma.call(a1f, b1f, c11f, saR1, sbC1, interleave=il, interleave_stride=2) b0nf = _b0n + sc_next = (_scn[0], _scn[1], _scn[2], _scn[3]) new_bufs = (an0, an1, ac0, ac1, bn0, bn1, bc0, bc1) # swap cur<->next - return a0nf, b0nf, (saR0_n, saR1_n, sbC0_n, sbC1_n), (c00f, c01f, c10f, c11f), new_bufs + return a0nf, b0nf, sc_next, (c00f, c01f, c10f, c11f), new_bufs bufs0 = (a_cur0, a_cur1, a_next0, a_next1, b_cur0, b_cur1, b_next0, b_next1) n_a = 2 * N_TILES_A n_b = 2 * N_TILES_B - _R = arith._to_raw - n_ga = N_TILES_A // _FP4_PACK # scale groups per A-scale (=len(saR0)) + n_ga = N_TILES_A // _FP4_PACK # scale groups per A half (=len(saR0)) n_gb = N_TILES_B // _FP4_PACK + n_sc = 2 * n_ga + 2 * n_gb # sc = (saR0,saR1,sbC0,sbC1) flattened + _R = arith._to_raw - def _flat_sc(sc4): - # sc4 = (saR0, saR1, sbC0, sbC1); each is a list of n_g i32 - out = [] - for s in sc4: - for v in s: - out.append(_R(v)) - return out + def _flat_sc(sc): + saR0, saR1, sbC0, sbC1 = sc + return [_R(v) for v in saR0] + [_R(v) for v in saR1] + [_R(v) for v in sbC0] + [_R(v) for v in sbC1] def _unflat_sc(flat): o = 0 @@ -564,13 +872,13 @@ def _unflat_sc(flat): o += n_gb sbC1 = list(flat[o : o + n_gb]) o += n_gb - return saR0, saR1, sbC0, sbC1 + return (saR0, saR1, sbC0, sbC1) - n_sc = 2 * n_ga + 2 * n_gb + # Carry = a0/b0 fragments + VGPR scale carry (scale[kc]) + 4 accumulator groups. init_state = ( _flat_frag(a0_frag) + _flat_frag(b0_frag) - + _flat_sc((saR0, saR1, sbC0, sbC1)) + + _flat_sc(sc0) + [_R(x) for x in c00_frag] + [_R(x) for x in c01_frag] + [_R(x) for x in c10_frag] @@ -616,7 +924,8 @@ def _unflat_sc(flat): off += n_a b0_frag = _unflat_frag(state[off : off + n_b], N_TILES_B) off += n_b - saR0, saR1, sbC0, sbC1 = _unflat_sc(state[off : off + n_sc]) + # VGPR carry at loop exit = scale[K_ITERS-2] (each iter advances sc by 2). + sc = _unflat_sc(state[off : off + n_sc]) off += n_sc c00_frag = list(state[off : off + N_ACCUMS]) off += N_ACCUMS @@ -627,32 +936,49 @@ def _unflat_sc(flat): c11_frag = list(state[off : off + N_ACCUMS]) off += N_ACCUMS - # Tail step K_ITERS - 2 (scale carried from loop's last prefetch). + # Tail step K_ITERS - 2: scale[K_ITERS-2] is the carried VGPR sc (no read). + # Read scale[K_ITERS-1] into the next carry in the c00 MFMA shadow. + saR0, saR1, sbC0, sbC1 = sc + _scn = [None, None, None, None] + _rd_scn = _read_scale_thunks(fx.Int32(K_ITERS - 1), _scn) + _b1 = [None] * N_TILES_B + _a1 = [None] * N_TILES_A wait_barrier((2 * N_TILES_A) + (2 * N_TILES_B)) - b1_frag = b_s2r.load(b_cur1, preshuffled=True) - c00_frag = _do_quad(a0_frag, b0_frag, c00_frag, saR0, sbC0) - a1_frag = a_s2r.load(a_cur1) - c01_frag = _do_quad(a0_frag, b1_frag, c01_frag, saR0, sbC1) + il = ( + _s2r_thunks(b_s2r, b_cur1, _b1, N_TILES_B, True) + + _s2r_thunks(a_s2r, a_cur1, _a1, N_TILES_A, False) + + _rd_scn + ) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag, saR0, sbC0, interleave=il, interleave_stride=2) + b1_frag = _b1 + a1_frag = _a1 + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag, saR0, sbC1) + _a0n = [None] * N_TILES_A + _b0n = [None] * N_TILES_B wait_barrier((1 * N_TILES_A) + (1 * N_TILES_B)) - a0_frag = a_s2r.load(a_next0) - c10_frag = _do_quad(a1_frag, b0_frag, c10_frag, saR1, sbC0) - b0_frag = b_s2r.load(b_next0, preshuffled=True) - c11_frag = _do_quad(a1_frag, b1_frag, c11_frag, saR1, sbC1) + il = _s2r_thunks(a_s2r, a_next0, _a0n, N_TILES_A, False) + _s2r_thunks(b_s2r, b_next0, _b0n, N_TILES_B, True) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, saR1, sbC0, interleave=il, interleave_stride=2) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, saR1, sbC1) + a0_frag = _a0n + b0_frag = _b0n a_cur0, a_next0 = a_next0, a_cur0 a_cur1, a_next1 = a_next1, a_cur1 b_cur0, b_next0 = b_next0, b_cur0 b_cur1, b_next1 = b_next1, b_cur1 - # Tail step K_ITERS - 1. - saR0, saR1, sbC0, sbC1 = _load_scales(K_ITERS - 1) + # Tail step K_ITERS - 1: scale[K_ITERS-1] read into the carry above. + _b1 = [None] * N_TILES_B + _a1 = [None] * N_TILES_A wait_barrier(0) - b1_frag = b_s2r.load(b_cur1, preshuffled=True) - a1_frag = a_s2r.load(a_cur1) - c00_frag = _do_quad(a0_frag, b0_frag, c00_frag, saR0, sbC0) - c01_frag = _do_quad(a0_frag, b1_frag, c01_frag, saR0, sbC1) - c10_frag = _do_quad(a1_frag, b0_frag, c10_frag, saR1, sbC0) - c11_frag = _do_quad(a1_frag, b1_frag, c11_frag, saR1, sbC1) + saR0, saR1, sbC0, sbC1 = (_scn[0], _scn[1], _scn[2], _scn[3]) + il = _s2r_thunks(b_s2r, b_cur1, _b1, N_TILES_B, True) + _s2r_thunks(a_s2r, a_cur1, _a1, N_TILES_A, False) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag, saR0, sbC0, interleave=il, interleave_stride=2) + b1_frag = _b1 + a1_frag = _a1 + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag, saR0, sbC1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, saR1, sbC0) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, saR1, sbC1) store_c.store(c00_frag, sa_R0, sb_C0) store_c.store(c01_frag, sa_R0, sb_C1)