direct lds

extra/gemm/hk_gemm_frag.py

@@ -0,0 +1,395 @@
+"""
+HipKittens hk_bf16_gemm (extra/thunder/amd/gemm_bf16.cpp) reimplemented with tinygrad UOps.
+
+C[M, N] (bf16) = A[M, K] @ B[N, K]^T, fp32 accumulation, exactly the kittens kernel shape:
+  - 256x256 output tile per workgroup, K_STEP=64
+  - 8 warps in a 2x4 grid, each warp owns a 128x64 warp-tile
+  - v_mfma_f32_16x16x32_bf16 on CDNA4 (gfx950) / v_wmma_f32_16x16x16_bf16 (wave32, gfx12) on RDNA4,
+    fp32 accumulators
+  - shared tiles As/Bs with the kittens st_16x32_s swizzle (16x32 subtiles of 1024B)
+  - K stages (STAGES=1: synchronous single buffer)
+
+Validated on gfx1201 hardware (exact for identity-B, rounding-level noise otherwise), rendered
+and compiled to gfx950 with comgr for assembly comparison against gemm_bf16.cpp.
+
+What is NOT expressible vs the kittens C++:
+  - explicit s_waitcnt vmcnt()/lgkmcnt() pipelining and s_setprio: tinygrad models async copy
+    overlap with slot dependencies and emits full workgroup barriers; instruction scheduling
+    is left to clang/LLVM
+  - direct-to-LDS global loads (buffer_load_lds): tinygrad goes global->reg->LDS
+
+Pipelining status: STAGES=2 gives the kittens-shaped double-buffered pipeline (2 x 64KB LDS
+like gemm_bf16.cpp, copies overlap the previous pair's mma's), written with FA/gemm_fragment
+conventions: LDS buffers are (2, tile) placeholders indexed by symbolic parity (ko % 2),
+which sidesteps static slot choice, fill iterations, predication and duplicate static stores.
+Validated on the CDNA4 emulator for all tile counts (amt = K//64 in {1..32}, odd/even),
+single- AND multi-workgroup (bit-close to stages=1 / to hippkittens at rounding level).
+
+Bug hunt notes (all fixed on this branch; they were entangled for a long time):
+  1. The double-buffered pipeline REGISTER-SPILLS (255+ VGPRs vs 166 for stages=1), and the
+     mock emulator aliased the spill (scratch) segment of ALL waves of a workgroup onto one
+     64-lane region. On real HW each wavefront owns a per-lane segment of the scratch ring
+     (indexed by (wave_id, lane)); waves trampled each other's spilled accumulators, giving
+     the "only the last wave's output survives" signature. emu.py now allocates per-wave
+     scratch buffers.
+  2. The remaining "shape-dependent" corruption (NaNs, mispositioned values in contiguous
+     copies feeding the GEMM) came from tinygrad's devectorizer fusing adjacent bf16 stores
+     into 32-bit stores with UNALIGNED (2-byte) granularity: legal on AMD FLAT/GLOBAL (the
+     hardware splits them), but the emulator floored misaligned addresses to the word below.
+     _mem_store now handles unaligned 32-bit (and wider) accesses byte-exactly.
+  3. memory_coalescing (late/coalesce.py) assumed a single static store per (buffer, index)
+     ("attempting multiple stores"); aliased stores (a double-buffered LDS slot written in a
+     prologue AND a loop body) are now simply kept scalar instead of asserting/merging.
+  4. pm_split_ranges may only split ranges WITHOUT hardware meaning (WEAK/REDUCE/LOOP);
+     splitting LOCAL/WARP/THREAD/GLOBAL/GROUP_REDUCE/UPCAST ranges scrambles the
+     logical<->hardware mapping of hand-written kernels such as this one.
+
+  RDNA4 (gfx12) uses 8-element accumulator fragments, so the 8x4 tile grid needs
+  256 fp32 acc registers per thread -> guaranteed spills (0.85 TF vs 96 TF default on
+  gfx1201). The kernel is right-sized for CDNA4 (fragsz 4 -> 128 acc regs).
+
+Lane layouts (RDNA4 verified with probing on gfx1201 hardware; CDNA from the mfma docs):
+  CDNA (64 thr/warp, 16x16x32):  A/B frag: tile-row = l%16, k = (l//16)*8+i (i in 0..7)
+  RDNA4 (32 thr/warp, 16x16x16): A/B frag: tile-row = l%16, k = (l//16)*4+(i%4)+8*(i//4) (i in 0..7)
+  both:                        acc frag: CDNA m=(l//16)*4+i (i<4) / RDNA4 m=(l//16)*8+i (i<8), n=l%16
+
+The RDNA4 fragment k-set {k0..3, k0+8..11} is not contiguous, so on RDNA4 the LDS column layout
+is block-permuted (4-element blocks within each 16-col group are stored as [0,2,1,3]) making
+every fragment 8 contiguous halves (one 16B chunk) on both archs; the copy path applies the
+same permutation.
+
+NOTE: thread ids come from UOp.special (like mi350x_uop_matmul.py), not an AxisType.LOCAL
+RANGE. (pm_split_ranges now only splits WEAK/REDUCE/LOOP ranges, so LOCAL ranges would
+survive too, but UOp.special is the sanctioned way to tag hardware lane ids.)
+NOTE 2: WMMA operand/accumulator fragments must carry the fragment length in their UOp shape.
+NOTE 3: swizzled addresses are written in provably-contiguous "base + vector-offset" form,
+otherwise the devectorizer emits scalar ds_read_u16/ds_write_b16.
+"""
+from tinygrad import Tensor, Device, dtypes
+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
+
+# ---- tile shape (identical to gemm_bf16.cpp; HK_TILE=128 overrides for small-LDS devices) ----
+BLOCK_M = BLOCK_N = getenv("HK_TILE", 256)
+K_STEP = 64
+WARPS_M, WARPS_N = 2, 4
+NUM_WARPS = WARPS_M * WARPS_N                                # 8
+WARP_TILE_M, WARP_TILE_N = BLOCK_M // WARPS_M, BLOCK_N // WARPS_N   # 128 x 64 (64 x 32 at HK_TILE=128)
+
+def arch_params(arch:str):
+  is_cdna = arch.startswith("gfx9")
+  if is_cdna:   # CDNA mfma 16x16x32 bf16: in-frag 8, acc 4 (16,16,16 on the acc side)
+    return dict(warp_threads=64, dims=(16,16,32), frag_in=8, frag_out=4,
+                acc_m=lambda l, i: (l//16)*4 + i, kperm=None, copy_vec=8)
+  # RDNA4 wmma 16x16x16 (wave32, gfx12 layout): in-frag 8 (permuted in LDS), acc 8
+  return dict(warp_threads=32, dims=(16,16,16), frag_in=8, frag_out=8,
+              acc_m=lambda l, i: (l//16)*8 + i, kperm=(0,2,1,3), copy_vec=4)
+
+# ---- kittens st_16x32_s swizzle (byte offset: off ^ (((off % 1024) >> 9) << 5)) ----
+# halves-index form: within a 1024B (16x32) subtile, halves-index bit4 ^= row bit3,
+# written in "base + vector-offset" form so the devectorizer can prove contiguity.
+def st_half_base(r, c, tile_cols:int):
+  """swizzled halves-index pre-vector-offset; c is the logical (permuted) column, 4-aligned."""
+  subtile_id = (r//16) * (tile_cols//32) + (c//32)
+  r16 = r % 16
+  flip = (r16 >> 3) & 1
+  return subtile_id*512 + r16*32 + (((c % 32) >> 2) ^ (flip << 2)) * 4
+
+def hk_bf16_gemm_kernel(C:UOp, A:UOp, B:UOp, *, arch:str, stages:int=1) -> UOp:
+  """C = A @ B^T ; A is (M,K), B is (N,K), C is (M,N). HipKittens tile shape."""
+  M, K = A.shape
+  N, K2 = B.shape
+  assert K == K2 and A.dtype == B.dtype == dtypes.bfloat16 and C.dtype == dtypes.bfloat16
+  assert not (M % BLOCK_M or N % BLOCK_N or K % K_STEP), f"dims must be multiples of {(BLOCK_M, BLOCK_N, K_STEP)}"
+
+  ap = arch_params(arch)
+  warp_threads, dims = ap["warp_threads"], ap["dims"]
+  FRAG_IN, FRAG_OUT, kperm, acc_m, CPV = ap["frag_in"], ap["frag_out"], ap["kperm"], ap["acc_m"], ap["copy_vec"]
+  NUM_THREADS = NUM_WARPS * warp_threads
+  TC_M, TC_N, TC_K = dims
+  MT, NT = WARP_TILE_M // TC_M, WARP_TILE_N // TC_N               # 8, 4 tiles per warp
+  # permute 4-half blocks within each 16-col group (RDNA4: [0,2,1,3] = swap middle blocks)
+  def perm_col(c):
+    if kperm is None: return c
+    return (c & ~15) | ((((c>>2) & 1) << 1 | ((c>>3) & 1)) << 2) | (c & 3)
+
+  bx, by = UOp.special(N//BLOCK_N, "gidx0"), UOp.special(M//BLOCK_M, "gidx1")
+  lane = UOp.special(warp_threads, "lidx0")
+  warp = UOp.special(NUM_WARPS, "lidx1")
+  warp_row, warp_col = warp // WARPS_N, warp % WARPS_N
+  tid = warp*warp_threads + lane
+
+  def smem(slot) -> UOp: return UOp.placeholder((BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+  As = [smem(2*i) for i in range(stages)]
+  Bs = [smem(2*i+1) for i in range(stages)]
+
+  # per-warp accumulator: (MT x NT) 16x16 tiles of FRAG_OUT fp32 per thread
+  acc = UOp.placeholder((MT, NT, FRAG_OUT), dtypes.float32, 12, AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))   # FA-style init: self-store, keeps the value flow loop-carried
+
+  # global -> LDS copy: CPV halves per op (16B on CDNA, 8B on RDNA4), thread-major coalescing
+  OPS_PER_TILE = BLOCK_M*K_STEP//CPV
+  OPR = K_STEP//CPV
+  def copy_tile(dst:UOp, src:UOp, base_row:UOp, base_col:UOp, slot:int) -> UOp:
+    ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot, AxisType.LOOP)
+    j = UOp.range(CPV, slot+1, AxisType.UPCAST)
+    chunk = ir*NUM_THREADS + tid
+    r, cb = chunk // OPR, chunk % OPR          # row, 4/8-col block
+    return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, base_col + cb*CPV + j]).end(ir, j)
+
+  def load_stage(sidx:int, ko, slot:int, barrier:bool) -> tuple[UOp, UOp]:
+    A_r = copy_tile(As[sidx], A, by*BLOCK_M, ko*K_STEP, slot)
+    B_r = copy_tile(Bs[sidx], B, bx*BLOCK_N, ko*K_STEP, slot+10)
+    bar = UOp.barrier(A_r, B_r) if barrier else UOp.group(A_r, B_r)
+    return As[sidx].after(bar), Bs[sidx].after(bar)
+
+  # ---- 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)
+    def one(ir:int) -> UOp:
+      chunk = ir*NUM_THREADS + tid
+      r, cb = chunk // OPR, chunk % OPR
+      return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, kt*K_STEP + cb*CPV + j])
+    return UOp.group(*[one(ir) for ir in range(NIT)]).end(j)
+
+  def compute(acc:UOp, A_l:UOp, B_l:UOp, afters:tuple[UOp, ...], pred:UOp|None=None, aoff:UOp=None, boff:UOp=None) -> UOp:
+    """One K_STEP=64 iteration: (K_STEP//TC_K) k-chunks unrolled, (MT x NT) mma each, like the kittens main loop.
+
+    pred (optional): a loop-range condition; accumulator stores are predicated on it so the
+    first (fill) iteration of a software pipeline can run the body with garbage LDS contents
+    without contaminating the accumulator."""
+    arow = warp_row*WARP_TILE_M + lane % 16    # fragment tile row in the LDS tile (m)
+    brow = warp_col*WARP_TILE_N + lane % 16    # (n)
+    ja = UOp.range(FRAG_IN, 701, AxisType.UPCAST)
+    jb = UOp.range(FRAG_IN, 702, AxisType.UPCAST)
+    acc_k = acc.after(*afters) if afters else acc
+    last_store = None
+    # in the permuted layout every fragment is 8 contiguous halves starting at an 8-aligned col
+    for kk in range(K_STEP//TC_K):
+      cc = kk*(TC_K//FRAG_IN) + (lane // 16)   # fragment chunk col (8 halves)
+      oa, ob = (aoff, boff) if aoff is not None else (None, None)
+      a_frags = [A_l[st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) if oa is None else
+                 A_l[oa + st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) for mt in range(MT)]
+      b_frags = [B_l[st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) if ob is None else
+                 B_l[ob + st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) for nt in range(NT)]
+      for mt in range(MT):
+        for nt in range(NT):
+          cur = acc_k[mt, nt]
+          out = UOp.wmma(a_frags[mt], b_frags[nt], cur, dims, 'AMD', warp_threads)
+          if pred is not None: out = pred.where(cur, out)
+          last_store = acc_k[mt, nt].store(out)
+          acc_k = acc_k.after(last_store)
+    return last_store
+
+  # ---- K loop ----
+  amt = cdiv(K, K_STEP)
+  _stages = stages
+  if _stages == 1:
+    ko = UOp.range(amt, 600, AxisType.LOOP)
+    A_l, B_l = load_stage(0, ko, 100, barrier=True)
+    last = compute(acc, A_l, B_l, afters=(ko,))
+    acc = acc.after(last.barrier().end(ko))
+  else:
+    # Double-buffered pipeline on FA/gemm_fragment conventions: each LDS buffer is a
+    # (2, tile) placeholder indexed by symbolic parity (ko % 2) -- no static slot choice,
+    # no duplicate static stores (memory_coalescing-safe), no fill iteration, no predication.
+    def smem2(slot) -> UOp: return UOp.placeholder((2*BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+    A_l, B_l = smem2(0), smem2(1)
+
+    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).
+
+      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
+      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: 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)
+  def store_i(i:int) -> UOp:
+    crow = by*BLOCK_M + warp_row*WARP_TILE_M + mt*16 + acc_m(lane, i)
+    ccol = bx*BLOCK_N + warp_col*WARP_TILE_N + nt*16 + lane % 16
+    return C[crow, ccol].store(acc[mt, nt, i].cast(dtypes.bfloat16))
+  out_st = UOp.group(*[store_i(i) for i in range(FRAG_OUT)])
+  return out_st.end(mt, nt).sink(arg=KernelInfo(name="hk_bf16_gemm",
+    estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*2)))
+
+def hk_bf16_gemm_tiny(a:Tensor, b:Tensor, stages:int=1) -> Tensor:
+  """C = a @ b.T for bf16 a (M,K), b (N,K) with the HipKittens-shaped tinygrad kernel."""
+  arch = Device[a.device].renderer.target.arch
+  c = Tensor.empty(a.shape[0], b.shape[0], dtype=dtypes.bfloat16, device=a.device)
+  return c.custom_kernel(a, b, fxn=lambda C, A, B: hk_bf16_gemm_kernel(C, A, B, arch=arch, stages=stages))[0]
+
+if __name__ == "__main__":
+  import numpy as np
+  from tinygrad import Device
+  M = N = K = 512
+  # exact test: B = identity -> C must equal A bit-exactly
+  a = Tensor.randn(M, K, dtype=dtypes.bfloat16).contiguous()
+  bid = Tensor(np.eye(K, N, dtype=np.float32), dtype=dtypes.bfloat16).contiguous()
+  cid = hk_bf16_gemm_tiny(a, bid, stages=getenv("STAGES", 1)).realize()
+  assert np.array_equal(cid.float().numpy(), a.float().numpy()), "identity test failed"
+  # real test: bf16 gemm vs fp32 reference, rounding-level noise
+  b = Tensor.randn(N, K, dtype=dtypes.bfloat16).contiguous()
+  c = hk_bf16_gemm_tiny(a, b, stages=getenv("STAGES", 1)).realize()
+  ref = (a @ b.T).float().realize()
+  err = (c.float() - ref).abs().max().item()
+  print(f"identity exact, random max err: {err:.5f}")
+
+  # ---- benchmark mode: kittens hk_bf16_gemm vs tinygrad stages={1,2} vs the default scheduled gemm ----
+  # run on real hardware with:  DEV=AMD:HIP:gfx950 DEBUG=2 HK_BENCH=1 python extra/gemm/hk_gemm_frag.py
+  # sizes via HK_SIZES="2048x2048x2048,4096x4096x4096" (default 2048 cubed), iteration count via ITERS=20.
+  # timings come from GlobalCounters.time_sum_s (sum of kernel times; same source as the DEBUG=2 'tm' column).
+  if getenv("HK_BENCH"):
+    from tinygrad.helpers import GlobalCounters
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    from tinygrad import Device
+    dev, iters, warm = Device.DEFAULT, getenv("ITERS", 20), 3
+    arch = Device[dev].renderer.target.arch
+    assert arch.startswith("gfx9"), "CDNA only"
+    def bench(label:str, fn, M:int, N:int, K:int) -> float:
+      try:
+        for _ in range(warm): fn()
+        Device[dev].synchronize()
+        GlobalCounters.reset()
+        import time
+        t0 = time.perf_counter()
+        for _ in range(iters): fn()
+        Device[dev].synchronize()
+        wall = time.perf_counter() - t0
+      except Exception as e:
+        print(f"    {label:32s}  unsupported/failed: {type(e).__name__}: {e}")
+        return float('nan')
+      # prefer kernel-side time (GlobalCounters matches the DEBUG=2 'tm' column); fall back to wall clock
+      ms = (GlobalCounters.time_sum_s if GlobalCounters.time_sum_s > 0 else wall) * 1e3 / iters
+      tf = 2*M*N*K / (ms * 1e-3) / 1e12
+      print(f"    {label:32s} {ms:9.3f} ms  {tf:8.1f} TFLOPS")
+      return tf
+    for (M, N, K) in [tuple(map(int, s.split("x"))) for s in getenv("HK_SIZES", "2048x2048x2048").split(",")]:
+      print(f"  size ({M},{N},{K}), grid {M//BLOCK_M}x{N//BLOCK_N} WGs, amt={K//K_STEP} k-tiles/WG")
+      np.random.seed(0)
+      An, Bn = np.random.randn(M, K), np.random.randn(K, N)
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous().realize()          # (M,K)
+      Bk = Tensor(Bn, dtype=dtypes.bfloat16).contiguous().realize()          # (K,N) for kittens
+      Bt = Bk.T.contiguous().realize()                                       # (N,K) for ours
+      tf_kc = bench("kittens hk_bf16_gemm (asm_gemm)", lambda: asm_gemm(A, Bk).realize(), M, N, K)
+      tf_s1 = bench("tiny stages=1", lambda: hk_bf16_gemm_tiny(A, Bt, stages=1).realize(), M, N, K)
+      tf_s2 = bench("tiny stages=2", lambda: hk_bf16_gemm_tiny(A, Bt, stages=2).realize(), M, N, K)
+      tf_df = bench("tinygrad default (a @ Bt.T)", lambda: (A @ Bt.T).realize(), M, N, K)
+      err = (hk_bf16_gemm_tiny(A, Bt, stages=2).float() - asm_gemm(A, Bk).float()).abs().max().item()
+      print(f"    correctness tiny-s2 vs kittens max diff: {err:.5f}")
+      for nm, tf in [("s1", tf_s1), ("s2", tf_s2), ("default", tf_df)]:
+        if tf == tf and tf_kc == tf_kc: print(f"    tiny {nm:8s}/kittens: {tf/tf_kc:6.2%}")
+  # match the real HipKittens hk_bf16_gemm on the (mock) CDNA4 emulator at small sizes.
+  # run from the repo root with:  DEV=MOCK+AMD:HIP:gfx950 HK_COMPARE=1 python extra/gemm/hk_gemm_frag.py
+  if getenv("HK_COMPARE"):
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    assert Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"), "needs CDNA4 (mock emulator or hardware)"
+    def compare(M:int, N:int, K:int, seed:int=0, identity:bool=False):
+      np.random.seed(seed)
+      An = np.random.randn(M, K)
+      Bn = np.eye(K, N) if identity else np.random.randn(K, N)      # (K,N) as expected by asm_gemm
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous()
+      B = Tensor(Bn, dtype=dtypes.bfloat16).contiguous()             # (K,N) for asm_gemm
+      c_hkc = asm_gemm(A, B).realize().float().numpy()               # real HipKittens hk_bf16_gemm
+      c_hkt = hk_bf16_gemm_tiny(A, B.T.contiguous(), stages=getenv("STAGES", 2)).realize().float().numpy()
+      ref64 = A.float().numpy().astype(np.float64) @ B.float().numpy()
+      tag = "ident" if identity else "rand "
+      print(f"({M},{N},{K}) {tag}: tiny-vs-kittens {np.abs(c_hkt-c_hkc).max():9.6f}  "
+            f"tiny-vs-fp64 {np.abs(c_hkt-ref64).max():9.6f}  kittens-vs-fp64 {np.abs(c_hkc-ref64).max():9.6f}")
+      assert np.abs(c_hkt - ref64).max() < 0.26, "tiny kernel must match fp64 at rounding level"
+      assert np.abs(c_hkt - c_hkc).max() < 0.51, "tiny kernel must match hipkittens"
+    # NOTE: hk_bf16_gemm requires K % 128 == 0 (its prologue+epilogue unconditionally touch
+    # k-tiles num_tiles-1 and num_tiles-2); at other K it reads wrong-but-in-bounds global
+    # memory on the emulator and on real hardware, so only K%128==0 sizes are checked here.
+    compare(256, 256, 128, seed=1)   # single workgroup
+    compare(256, 256, 256, seed=2)
+    compare(512, 512, 128, seed=3)   # multi workgroup
+    compare(256, 256, 128, identity=True)   # bit-exact check

test/mockgpu/amd/emu.py

@@ -6,7 +6,7 @@
 #   arg=3: lds - local data share
 #   arg=4: scratch - per-lane scratch memory
 from __future__ import annotations
-import ctypes, functools, re, platform, subprocess, tempfile
+import ctypes, functools, re, platform, subprocess, tempfile, os
 from typing import Callable
 
 # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode
@@ -372,11 +372,33 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals
   return _write_64bit(val, wfn, reg_or_addr, is_mem, *args) if bits == 64 else [wfn(reg_or_addr, _to_u32(val), *args)]
 
 def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, data_bits: int = 32) -> list[UOp]:
-  """Conditional memory store with sub-word support. Returns list of store UOps."""
+  """Conditional memory store with sub-word and unaligned support. Returns list of store UOps.
+
+  FLAT/GLOBAL accesses on AMD hardware are allowed to be unaligned, so a 32-bit store at a
+  byte offset of 1-3 spans two words (handled in 64-bit domain to keep shifts in range)."""
   adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
+  if data_bits > 32:   # wider stores decompose into dwords; each dword handles its own alignment
+    ws = val.cast(dtypes.uint64) if data_bits > 64 else val
+    return [s for i in range(data_bits // 32)
+            for s in _mem_store(mem, addr + UOp.const(adt, i * 4), ws >> UOp.const(ws.dtype, 32 * i) if i else ws, active, addr_bits, 32)]
   word_addr = addr >> UOp.const(adt, 2)
   idx = mem.index(word_addr.valid(active))
-  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
+  if data_bits == 32:
+    byte_off = (addr & UOp.const(adt, 3)).cast(dtypes.uint32)
+    is_unaligned = byte_off.ne(UOp.const(dtypes.uint32, 0))
+    if addr.divides(4) is not None: return [idx.store(active.where(_to_u32(val), idx))]
+    shift = byte_off * UOp.const(dtypes.uint32, 8)
+    val64 = _to_u32(val).cast(dtypes.uint64)
+    # word0 keeps its low byte_off*8 bits, gets val's low bits shifted in; word1 gets the rest
+    low_keep = (UOp.const(dtypes.uint32, 1) << shift) - UOp.const(dtypes.uint32, 1)
+    lo_bits = ((val64 << shift.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 0xFFFFFFFF)).cast(dtypes.uint32)
+    new_word0 = (idx & low_keep) | lo_bits
+    store0 = idx.store(active.where(is_unaligned.where(new_word0, _to_u32(val)), idx))
+    idx1 = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int64).valid(active & is_unaligned))
+    spill = (val64 >> (UOp.const(dtypes.uint64, 32) - shift.cast(dtypes.uint64))).cast(dtypes.uint32)
+    keep = UOp.const(dtypes.uint32, 0xFFFFFFFF) << shift
+    new_word1 = (idx1 & keep) | spill
+    return [store0, idx1.store((active & is_unaligned).where(new_word1, idx1))]
   # Sub-word store: read-modify-write with mask
   byte_pos = addr.cast(dtypes.uint32) & _c(3)
   byte_shift = byte_pos * _c(8)
@@ -2005,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):
@@ -2114,6 +2138,8 @@ def _decode_at(pc: int, arch: str):
 F32_INLINE = {240: 0x3f000000, 241: 0xbf000000, 242: 0x3f800000, 243: 0xbf800000,  # 0.5, -0.5, 1.0, -1.0
               244: 0x40000000, 245: 0xc0000000, 246: 0x40800000, 247: 0xc0800000, 248: 0x3e22f983}  # 2.0, -2.0, 4.0, -4.0, 1/(2*pi)
 
