From 30ca3f2af8d2c96e8d29a6580c1f076dfda2ba0c Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 29 Oct 2025 16:25:27 +0800 Subject: [PATCH] all double matmul (#12993) * fix more double matmuls * a few more * all double matmul passes * opts for flash attention * fix spec * comment --- test/test_rangeify.py | 23 ++++++++++++++++------- tinygrad/codegen/late/devectorizer.py | 7 +++++-- tinygrad/codegen/opt/postrange.py | 11 +++++++++-- tinygrad/schedule/rangeify.py | 4 ++++ tinygrad/uop/spec.py | 3 ++- tinygrad/uop/symbolic.py | 2 +- 6 files changed, 37 insertions(+), 13 deletions(-) diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 7a5059e80d..e25af7ff7b 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -39,14 +39,14 @@ class TestDoubleMatmul(unittest.TestCase): def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4))) def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4))) - @unittest.skip("doesn't work") def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4))) def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4))) - @unittest.skip("doesn't work") + def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2))) + def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2))) + def test_upcast_01_unroll_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4))) - @unittest.skip("doesn't work") def test_upcast_12_unroll_01(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4))) @@ -83,7 +83,7 @@ elif getenv("BIG") > 1: BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128 elif getenv("BIG") > 0: # bigger - BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64 + BS, HEADS, SEQLEN, EMB = 4, 32, 128, 128 else: BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8 @@ -130,9 +130,9 @@ class TestPcontig(unittest.TestCase): print(f"mse: {mse}") self.assertLessEqual(mse, 1e-6) - def test_flash_attention(self): - with Context(PCONTIG=2, DEBUG=2): - ret = fa().realize() + def test_flash_attention(self, opts=None): + with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)): + ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize() print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS") with Context(DEBUG=2): cmp = fa().realize() @@ -142,6 +142,15 @@ class TestPcontig(unittest.TestCase): print(f"mse: {mse}") self.assertLessEqual(mse, 1e-6) + def test_flash_attention_opt(self): + opts = () + # columns in top matrix + opts += (Opt(OptOps.UPCAST, 0, 4),) + # columns in bottom matrix + opts += (Opt(OptOps.UPCAST, 3, 4),) + # rows in all the matrix + opts += (Opt(OptOps.UPCAST, 4, 4),) + self.test_flash_attention(opts) # *** non CI rangeify tests below this line *** diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 65e21096c7..b7d9c81bdd 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -233,9 +233,10 @@ def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp): def no_vectorized_index_broadcast(buf:UOp, cast:UOp, bcast:UOp, idx:UOp): cnt = cast.dtype.count - precnt = len(bcast.src) + precnt = bcast.dtype.vcount + input_gep = bcast.arg if bcast.op is Ops.GEP else ([0]*precnt) gep_arg = tuple(flatten([range(precnt) for _ in range(cnt)])) - sum_arg = tuple(flatten([[i]*precnt for i in range(cnt)])) + sum_arg = tuple(flatten([[i+y for y in input_gep] for i in range(cnt)])) return buf.broadcast(cnt*precnt).index(idx.gep(gep_arg)*cnt+UOp.const(dtypes.index.vec(cnt*precnt), sum_arg)) devectorize_buf_and_index = PatternMatcher([ @@ -243,6 +244,8 @@ devectorize_buf_and_index = PatternMatcher([ (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index), (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").broadcast(name="bcast").index(UPat.var("idx")), no_vectorized_index_broadcast), + (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").gep(name="bcast").index(UPat.var("idx")), + no_vectorized_index_broadcast), ]) devectorize = PatternMatcher([ diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index c2a82d5c85..c0ac187252 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -63,8 +63,15 @@ class Scheduler: self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range") return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1) - def _globalizable_rngs(self) -> list[UOp]: + def _output_rngs(self) -> list[UOp]: return flatten([list(UOp.sink(*s.src[1:]).ranges) for s in self.ast.src if s.op is Ops.END]) + def _globalizable_rngs(self) -> list[UOp]: + ret = self._output_rngs() + # exclude any output ranges from global that don't appear in all BUFFERIZE + for x in self.ast.toposort(): + if x.op is Ops.BUFFERIZE: + ret = [r for r in ret if r in x.ranges] + return ret def convert_loop_to_global(self): if not self.ren.has_local: return None @@ -75,7 +82,7 @@ class Scheduler: self.ast = self.ast.substitute(dict(zip(self.rngs, rng))) def colors(self) -> list[str]: - output_rngs = self._globalizable_rngs() + output_rngs = self._output_rngs() ret = [] for x,r in zip(self.axis_types, self.rngs): if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 83218d3989..86dd580dd7 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -444,6 +444,10 @@ rangeify_codegen = PatternMatcher([ (UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True).broadcast(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True), lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))), + (UPat(Ops.AFTER, name="a").gep(name="b"), lambda a,b: a.gep(b.arg)), + (UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True).gep(name="dg").f(Ops.INDEX, name="idx", allow_any_len=True), + lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else + idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))), ]) def remove_metadata_tags(ctx:LocalAddBufferContext, x:UOp): diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 72a81f2019..d882f4206c 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -230,8 +230,9 @@ full_spec = PatternMatcher([ # in progress MSTACK may lose device (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True), - # temp VECTORIZEs during rewrite have the wrong dtype + # temp VECTORIZE/INDEX during rewrite have the wrong dtype (UPat(Ops.VECTORIZE), lambda: True), + (UPat(Ops.INDEX), lambda: True), # all loads/stores (UPat((Ops.LOAD, Ops.STORE)), lambda: True), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index e07ba57f58..c6405db546 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -265,7 +265,7 @@ gep_pushing = PatternMatcher([ # push all GEPs through ALUs (fix arange stuff) (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, name='gep'), lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \ - if not isinstance(gep.dtype, PtrDType) else None), + if not isinstance(gep.dtype, PtrDType) and not isinstance(alu.dtype, PtrDType) else None), # CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later) (UPat(Ops.CAT, name="x"), lambda x: UOp(Ops.VECTORIZE, x.dtype, tuple(y.gep(i) for y in x.src for i in range(y.dtype.count))) \ if not isinstance(x.dtype, PtrDType) else None),