nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

direct lds

Browse Source
This commit is contained in:
geohot
2026-07-31 09:03:32 -07:00
parent 1cfaa385d6
commit 9540700bbc
5 changed files with 129 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/hk_gemm_frag.py
+80 -32
View File
@@ -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)
test/mockgpu/amd/emu.py
+3 -1
View File
@@ -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):
tinygrad/codegen/late/coalesce.py
+3 -1
View File
@@ -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()
tinygrad/renderer/cstyle.py
+42 -1
View File
@@ -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
+1 -1
View File
@@ -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,