+_inst_hist: dict = {}
+
 class WaveState:
   __slots__ = ('vgpr_buf', 'sgpr_buf', 'accvgpr_buf', '_vgpr_mv', '_sgpr_mv', 'n_lanes', 'wave_size')
 
@@ -2209,7 +2235,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
   # Use Buffer objects with external_ptr=0 for vmem
   vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
   lds_buf = Buffer('CPU', max(lds_size // 4, 1), dtypes.uint32).ensure_allocated()
-  scratch_buf = Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated() if scratch_size else None
+  ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))
+  # NOTE: scratch (private/spill) memory is per-wavefront; buffers are allocated in the wave loop below
 
   # Initialize SQTT encoder — emits packets inline as instructions execute (only when profiling)
   if PROFILE:
@@ -2227,6 +2254,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
         print(colored(msg, 'green') if len(_canonical_runner_cache) > prev_len else msg)
     return program[pc]
 
+  if os.getenv("EMU_TRACE_INST") or os.getenv("EMU_WATCH_PC") is not None:
+    print(f"[emu-dispatch] gx={gx} gy={gy} gz={gz} lx={lx} ly={ly} lz={lz} scratch={scratch_size}", flush=True)
   # Set DAZ+FTZ during emulator execution, restore afterward to avoid breaking hypothesis tests
   # Only trace the first workgroup (like real HW traces one CU/SIMD), subsequent workgroups run but don't add to trace
   tracing = bool(PROFILE)
