From 7a95d19e7ed22a718bbe34836ad8ed5bc15f7bfd Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 25 Aug 2025 20:35:55 -0700 Subject: [PATCH] sort of wmma --- test/test_rangeify.py | 32 +++++++++---- test/test_tiny.py | 5 ++- tinygrad/codegen/__init__.py | 5 ++- tinygrad/codegen/gpudims.py | 75 ++++++++++++++++++++++++++++++- tinygrad/codegen/late/expander.py | 9 ++-- tinygrad/codegen/opt/__init__.py | 2 +- tinygrad/codegen/opt/kernel.py | 2 +- tinygrad/codegen/opt/postrange.py | 7 +++ tinygrad/runtime/ops_null.py | 1 + tinygrad/schedule/rangeify.py | 5 ++- tinygrad/uop/ops.py | 7 +++ 11 files changed, 129 insertions(+), 21 deletions(-) diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 88c388095e..9f00775e5c 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor -from tinygrad.helpers import RANGEIFY +from tinygrad.helpers import RANGEIFY, Context, GlobalCounters N = 256 @@ -96,14 +96,28 @@ class TestRangeify(unittest.TestCase): out.realize() def test_flash_attention(self): - BS = 4 - HEADS = 2 - MATDIM = 16 - EMB = 8 - q = Tensor.empty(BS, HEADS, MATDIM, EMB) - k = Tensor.empty(BS, HEADS, MATDIM, EMB) - v = Tensor.empty(BS, HEADS, MATDIM, EMB) - q.scaled_dot_product_attention(k, v).realize() + #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8 + BS, HEADS, SEQLEN, EMB = 4, 16, 128, 64 + + # llama 8B + #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128 + + def fa(): + Tensor.manual_seed(1337) + with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)] + return q.scaled_dot_product_attention(k, v).realize() + + with Context(DEBUG=4): + GlobalCounters.reset() + ret = fa() + with Context(RANGEIFY=0): + with Context(DEBUG=2): + GlobalCounters.reset() + cmp = fa() + with Context(DEBUG=0): + mse = ((cmp-ret)**2).sum().item() + print(f"mse: {mse}") + self.assertLessEqual(mse, 1e-6) from tinygrad import dtypes from tinygrad.uop.ops import UOp diff --git a/test/test_tiny.py b/test/test_tiny.py index e86a726944..b784a869e1 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -30,7 +30,10 @@ class TestTiny(unittest.TestCase): def test_gemm(self, N=64, out_dtype=dtypes.float): a = Tensor.ones(N,N).contiguous() b = Tensor.eye(N).contiguous() - self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N)) + lst = (out:=a@b).tolist() + for y in range(N): + for x in range(N): + self.assertEqual(lst[y][x], 1.0) if IMAGE < 2: self.assertEqual(out.dtype, out_dtype) # *** randomness *** diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index eeae493a6f..f35c6ff683 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -9,7 +9,7 @@ from tinygrad.renderer import Renderer # import all pattern matchers here from tinygrad.codegen.lowerer import pm_lowerer, get_index from tinygrad.codegen.quantize import pm_quant -from tinygrad.codegen.gpudims import pm_add_gpudims +from tinygrad.codegen.gpudims import pm_add_gpudims, pm_tensor_cores from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing from tinygrad.uop.decompositions import get_late_rewrite_patterns from tinygrad.codegen.late.expander import migrate_indexing, expander @@ -63,6 +63,9 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize")) ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True)) + # add tensor cores + if _RANGEIFY: ret.append(RewriteStep(pm_tensor_cores, lambda _: ({}, opts), name="tensor cores")) + if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast")) # ** expander (expand_rewrite) ** diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index d510f8aba4..08d20fa4ed 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -1,9 +1,10 @@ import math, functools, operator -from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType +from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite from tinygrad.helpers import all_int, partition, flatten, prod, dedup from tinygrad.dtype import dtypes, AddrSpace from tinygrad.shape.view import get_contraction from tinygrad.renderer import Renderer +from tinygrad.codegen.opt.tc import TensorCore def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]): # TODO: symbolic shape @@ -114,7 +115,7 @@ def fix_group_for_reduce(x:UOp): # do only the non grouped reduces early ret = x.replace(src=(x.src[0],)+tuple(reduce_r)) reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr] - buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop) + buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=AddrSpace.LOCAL).index(*upstream_locals, *reduce_loop) # gate with an if on the store + do the final reduce buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf)) @@ -133,3 +134,73 @@ pm_add_gpudims = PatternMatcher([ # fix group for reduce (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce), ]) + +def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp): + # tensor cores have three ranges. X, Y, and REDUCE + in0_ranges = [u for u in in0.ranges if u not in in1.ranges] + in1_ranges = [u for u in in1.ranges if u not in in0.ranges] + if len(in0_ranges) != 1 or len(in1_ranges) != 1: return None + in0_range, in1_range = in0_ranges[0], in1_ranges[0] + + # confirm the dtype and size is good + tc_opts: list[TensorCore] = [] + for tc in ctx[1].tensor_cores: + if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in: + if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])): + tc_opts.append(tc) + if len(tc_opts) == 0: return None + tc = tc_opts[0] + + # create the new ranges as speced by the tensor core + old_range = [in0_range, in1_range, r_range] + new_range = [r.replace(src=(r.src[0]//tc.dims[i],)) for i,r in enumerate(old_range)] + new_reduce_range = new_range[2] + tc_range = 9050 + r_range.arg[0]*100 + red_ranges = [] + + ne: list[UOp] = [] + for o in tc.opts: + lrange = UOp.range(dtypes.int, 2, tc_range, AxisType.UPCAST if o[0] == "u" else AxisType.LOCAL) + ne.append(lrange) + tc_range += 1 + new_range[int(o[1])] = (2 * new_range[int(o[1])]) + lrange + for _, amt in tc.get_reduce_axes(): + lrange = UOp.range(dtypes.int, amt, tc_range, AxisType.UNROLL) + ne.append(lrange) + red_ranges.append(lrange) + tc_range += 1 + new_range[2] = (amt * new_range[2]) + lrange + tne = [x.replace(tag=1) for x in ne] + + # replace ranges in other parts of the graph + for x,y in zip(old_range, new_range): ctx[0][x] = y + + # apply the swizzled ranges to the srcs + srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)] + srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))] + + # get reduce/upcast axes for the tensor cores + ned = dict(zip(tc.base_shape_str(), ne)) + tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))]) + base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()]) + tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)]) + + # construct the op + # TODO: remove tc_upcast_axes from the arg + wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].device, tc.threads, tc_upcast_axes, tc_reduce_axes) + wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=( + UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]), + UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]), + UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0))+tuple(red_ranges), arg=wmma_arg) + tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2]) + return tc_uop.reduce(new_reduce_range, arg=Ops.ADD) + +from tinygrad.codegen.opt.postrange import pm_flatten_range + +pm_tensor_cores = PatternMatcher([ + ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores), + + # replace range + #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)), + (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range)), +]) \ No newline at end of file diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index 46fb5d990e..1e4961e68d 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -50,7 +50,7 @@ def do_expand(root:UOp): if root.op is Ops.IF or src.op is Ops.IF: # for the first arg of IF, just pass them through ignoring UNROLLS new_srcs.append(src) - elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1): + elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3): # for any range args of STORE/REDUCE, pass them through new_srcs.append(src) elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType): @@ -89,9 +89,10 @@ expander = PatternMatcher([ (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand), (UPat(Ops.CONTRACT, name="con"), do_contract), - # BARRIERs aren't actually expanded - (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)), - lambda ex: UOp(Ops.UNROLL, src=(UOp(Ops.BARRIER, src=ex.src),)*len(ex.src), arg=ex.arg)), + # BARRIERs aren't actually expanded, now they even block expands + (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)), lambda ex: UOp(Ops.BARRIER, src=ex.src)), + #(UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)), + # lambda ex: UOp(Ops.UNROLL, src=(UOp(Ops.BARRIER, src=ex.src),)*len(ex.src), arg=ex.arg)), # empty UNROLL is NOOP (UPat(Ops.UNROLL, src=(UPat.var('x'),), arg=()), lambda x: x), # UNROLL GEP (needed for WMMA, generalize this) -> vectorized ALU diff --git a/tinygrad/codegen/opt/__init__.py b/tinygrad/codegen/opt/__init__.py index 38c430d80c..09fd766079 100644 --- a/tinygrad/codegen/opt/__init__.py +++ b/tinygrad/codegen/opt/__init__.py @@ -37,7 +37,7 @@ pm_get_optimization = PatternMatcher([ def apply_opt(ast:UOp, renderer:Renderer): k = Kernel(ast, opts=renderer) - if ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply) + k.apply_opts(ast.arg.opts_to_apply) ret = k.get_optimized_ast() if __debug__: type_verify(list(ret.toposort())) return ret diff --git a/tinygrad/codegen/opt/kernel.py b/tinygrad/codegen/opt/kernel.py index 664ff573f1..14a50b8c03 100644 --- a/tinygrad/codegen/opt/kernel.py +++ b/tinygrad/codegen/opt/kernel.py @@ -86,7 +86,7 @@ class Kernel: # group simplifies self.simplify_ones() - #self.simplify_merge_adjacent() + self.simplify_merge_adjacent() # axis types global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 8864303b43..d4c4ff15d3 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -1,4 +1,5 @@ import math +from tinygrad.dtype import AddrSpace from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite, _substitute from tinygrad.uop.symbolic import symbolic from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps @@ -29,12 +30,18 @@ class RKernel(Kernel): self.replaces = {} if self.opts.has_local: store_rngs = self.ast.src[0].src[2:] + + # filter any not in local stores + local_store_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL] + for ls in local_store_rngs: store_rngs = [x for x in store_rngs if x in ls] + store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else [] rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rng] self.replaces.update(dict(zip(self.rng, rng))) self.rng = rng def simplify_merge_adjacent(self): + return # NOTE: this one is better than the one in kernel.py, which is kind of a problem terminators = [u for u in self.ast.toposort() if u.op in {Ops.REDUCE, Ops.STORE}] termination = {} diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index 10ae40f4ed..8223e5ddee 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -7,6 +7,7 @@ class NullRenderer(CStyleLanguage): device = "NULL" has_local = False float4 = "float4" + barrier = "// BARRIER" code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"} class NullProgram: diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 480ee24034..dca9c3021f 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -338,13 +338,14 @@ def bufferize_to_store(x:UOp): shape = tuple([int(r.vmax+1) for r in rngs]) size = prod(shape) assert size > 0, f"no zero sized buffers {shape}" - sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, tuple) else x.arg[0]) + sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, AddrSpace) else x.arg) if x.src[0].op is Ops.ASSIGN: assign_target, assign_src = x.src[0].src assert assign_target.op is Ops.INDEX return assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=sdtype) + # NOTE: the DEFINE_LOCAL needs to be disambiguated here if sdtype.addrspace == AddrSpace.GLOBAL: buf = UOp.new_buffer(x.arg, size, x.dtype) - else: buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=x.arg[1]) + else: buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=UOp.unique().arg) return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype) pm_add_buffers = pm_mops+PatternMatcher([ diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 91c9e0c23c..b8989c1ae8 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -142,6 +142,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}: return None + if self.op is Ops.BARRIER: return None if self.op in GroupOp.Block: return None from tinygrad.shape.shapetracker import ShapeTracker # VIEW and MovementOps define a new ShapeTracker from the arg @@ -211,6 +212,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass): ret.update(self.src[1].ranges) for s in self.src[2:]: if s in ret: del ret[s] + elif self.op in {Ops.WMMA}: + ret = self.src[0].ranges.copy() + ret.update(self.src[1].ranges) + ret.update(self.src[2].ranges) + for s in self.src[3:]: + if s in ret: del ret[s] else: ret = {} for s in self.src: ret.update(s.ranges)