diff --git a/docs/abstractions4.py b/docs/abstractions4.py index c8167b3ce9..d28a3e3a34 100644 --- a/docs/abstractions4.py +++ b/docs/abstractions4.py @@ -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) diff --git a/extra/gemm/amd_copy_matmul.py b/extra/gemm/amd_copy_matmul.py index 5d9b4c701a..cd86c7cea0 100644 --- a/extra/gemm/amd_copy_matmul.py +++ b/extra/gemm/amd_copy_matmul.py @@ -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, diff --git a/extra/gemm/amd_flash_attention.py b/extra/gemm/amd_flash_attention.py index 73740d8165..7e134208d8 100644 --- a/extra/gemm/amd_flash_attention.py +++ b/extra/gemm/amd_flash_attention.py @@ -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) diff --git a/extra/gemm/amd_uop_matmul.py b/extra/gemm/amd_uop_matmul.py index e83233b04a..8db8a608ee 100644 --- a/extra/gemm/amd_uop_matmul.py +++ b/extra/gemm/amd_uop_matmul.py @@ -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) diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index 276c44615b..e7b00c8b61 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -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] diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index 08f847dbd3..dca8dc375d 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -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:]) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 39eb47c359..3a8d7a0583 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index fdd0601072..80833ba322 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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):