@@ -2237,18 +2266,27 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
         for gidx in range(gx):
           # Initialize all wavefronts for this workgroup
           waves: list[tuple[WaveState, list]] = []
+          wave_scratch_bufs: list[Buffer] = []   # keep alive for the dispatch
           for wave_start in range(0, total_threads, wave_size):
             st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data,
                             wave_size)
+            # each wavefront owns its private (spill) scratch segment: on real HW the ring is indexed by
+            # (wave_id, lane), and scratch addresses here are lane*stride within the wave's segment.
+            if scratch_size:
+              wave_scratch_bufs.append(Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated())
+              ctypes.memset(wave_scratch_bufs[-1]._buf.va_addr, 0, scratch_size * wave_size)
             c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
                       ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
-                      ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0),
+                      ctypes.c_uint64(wave_scratch_bufs[-1]._buf.va_addr if scratch_size else 0),
                       ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]
             waves.append((st, c_bufs))
 
           # Execute wavefronts with barrier synchronization
           # Each wave runs until it hits s_barrier or s_endpgm. When all waves have stopped, release barrier waves.
           done = [False] * len(waves)
+          _exec_dumped: set = set()
+          _trace_on = bool(os.getenv("EMU_TRACE_INST")) and total_threads >= int(os.getenv("EMU_TRACE_MIN_THREADS", "1"))
+          if _trace_on: _inst_hist.clear()
           for total_inst in range(10_000_000):
             if all(done): break
             for wi, (st, c_bufs) in enumerate(waves):
