[gfx1250] Add 256x256 compute-bound GEMM (WIP)

kernels/gemm/gemm_a8w4_mxscale_gfx1250.py

@@ -52,7 +52,7 @@ def launch_gemm_a8w4_mxscale(
     K_WS = tile_k // WMMA_K
     PACK_TK = tile_k // 2  # B row bytes per K-tile (FP4 packed 2/byte)
     SC_WORDS = tile_k // 4  # scale i32 words per super-row per K-tile
-    SA_SUPERS = tile_m // 32
+    SA_SUPERS = max(1, tile_m // 32)  # a sub-super tile still stages its whole containing super-row
     SB_SUPERS = tile_n // 32
     warp_tile_m = tile_m // m_warp
     warp_tile_n = tile_n // n_warp
@@ -230,6 +230,8 @@ def load_b(buf, wn, ks):
 
         def load_sa(buf, wm, ks):
             row = wmb + wm * 16 + lane16
+            if const_expr(tile_m < 32):
+                row = row + blk_m % 32  # sub-super tile: its rows sit mid-super-row
             word = (row // 32) * SC_WORDS + ks * 32 + (row % 32)
             return lds_load_b32(buf, fx.Int64(SA_OFF + word * 4))[0]
 

kernels/gemm/gemm_a8w8_splitk_reduce_gfx1250.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Split-K partial reduction epilogue for the gfx1250 A8W8 GEMM."""
+
+import functools
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr import gpu, ptrtoint, range_constexpr
+from flydsl.expr.typing import T
+
+DEFAULT_BLOCK = 128
+VEC = 8  # 16 bytes per thread and slice: one BufferCopy128b in/out.
+
+
+def compile_gemm_a8w8_splitk_reduce(
+    *,
+    split_k: int,
+    out_dtype_str: str = "bf16",
+    block: int = DEFAULT_BLOCK,
+    unroll: int = 0,
+):
+    if split_k <= 1:
+        raise ValueError(f"split_k must be greater than one, got {split_k}")
+    if out_dtype_str not in ("bf16", "f16"):
+        raise ValueError(f"unsupported output dtype {out_dtype_str!r}")
+    if block not in (64, 128, 256, 512):
+        raise ValueError(f"block must be one of 64, 128, 256, or 512, got {block}")
+    unroll = unroll or 256 // block
+    if unroll <= 0:
+        raise ValueError(f"unroll must be positive, got {unroll}")
+    return _compile_gemm_a8w8_splitk_reduce(split_k, out_dtype_str, block, unroll)
+
+
+@functools.lru_cache(maxsize=32)
+def _compile_gemm_a8w8_splitk_reduce(split_k: int, out_dtype_str: str, block: int, unroll: int):
+    is_f16 = out_dtype_str == "f16"
+    tile = block * VEC
+    span = tile * unroll
+
+    @flyc.kernel(known_block_size=[block, 1, 1])
+    def reduce_kernel(partials: fx.Pointer, out: fx.Pointer, i32_total: fx.Int32):
+        elem = fx.Float16 if is_f16 else fx.BFloat16
+        vec_f32, vec_out = T.vec(VEC, T.f32), T.vec(VEC, elem.ir_type)
+        tid, blk = gpu.thread_id("x"), gpu.block_id("x")
+        atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem)
+        slice_bytes = fx.Int64(i32_total) * fx.Int64(2)
+
+        def _view(ptr_i64):
+            ptr_ty = fx.PointerType.get(elem.ir_type, address_space=fx.AddressSpace.Global, alignment=2)
+            view = fx.make_view(
+                fx.inttoptr(ptr_ty, ptr_i64),
+                fx.make_layout((1, i32_total), (i32_total, 1)),
+            )
+            return fx.rocdl.make_buffer_tensor(view, num_records_bytes=slice_bytes)
+
+        partial_base = fx.Int64(ptrtoint(partials))
+        partial_bufs = [_view(partial_base + fx.Int64(s) * slice_bytes) for s in range_constexpr(split_k)]
+        out_buf = _view(fx.Int64(ptrtoint(out)))
+
+        tile_mn, tv_layout = fx.make_layout_tv(
+            fx.make_layout((1, block), (1, 1)),
+            fx.make_layout((1, VEC), (1, 1)),
+        )
+        thread_copy = fx.make_tiled_copy(atom, tv_layout, tile_mn).get_slice(tid)
+        base = fx.Int32(blk) * fx.Int32(unroll)
+
+        def _part(buf, u):
+            return fx.slice(fx.zipped_divide(buf, tile_mn), (None, (0, base + u)))
+
+        srcs = [[thread_copy.partition_S(_part(buf, u)) for buf in partial_bufs] for u in range_constexpr(unroll)]
+        frags = [[fx.make_fragment_like(src) for src in row] for row in srcs]
+
+        # Maximize load-to-use distance: issue every slice load before arithmetic.
+        for u in range_constexpr(unroll):
+            for s in range_constexpr(split_k):
+                fx.copy(atom, srcs[u][s], frags[u][s])
+
+        for u in range_constexpr(unroll):
+            acc = fx.Vector(fx.memref_load_vec(frags[u][0])).extf(vec_f32)
+            for s in range_constexpr(1, split_k):
+                acc = acc + fx.Vector(fx.memref_load_vec(frags[u][s])).extf(vec_f32)
+            dst = thread_copy.partition_D(_part(out_buf, u))
+            out_frag = fx.make_fragment_like(dst)
+            fx.memref_store_vec(acc.truncf(vec_out), out_frag)
+            fx.copy(atom, out_frag, dst)
+
+    @flyc.jit
+    def launch(partials: fx.Pointer, out: fx.Pointer, i32_total: fx.Int32, stream: fx.Stream):
+        n_tiles = (i32_total + fx.Int32(span - 1)) // fx.Int32(span)
+        reduce_kernel(partials, out, i32_total).launch(
+            grid=(n_tiles, 1, 1),
+            block=(block, 1, 1),
+            stream=stream,
+        )
+
+    return launch
+
+
+__all__ = ["compile_gemm_a8w8_splitk_reduce"]

python/flydsl/expr/rocdl/tdm_ops.py

@@ -57,6 +57,7 @@
     "tensor_store_gather",
     "tensor_store_2d",
     "tensor_wait",
+    "update_tensor_descriptor_2d_lds_addr",
     "update_tensor_descriptor_2d_addr_lo",
     "update_tensor_gather_descriptor_addr_lo",
     "update_tensor_descriptor_2d_addr_lo_hi",
@@ -202,6 +203,24 @@ def _i32_const(v: int) -> ir.Value:
     return _unwrap(std_arith.ConstantOp(i32, ir.IntegerAttr.get(i32, v)).result)
 
 
+def _extract_lds_base_index(value) -> ir.Value:
+    """Extract a builtin memref or Fly shared view as an LDS byte index."""
+    raw = _unwrap(as_ir_value(value))
+    try:
+        ir.MemRefType(raw.type)
+        return _unwrap(memref_dialect.extract_aligned_pointer_as_index(raw))
+    except ValueError:
+        pass
+
+    from ..._mlir.dialects import fly as _fly_d
+
+    ptr_type = ir.Type.parse("!llvm.ptr<3>")
+    ptr = _fly_d.extract_aligned_pointer_as_index(ptr_type, raw)
+    i64 = ir.IntegerType.get_signless(64)
+    ptr_i64 = llvm_dialect.ptrtoint(i64, ptr)
+    return _unwrap(std_arith.IndexCastOp(ir.IndexType.get(), ptr_i64).result)
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -356,7 +375,7 @@ def make_tensor_descriptor_2d(
     glb_addr_i64 = glb_base_i64 + glb_byte_off_i64
 
     # -- LDS address (byte address within shared memory) --
-    lds_base_idx = _ArithValue(memref_dialect.extract_aligned_pointer_as_index(lds_memref))
+    lds_base_idx = _ArithValue(_extract_lds_base_index(lds_memref))
     # Compute padded LDS stride (elements) for the outer dim
     if pad_interval > 0 and pad_amount > 0:
         lds_inner_stride = inner_tile + pad_amount  # padded row width
@@ -835,6 +854,25 @@ def tensor_store_gather(
 # ---------------------------------------------------------------------------
 
 
+@dsl_loc_tracing
+def update_tensor_descriptor_2d_lds_addr(
+    desc: TDMDescriptor2D,
+    new_lds_addr,
+) -> TDMDescriptor2D:
+    """Return a 2-D descriptor with its LDS address replaced."""
+    from ..._mlir.dialects import vector as _vector_dialect
+
+    return TDMDescriptor2D(
+        dgroup0=_vector_dialect.InsertOp(
+            _raw(new_lds_addr),
+            _raw(desc.dgroup0),
+            static_position=[1],
+            dynamic_position=[],
+        ).result,
+        dgroup1=desc.dgroup1,
+    )
+
+
 def _replace_dgroup0_addr_lo(dgroup0, new_addr_lo):
     """Return a new vector<4xi32> with lane 2 replaced by ``new_addr_lo``."""
     from ..._mlir.dialects import vector as _vector_dialect
@@ -941,11 +979,11 @@ def update_tensor_gather_descriptor_addr_lo(
 # Lane 3 layout (matching ``make_tensor_descriptor_2d`` /
 # ``make_tensor_gather_dgroup0``):
 #     [31:30]  type field (always set to 2 = ``0b10`` via ``| (1 << 31)``)
-#     [29:0]   addr_hi[29:0] (high 32 bits of the original 64-bit address;
-#              only [15:0] are meaningful for 48-bit AMDGPU virtual addresses)
+#     [29:25]  reserved (zero for a valid gfx1250 D#)
+#     [24:0]   global_addr[56:32] (gfx1250 uses a 57-bit GPU virtual address)
 # ---------------------------------------------------------------------------
 
-# Mask covering the address bits in lane 3 (everything except the type field).
+# Mask covering the non-type payload in lane 3, including reserved zero bits.
 _TDM_ADDR_HI_MASK = 0x3FFFFFFF  # bits [29:0]
 # Mask covering the type-field bits at the top of lane 3.
 _TDM_ADDR_HI_FLAG_MASK = 0xC0000000  # bits [31:30]