diff --git a/extra/gemm/hk_gemm_frag.py b/extra/gemm/hk_gemm_frag.py index 0497342520..44c99b0454 100644 --- a/extra/gemm/hk_gemm_frag.py +++ b/extra/gemm/hk_gemm_frag.py @@ -66,7 +66,7 @@ NOTE 3: swizzled addresses are written in provably-contiguous "base + vector-off otherwise the devectorizer emits scalar ds_read_u16/ds_write_b16. """ from tinygrad import Tensor, Device, dtypes -from tinygrad.uop.ops import UOp, AxisType, KernelInfo +from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo from tinygrad.dtype import AddrSpace from tinygrad.renderer import Estimates from tinygrad.helpers import getenv, cdiv @@ -147,6 +147,14 @@ def hk_bf16_gemm_kernel(C:UOp, A:UOp, B:UOp, *, arch:str, stages:int=1) -> UOp: # ---- pipelined path (stages=2) ---- NIT = cdiv(OPS_PER_TILE, NUM_THREADS) # copy ops per thread per tile + def setprio(n:int, slot:int) -> UOp: + """__builtin_amdgcn_s_setprio(n), like gemm_bf16.cpp: raise warp priority for the mma phase + so global/LDS traffic of the other waves doesn't starve issue slots.""" + # distinct src slot per call site so identical-priority instructions at different k-tiles + # don't get UOp-hash-deduped into one placement (s_setprio is position-sensitive) + return UOp(Ops.CUSTOMI, dtypes.void, src=(UOp.const(dtypes.weakint, slot), UOp.const(dtypes.weakint, n)), + arg="__builtin_amdgcn_s_setprio({1}); // {0}") + def gload_write_tile(dst:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp: """store one global tile into an LDS slot (loads and stores share the vec range j).""" j = UOp.range(CPV, slot, AxisType.UPCAST) @@ -202,46 +210,86 @@ def hk_bf16_gemm_kernel(C:UOp, A:UOp, B:UOp, *, arch:str, stages:int=1) -> UOp: TILE_ELEMS = BLOCK_M * K_STEP def copy_stage(dst:UOp, slot_off:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp: - """store one global tile into dst + slot_off (flat element offset -- slot_off = parity*TILE_ELEMS).""" - j = UOp.range(CPV, slot, AxisType.UPCAST) + """store one global tile into dst + slot_off (flat element offset -- slot_off = parity*TILE_ELEMS). + + Each thread's 8-element chunk is a single buffer_load_lds direct-to-LDS instruction + (the kittens '... offen lds' fill path), emitted via Ops.CUSTOMI so it bypasses the + devectorizer (a SHRINK store of a SHRINK load gets expanded to scalars before render).""" ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot+1, AxisType.LOOP) chunk = ir*NUM_THREADS + tid r, cc = chunk // OPR, chunk % OPR - return dst[slot_off + st_half_base(r, perm_col(cc*CPV), K_STEP) + j].store(src[base_row + r, kt*K_STEP + cc*CPV + j]).end(ir, j) + if getenv("HK_G2L", 0) == 3: + # direct-to-LDS fill (kittens '... offen lds' path): the hardware writes each lane's + # chunk to the lane-linear LDS address (M0 + lane*size), so the swizzle is moved to + # the GLOBAL side: lane q's 16B chunk fetches the matrix element that st_half_base + # maps to the tile-linear position q. Verified bijective; the fragment-read layout + # (and therefore the read swizzle) is unchanged. + chunk = ir*NUM_THREADS + tid + p_ = chunk * CPV # tile-linear halves position of this lane's chunk + sub = p_ >> 9 # 16x32 subtile id (512 halves) + r16 = (p_ & 511) >> 5 + flip = (r16 >> 3) & 1 + cb = ((p_ & 31) >> 2) ^ (flip << 2) + r_ = (sub >> 1) * 16 + r16 + c_ = cb*4 + (sub & 1) * 32 # global column (8-aligned) + off_g = (base_row + r_) * K + kt*K_STEP + c_ + lds_el = slot_off + ir*NUM_THREADS*CPV # elements; &buf[el*8] = chunk base byte addr + # feed the raw PARAM (unwrapping the scheduler's RESHAPE view, which would otherwise + # live unfused into the program and fail spec: 'movement ops not allowed in programs'). + prm = src + while prm.op is not Ops.PARAM and len(prm.src): prm = prm.src[0] + nbytes = prm.max_numel() * prm.dtype.itemsize + gname = f"data{prm.arg.slot}_{prm.max_numel()}" + return UOp(Ops.CUSTOMI, dtypes.void, src=(prm, dst, lds_el, off_g), + arg=(f"llvm_amdgcn_raw_buffer_load_lds(make_srsrc_((void*){gname}, {nbytes}), " + f"(as3_uint32_ptr)(&({{1}}[({{2}})])), {CPV*2}, ((unsigned)({{3}}))*2U, 0, 0, 0);")).end(ir) + # default: elementwise global->LDS stores + off_l = slot_off + st_half_base(r, perm_col(cc*CPV), K_STEP) + off_g = (base_row + r) * K + kt*K_STEP + cc*CPV + j = UOp.range(CPV, slot, AxisType.UPCAST) + return dst[off_l + j].store(src[base_row + r, kt*K_STEP + cc*CPV + j]).end(ir, j) ZERO = UOp.const(dtypes.weakint, 0) # prologue: tile 0 into slot 0 of both buffers, barrier before first read g0 = UOp.group(copy_stage(A_l, ZERO, A, by*BLOCK_M, ZERO, 100), copy_stage(B_l, ZERO, B, bx*BLOCK_N, ZERO, 110)) bar0 = UOp.barrier(g0) - # Double-buffered pipeline with symbolic-parity slot indexing (FA/gemm_fragment - # conventions): slot ko%2 holds k-tile ko; compute of tile ko overlaps the prefetch - # copy of tile ko+1 into the other slot. No fill iteration, no predication, one static - # store per slot scope, add_war_barrier protects each hand-off. - # NOTE: this requires pm_split_ranges to split the ko LOOP range at the (ko % 2) - # boundary (ko -> 2 inner iterations with static parity per body). Without the split - # the parity stays symbolic, the LDS swizzle decomposition masks it with (x & 8191)<<2 - # style chains, the devectorizer scalarizes all fragment reads/copies - # (DS_READ_U16/DS_WRITE_B16 -> mis-addressed + 2 mfma's lost) and the kernel misfires. - ko = UOp.range(amt, 600, AxisType.LOOP) - pr, pn = ko % 2, (ko+1) % 2 # slot of the tile being computed / being prefetched - kt_next = UOp.minimum(ko+1, amt-1) # clamped tail prefetch (its data is unused) - # the prefetch of the NEXT tile is ordered only against the iteration chain (not the - # compute's fragment stores), so it can be issued BEFORE the wmma's and overlap them: - # its write slot is the one compute read TWO iterations ago (slot (ko-1)%2 == (ko+1)%2), - # and the loop-end raw/war barrier covers the hand-off (write(ko) -> read(ko+1)). - # flat slot indexing (parity * TILE_ELEMS) keeps the swizzled fragment base in - # "chunk base + vector offset" form so the devectorizer keeps ds_read/ds_write_b128; - # A_l[pr][...] (leading-dim select on the (2, tile) view) scalarizes to ds_read_u16. - pa, pb = A_l.after(bar0, ko), B_l.after(bar0, ko) - ga = UOp.group(copy_stage(pa, pn*TILE_ELEMS, A, by*BLOCK_M, kt_next, 130), - copy_stage(pb, pn*TILE_ELEMS, B, bx*BLOCK_N, kt_next, 140)) - last = compute(acc, pa, pb, afters=(ko,), aoff=pr*TILE_ELEMS, boff=pr*TILE_ELEMS) - # one barrier per k-tile hand-off (like stages=1): covers write(ko)->read(ko+1), and - # read(ko)->write(ko+1) is safe because (ko+1)%2 slot was fully read at iteration ko-1 - # which is closed by the ko-1 barrier. The prefetch rides IN FRONT of the wmma's and - # overlaps them; there is no barrier between the copy and the compute in an iteration. - acc = acc.after(UOp.group(last, ga).barrier().end(ko)) + # Double-buffered pipeline: slot ko%2 holds k-tile ko; the prefetch copy of tile ko+1 + # (into the other slot) rides IN FRONT of the wmma's and overlaps them; one barrier per + # k-tile hand-off covers write(ko)->read(ko+1) [and read(ko)->write(ko+1) is closed by + # the ko-1 barrier already]. The parities/offsets are static python constants when + # HK_UNROLL (default on): straight-line like the kittens main loop; the rolled variant + # uses pm_split_ranges to split the ko LOOP range at the (ko % 2) boundary. + if getenv("HK_UNROLL", 1) and amt % (UN := getenv("HK_UNROLL_U", 8)) == 0: + # outer rolled loop of amt//U iterations, U python-unrolled k-tiles inside: nearly the + # kittens straight-line node shape (one barrier per k-tile) at a fraction of the + # full-unroll uop count (full unroll of amt=64 needs ~12 min of schedule time; U=8 + # keeps every tile's prefetch + compute + hand-off barrier but stays seconds). + ko_o = UOp.range(amt // UN, 600, AxisType.LOOP) + pa, pb = A_l.after(bar0, ko_o), B_l.after(bar0, ko_o) + for i in range(UN): + kt = ko_o * UN + i + pr, pn = (i % 2) * TILE_ELEMS, ((i + 1) % 2) * TILE_ELEMS + kt_next = UOp.minimum(kt + 1, amt - 1) + ga0 = UOp.group(copy_stage(pa, UOp.const(dtypes.weakint, pn), A, by*BLOCK_M, kt_next, 300 + 4*i), + copy_stage(pb, UOp.const(dtypes.weakint, pn), B, bx*BLOCK_N, kt_next, 302 + 4*i)) + sp_hi = setprio(1, 300 + 4*i) # kittens: raised prio for the mma phase + last = compute(acc, pa, pb, afters=(ko_o, sp_hi), aoff=UOp.const(dtypes.weakint, pr), boff=UOp.const(dtypes.weakint, pr)) + sp_lo = setprio(0, 301 + 4*i) + handoff = UOp.group(last, sp_lo, ga0).barrier() + acc = acc.after(handoff) + pa, pb = A_l.after(handoff, ga0), B_l.after(handoff, ga0) + acc = acc.after(UOp.group(handoff).end(ko_o)) + else: + ko = UOp.range(amt, 600, AxisType.LOOP) + pr, pn = ko % 2, (ko+1) % 2 # slot of the tile being computed / being prefetched + kt_next = UOp.minimum(ko+1, amt-1) # clamped tail prefetch (its data is unused) + pa, pb = A_l.after(bar0, ko), B_l.after(bar0, ko) + ga = UOp.group(copy_stage(pa, pn*TILE_ELEMS, A, by*BLOCK_M, kt_next, 130), + copy_stage(pb, pn*TILE_ELEMS, B, bx*BLOCK_N, kt_next, 140)) + sp_hi = setprio(1, 150) + last = compute(acc, pa, pb, afters=(ko, sp_hi), aoff=pr*TILE_ELEMS, boff=pr*TILE_ELEMS) + acc = acc.after(UOp.group(last, setprio(0, 151), ga).barrier().end(ko)) # ---- epilogue: per-thread fragment stores, cast to bf16 (scalar per fragment element) ---- mt, nt = UOp.range(MT, 801, AxisType.LOOP), UOp.range(NT, 802, AxisType.LOOP) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 1264a76b97..17ca2d96fd 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -2027,7 +2027,9 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: stores: list[UOp] = [] if is_lds and not is_store: - # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size + # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size. + # HW never takes a per-lane LDS address: kittens' direct fill sets M0 (s_mov_b32 m0, sN) + # before every lds instruction, giving lane-linear chunks with the swizzle on the GLOBAL side. lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF) lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4) for i in range(n_dwords): diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 94a4790162..f3a67b1723 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -106,7 +106,9 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: # TODO: this should handle images too, it's just memory coalescing if u.op in {Ops.LOAD, Ops.STORE}: assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores" - assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}" + # movement-op-wrapped accesses (e.g. a REG placeholder store through a RESHAPE) aren't + # index-addressed; there's nothing to coalesce for them + if u.src[0].op is not Ops.INDEX: continue buf, idx_u = u.src[0].src if buf.addrspace == AddrSpace.REG: continue idx, valid = idx_u.get_idx(), idx_u.get_valid() diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index ce6d6330de..22f4fde3fe 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -236,7 +236,7 @@ class CStyleLanguage(Renderer): assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}" if u.op in {Ops.ENDIF, Ops.END}: depth -= 1 - if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ + if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK} or \ (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \ (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \ (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): @@ -474,6 +474,42 @@ class NVCCRenderer(CUDARenderer): def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar()) def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})" +def _g2l_parts(u:UOp) -> tuple[UOp, UOp, UOp, UOp]|None: + """STORE(local[li]) <- LOAD(global[gi]) (scalar INDEX or vec SHRINK): a global->shared copy expressible + as one buffer_load_lds (direct-to-LDS) instruction on gfx9.4+. Returns (buf, lidx, gbuf, gidx).""" + if u.op is not Ops.STORE or len(u.src) != 2: return None + li, ld = u.src + if li.op is Ops.INDEX and li.addrspace == AddrSpace.LOCAL and len(li.src) == 2: buf, idx = li.src + elif li.op is Ops.SHRINK and li.src[1].dtype is not None and li.src[0].addrspace == AddrSpace.LOCAL: buf, idx = li.src[0], li.src[1] + else: return None + if ld.op is not Ops.LOAD or len(ld.src) != 1: return None + gi = ld.src[0] + if gi.op is Ops.INDEX and gi.addrspace == AddrSpace.GLOBAL and len(gi.src) == 2: gbuf, gidx = gi.src + elif gi.op is Ops.SHRINK and gi.src[0].addrspace == AddrSpace.GLOBAL: gbuf, gidx = gi.src[0], gi.src[1] + else: return None + if li.dtype.scalar() != ld.dtype.scalar(): return None + return buf, idx, gbuf, gidx + +def _g2l_match(u:UOp) -> bool: return _g2l_parts(u) is not None + +def _render_g2l_lds(ctx, u:UOp) -> str|None: + if (parts := _g2l_parts(u)) is None: return None + buf, idx, gbuf, gidx = parts + sz = u.src[0].dtype.itemsize # whole copy size in bytes (16 for an 8xbf16 chunk) + esz = u.src[0].dtype.scalar().itemsize + return (f"llvm_amdgcn_raw_buffer_load_lds(make_srsrc_((void*){ctx[gbuf]}, {gbuf.max_numel()*gbuf.dtype.itemsize}), " + f"(as3_uint32_ptr)(&({ctx[buf]}[({ctx[idx]})])), {sz}, ((unsigned)({ctx[gidx]}))*{esz}U, 0, 0, 0);") + +G2L_LDS_DECLS = [ + "typedef int int32x4_t __attribute__((ext_vector_type(4)));", + "typedef __attribute__((address_space(3))) unsigned* as3_uint32_ptr;", + ("extern __attribute__((device)) void\n" + "llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, as3_uint32_ptr lds_ptr, int size, int voffset, int soffset, int offset, int aux)\n" + ' __asm("llvm.amdgcn.raw.buffer.load.lds");'), + """static inline __attribute__((device)) int32x4_t make_srsrc_(const void* p, unsigned rb) { + int32x4_t r = {(int)(unsigned long)p, (int)(((unsigned long)p)>>32), (int)rb, 0x110000}; + return r;\n}"""] + class HIPRenderer(CStyleLanguage): shared_max = 65536 # NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation @@ -491,6 +527,8 @@ class HIPRenderer(CStyleLanguage): if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast if self.is_cdna(target.arch): self.string_rewrite = PatternMatcher([ + # direct global->LDS copies (buffer_load_lds), skipping the register round-trip + (UPat(Ops.STORE, name="st"), lambda ctx,st: _render_g2l_lds(ctx, st) if getenv("HK_G2L") else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}," f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"), @@ -536,6 +574,9 @@ class HIPRenderer(CStyleLanguage): def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: prefix, ockl = [], [] + g2l_used = any(_g2l_match(u) for u in uops) or \ + any(u.op is Ops.CUSTOMI and isinstance(u.arg, str) and u.arg.startswith("llvm_amdgcn_raw_buffer_load_lds") for u in uops) + if self.is_cdna(self.target.arch) and g2l_used: prefix += G2L_LDS_DECLS type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" } used_dtypes = uops_to_dtypes(uops) if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops): diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 5298129d0a..ed07d8804c 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -91,7 +91,7 @@ spec_shared = PatternMatcher([ isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)), # GROUP of stores (or groups, or NOOPs) - (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END))), lambda: True), + (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END, Ops.CUSTOMI))), lambda: True), # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.INDEX,