@@ -2259,9 +2297,23 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                 if pc == ENDPGM_PC:
                   done[wi] = True
                   if tracing: sqtt_finish(wi)
+                  if os.getenv("EMU_EXEC_DUMP") and wi not in _exec_dumped:
+                    _exec_dumped.add(wi)
+                    print(f"[exec-dump] gid=({gidx},{gidy},{gidz}) wave{wi} "
+                          f"exec_lo={st._read_sgpr(EXEC_LO.offset):08x} exec_hi={st._read_sgpr(EXEC_LO.offset+1):08x}")
                   break
                 fxn, globals_list, is_barrier, inst = _ensure_compiled(pc)
                 if DEBUG >= 5: print(f"  exec gid=({gidx},{gidy},{gidz}) w={wi} PC={pc - lib}: {inst!r}", flush=True)
+                if _trace_on:
+                  key = (gidx, gidy, gidz, wi)
+                  _inst_hist.setdefault(key, []).append((pc - lib, type(inst).__name__, getattr(inst, 'op', None) and inst.op.name))
+                wpc, wwave = os.getenv("EMU_WATCH_PC"), int(os.getenv("EMU_WATCH_WAVE", "0"))
+                if wpc is not None and (int(wpc) < 0 or (pc - lib) == int(wpc)) and wi == wwave and gidx == gidy == gidz == 0:
+                  if int(wpc) < 0 and total_inst < 800: print(f"[watch-stream] pc={pc-lib} {inst!r}", flush=True)
+                  for rv in os.getenv("EMU_WATCH_VGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_vgpr(int(rv[1:]), l) for l in range(8)], flush=True)
+                  for rv in os.getenv("EMU_WATCH_SGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_sgpr(int(rv[1:]))], flush=True)
                 fxn(*[c_bufs[g] for g in globals_list])
                 if tracing:
                   inst_op = inst.op.value if hasattr(inst, 'op') else 0
@@ -2269,9 +2321,26 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                 if is_barrier: break  # s_barrier hit: PC already advanced past it, pause this wave
               else: raise RuntimeError("exceeded 1M instructions in single wave, likely infinite loop")
             # All waves have either hit barrier or endpgm — release barrier waves for next round
+            if os.getenv("EMU_DUMP_ROUNDS") and lds_size > 0 and total_inst < int(os.getenv("EMU_DUMP_ROUNDS")):
+              import numpy as _np
+              lds_words = _np.frombuffer((ctypes.c_uint32 * (lds_size//4)).from_address(lds_buf._buf.va_addr), dtype=_np.uint32)
+              nz = _np.argwhere(lds_words != 0)
+              print(f"[emu-dump] round={total_inst} lds nonzero words={len(nz)}",
+                    (f"first16={[f'w{w}:0x{lds_words[w]:08x}' for w in nz[:16].flatten()]}" if len(nz) else ""), flush=True)
+              for wi2, (st2, _) in enumerate(waves):
+                v = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.vgpr_buf._buf.va_addr), dtype=_np.uint32)
+                if st2.wave_size == 64:
+                  av = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.accvgpr_buf._buf.va_addr), dtype=_np.uint32)
+                else: av = _np.zeros(1, dtype=_np.uint32)
+                print(f"[emu-dump]   wave{wi2} vgpr nonzero={int((v!=0).sum())} accvgpr nonzero={int((av!=0).sum())}", flush=True)
           else: raise RuntimeError("exceeded 10M total scheduling rounds")
           tracing = False  # only trace the first workgroup
 
+          if _trace_on:
+            import pickle
+            tag = os.environ["EMU_TRACE_INST"]
+            with open(f"/tmp/emu_trace_{tag}.pkl", "wb") as f: pickle.dump(dict(_inst_hist), f)
+            os.environ.pop("EMU_TRACE_INST", None)   # only dump for the first matching dispatch
           # Reset LDS for next workgroup
           if lds_size > 0: ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))
 

