geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions 7 Activity

make .barrier implicit (kimi) (#17275)

Browse Source 
* make .barrier implicit (kimi)

* simplier

* lil

* remove tinygrad stock barriers

* readable

* lil
This commit is contained in:
George Hotz
2026-07-28 22:34:57 -07:00
committed by GitHub
parent 2f8f2d2d37
commit 451120c6e1
8 changed files with 51 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/abstractions4.py
+1 -1
View File
@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):
# store all the per lane accumulators to LOCAL
local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
local_accs = local_accs.after(local_accs[lane].store(acc[0]))
# accumulate LOCALs into a single per CU accumulator
late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
extra/gemm/amd_copy_matmul.py
+4 -4
View File
@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
B_copy = B_local.permute((1,0)) if use_wmma else B_local
A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
barrier = UOp.barrier(A_store, B_store)
A_local, B_local = A_local.after(barrier), B_local.after(barrier)
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
# -- COMPUTE --
lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
b_frag = b_frag.reshape(1, TN).expand(TM, TN)
acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
# store accumulator and loop
acc = acc.after(acc_store.end(k).barrier().end(k_tile))
# store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
acc = acc.after(acc_store.end(k).end(k_tile))
# store accumulator to output (unified)
c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
extra/gemm/amd_flash_attention.py
+7 -7
View File
@@ -84,9 +84,9 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
Q_lds = Q_lds.after(qk_load_barrier)
KV_lds_k = KV_lds.after(qk_load_barrier)
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
Q_lds = Q_lds.after(UOp.group(Q_store, K_store))
KV_lds_k = KV_lds.after(UOp.group(Q_store, K_store))
# -- S = Q @ K^T via WMMA (re-init each n_tile) --
S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
@@ -147,9 +147,9 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
# load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
P_lds = P_lds.after(pv_barrier)
KV_lds_v = KV_lds.after(pv_barrier)
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
P_lds = P_lds.after(UOp.group(P_store, V_store))
KV_lds_v = KV_lds.after(UOp.group(P_store, V_store))
# -- acc += P @ V via WMMA --
k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
@@ -161,7 +161,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
# end KV tile loop
n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).end(n_tile)
acc = acc.after(n_tile_end)
l_i = l_i.after(n_tile_end)
m_i = m_i.after(n_tile_end)
extra/gemm/amd_uop_matmul.py
+3 -4
View File
@@ -66,9 +66,8 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
# TODO: can we automate barrier?
barrier = UOp.barrier(A_local_store, B_local_store)
A_local, B_local = A_local.after(barrier), B_local.after(barrier)
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
# open inner k range
k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -102,7 +101,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
# Close k, sync, and close K tiles
sink = sink.end(k).barrier().end(k_tile_range)
sink = sink.end(k).end(k_tile_range)
# ---------------------------
# REG -> GLOBAL (epilogue)
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+2 -2
View File
@@ -36,13 +36,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
lmax_val = lmax.after(lmax_store.end(it))[0]
lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
lds = lds.after(lds[tid].store(lmax_val).barrier())
lds = lds.after(lds[tid].store(lmax_val))
step = THREADS_PER_WG // 2
while step:
active = tid < step
other = lds[(tid + step).valid(active)].load()
lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
step //= 2
device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
extra/thunder/tiny/tk/group.py
+3 -3
View File
@@ -224,7 +224,7 @@ class Group:
# store to shared memory
red_local_store = red_local[self.laneid].store(red_reg[0])
red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
red_local = red_local.after(red_local_store).reshape(red_local.shape)
# reduce from shared memory
for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -258,7 +258,7 @@ class Group:
# store to shared memory
red_local_store = red_local[self.laneid].store(red_reg[0])
red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
red_local = red_local.after(red_local_store).reshape(red_local.shape)
# reduce from shared memory
for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -342,7 +342,7 @@ class Group:
if src.dtype != dst.dtype:
src_load = src_load.cast(dst.dtype)
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
dst_store = dst_store.end(height, width, outer, inner).barrier()
dst_store = dst_store.end(height, width, outer, inner)
elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
srcf = src.flatten()
row_stride = prod(src.shape[axis+1:])
tinygrad/codegen/__init__.py
+30 -1
View File
@@ -242,7 +242,7 @@ pm_add_loads = PatternMatcher([
def add_local_buffer(ctx, x:UOp):
buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace)
return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier())
return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]))
pm_add_local_buffers = PatternMatcher([
(UPat(Ops.STAGE, name="x"), add_local_buffer),
@@ -255,6 +255,32 @@ pm_cast_float_alu = PatternMatcher([
lambda u,x: u.replace(src=(x.cast(u.dtype),)) if x.dtype != u.dtype else None),
])
def _is_local_store(x:UOp): return x.op is Ops.STORE and x.addrspace is AddrSpace.LOCAL
def add_raw_barrier(after:UOp):
# loads from a LOCAL buffer that depend (via AFTER) on stores to LOCAL memory need a workgroup barrier
if after.addrspace is not AddrSpace.LOCAL: return None
# one toposort over all the deps
deps = UOp.sink(*after.src[1:]).backward_slice
if not any(_is_local_store(x) for x in deps) or any(x.op is Ops.BARRIER for x in deps): return None
return after.src[0].after(UOp(Ops.BARRIER, src=after.src[1:]))
def add_war_barrier(end:UOp):
# a LOCAL buffer stored and loaded in the same loop needs a barrier at the end of the loop body
rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.LOOP) and r.vmax > 0]
if not rngs or end.src[0].op is Ops.BARRIER: return None
sl = end.src[0].backward_slice_with_self
# only stores that are inside this loop body (not in the backward slice through AFTER chains from other loops)
store_bufs = {x.buf_uop for x in sl if _is_local_store(x) and any(r in x.ranges for r in rngs)}
# a load whose buffer matches a local store's buffer is necessarily a local load
if not (loads:=[x for x in sl if x.op is Ops.LOAD and x.src[0].buf_uop in store_bufs]): return None
return end.replace(src=(UOp(Ops.BARRIER, src=(end.src[0], *loads)),)+end.src[1:])
pm_implicit_barriers = PatternMatcher([
(UPat(Ops.AFTER, name="after"), add_raw_barrier),
(UPat(Ops.END, name="end"), add_war_barrier),
])
def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
if DEBUG >= 5: print(pyrender(ast))
@@ -344,6 +370,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends
sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
# add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")
# this was the linearizer
sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
tinygrad/schedule/rangeify.py
+1 -1
View File
@@ -390,7 +390,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
# handle locals
buf = UOp.placeholder((size,), dtype, next(ctx), AddrSpace.LOCAL)
do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
return buf.after(do_store.barrier()).cast(x.dtype)
return buf.after(do_store).cast(x.dtype)
# collapse any BUFFERIZE to single input BUFFERIZE
def flatten_bufferize(x:UOp):