test/null/test_uop_graph.py

@@ -248,6 +248,24 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
 
+  def test_coalesce_aliased_stores(self):
+    from tinygrad.codegen.late.coalesce import memory_coalescing
+    from tinygrad import Device
+    # two distinct stores to the same INDEX (e.g. a double-buffered LDS slot written in a
+    # prologue and a loop body) can't be merged: merging would drop one write
+    lbuf = UOp.placeholder((64,), dtypes.half, 0, AddrSpace.LOCAL)
+    r = UOp.range(4, 0, AxisType.LOOP)
+    gbuf = UOp.placeholder((64,), dtypes.half, 1, AddrSpace.GLOBAL)
+    s_a = lbuf.index(r*2).store(gbuf.index(r*2).load())
+    s_b = lbuf.index(r*2).store(gbuf.index(r*2+64).load())   # aliases s_a
+    s_c = lbuf.index(r*2+1).store(gbuf.index(r*2+32).load()) # adjacent, unique -> coalesceable with nothing
+    out = memory_coalescing(UOp.sink(s_a, s_b, s_c).end(r), Device["NULL"].renderer)
+    stores = [u for u in out.toposort() if u.op is Ops.STORE]
+    # both aliased stores survive (both datas preserved), nothing merged across them
+    self.assertEqual(len([u for u in stores if u.src[0].src[1] is not None and u.src[0].op is Ops.INDEX]), 3)
+    datas = sorted([u.src[1] for u in stores], key=str)
+    self.assertEqual(len(datas), 3)
+
   def test_devectorize_derives_lane_dtype(self):
     from tinygrad.codegen import do_devectorize
     # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit

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()
@@ -137,8 +139,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
       # TODO: a better way to get this than ctx
       lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
     lengths.append(1)  # worst case, it's not folded
+    # stores that alias (same buf+idx+valid from multiple store sites, e.g. a double-buffered LDS
+    # slot written in a prologue and a loop body) can't be merged: merging would drop one write.
+    # keep those scalar and only coalesce the unique ones.
+    keys = [k for k in sorted(offsets.keys()) if op is Ops.LOAD or len(offsets[k]) == 1]
     # do the grouping
-    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
+    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(keys), lambda x: x[1]-x[0])]
     for full_grp in grouped_offsets:
       while len(full_grp):
         offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(None, full_grp[0])
@@ -148,10 +154,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
         offset = offset.valid(valid) if valid is not None else offset
         idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(None, len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
-          datas = []
-          for i,g in enumerate(grp):
-            assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
-            datas.append(offsets[g][0].src[1])
+          datas = [offsets[g][0].src[1] for g in grp]
           store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
           for i,g in enumerate(grp): replacements[offsets[g][0]] = store
         else:

tinygrad/codegen/simplify.py

@@ -59,9 +59,13 @@ pm_simplify_ranges = PatternMatcher([
   (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))),
 ])
 
+SPLITTABLE_TYPES = {AxisType.WEAK, AxisType.REDUCE, AxisType.LOOP}
+
 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
-  # ranges that aren't looped over can't be split
-  if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \
+  # ranges that aren't looped over can't be split. ranges with hardware meaning are never split
+  # (LOCAL/WARP/THREAD/GLOBAL/GROUP_REDUCE/DEVICE map to launch dims; UPCAST/UNROLL are vector
+  # widths): splitting them scrambles the logical<->hardware mapping of hand-written kernels.
+  if r not in ctx and r.arg[-1] in SPLITTABLE_TYPES \
     and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
 
 def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:

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):

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,