From 5bab842337082f94c2deeaa186df93afef126733 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 1 Aug 2025 15:12:17 -0700 Subject: [PATCH] beautiful mnist --- test/test_tiny.py | 6 +++ tinygrad/codegen/rangeify.py | 81 ++++++++++++++-------------------- tinygrad/schedule/kernelize.py | 2 + tinygrad/uop/ops.py | 4 +- 4 files changed, 43 insertions(+), 50 deletions(-) diff --git a/test/test_tiny.py b/test/test_tiny.py index 0d3f3a21af..e72e7a8320 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -37,6 +37,12 @@ class TestTiny(unittest.TestCase): a = Tensor.eye(4, dtype=dtypes.int) self.assertListEqual(a.tolist(), [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) + def test_conv(self, N=32): + a = Tensor.ones(1,4,N,N).contiguous() + w1 = Tensor.ones(16,4,3,3).contiguous() + out = a.conv2d(w1) + self.assertTrue(all([x == 36.0 for x in out.contiguous().flatten().tolist()])) + # *** randomness *** def test_random(self): diff --git a/tinygrad/codegen/rangeify.py b/tinygrad/codegen/rangeify.py index b244fb8754..ef7caecf85 100644 --- a/tinygrad/codegen/rangeify.py +++ b/tinygrad/codegen/rangeify.py @@ -15,6 +15,18 @@ def rangify_store(ctx:list[int], x:UOp): mm2 = UOp(Ops.INDEX, dtype=x.src[0].dtype, src=(x.src[1],)+tuple(ranges)) return UOp(Ops.STORE, src=(mm, mm2)+tuple([x for x in ranges if x.op is not Ops.CONST]), tag=1) +def map_reduce(ctx:list[int], x:UOp, r:UOp): + rngs = list(x.src[1:]) + new_ranges = [] + for i,s in enumerate(r.src[0].shape): + if i in r.arg[1]: + assert rngs[i].op == Ops.CONST + rngs[i] = UOp.range(dtypes.int, s, ctx[0]) + new_ranges.append(rngs[i]) + ctx[0] += 1 + mm = UOp(Ops.INDEX, r.src[0].dtype, src=(r.src[0],)+tuple(rngs)) + return UOp(Ops.REDUCE, r.dtype, src=(mm,)+tuple(new_ranges), arg=r.arg[0]) + def map_reshape(x:UOp, r:UOp): acc = 1 to_sum = [] @@ -33,35 +45,6 @@ def map_reshape(x:UOp, r:UOp): ret = UOp.sink(*ret).simplify().src[::-1] if len(ret) else () return UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)) -def map_expand(x:UOp, r:UOp): - inp_shape, exp_shape = x.src[0].src[0].shape, x.src[0].shape - ret = list(x.src[1:]) - exp_ranges = [] - for i,(x,y) in enumerate(zip(inp_shape, exp_shape)): - if x != y: - exp_ranges.append(ret[i]) - ret[i] = UOp.const(dtypes.int, 0) - return UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)) - -def map_permute(x:UOp, r:UOp): - ret = x.src[1:] - # argsort or not? - perm = argsort(x.src[0].arg) - ret = tuple([ret[p] for p in perm]) - return UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)) - -def map_shrink(x:UOp, r:UOp): - ret = list(x.src[1:]) - for i,(s,(ss,se)) in enumerate(zip(r.src[0].shape, r.arg)): - if ss != 0: ret[i] = ret[i] + ss - return UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)) - -def map_flip(x:UOp, r:UOp): - ret = list(x.src[1:]) - for i,(s,a) in enumerate(zip(r.shape, r.arg)): - if a: ret[i] = (s-1)-ret[i] - return UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)) - def map_pad(x:UOp, r:UOp): ret = list(x.src[1:]) bigwhere = UOp.const(dtypes.bool, True) @@ -78,29 +61,25 @@ def map_pad(x:UOp, r:UOp): # PAD is with 0 return bigwhere.simplify().where(UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)), UOp.const(r.dtype, 0)) -def map_reduce(ctx:list[int], x:UOp, r:UOp): - rngs = list(x.src[1:]) - new_ranges = [] - for i,s in enumerate(r.src[0].shape): - if i in r.arg[1]: - assert rngs[i].op == Ops.CONST - rngs[i] = UOp.range(dtypes.int, s, ctx[0]) - new_ranges.append(rngs[i]) - ctx[0] += 1 - mm = UOp(Ops.INDEX, r.src[0].dtype, src=(r.src[0],)+tuple(rngs)) - return UOp(Ops.REDUCE, r.dtype, src=(mm,)+tuple(new_ranges), arg=r.arg[0]) pm_rangeify = PatternMatcher([ - # TODO: handle MAP on STORE + # TODO: handle INDEX on STORE (UPat(Ops.STORE, name="x"), rangify_store), - (UPat(Ops.INDEX, src=(UPat(Ops.PERMUTE, name="r"),), allow_any_len=True, name="x"), map_permute), - (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, name="r"),), allow_any_len=True, name="x"), map_reshape), - (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="x"), map_expand), - (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="r"),), allow_any_len=True, name="x"), map_shrink), - (UPat(Ops.INDEX, src=(UPat(Ops.FLIP, name="r"),), allow_any_len=True, name="x"), map_flip), - (UPat(Ops.INDEX, src=(UPat(Ops.PAD, name="r"),), allow_any_len=True, name="x"), map_pad), (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="r"),), allow_any_len=True, name="x"), map_reduce), - (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),)), lambda c: c), + + # this is like the definitions of these + (UPat(Ops.INDEX, src=(UPat(Ops.PERMUTE, name="r"),), allow_any_len=True, name="x"), + lambda r,x: UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple([x.src[1+p] for p in argsort(x.src[0].arg)]))), + (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="r"),), allow_any_len=True, name="x"), + lambda r,x: UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple([a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(x.src[1:], r.arg)]))), + (UPat(Ops.INDEX, src=(UPat(Ops.FLIP, name="r"),), allow_any_len=True, name="x"), + lambda r,x: UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple([((s-1)-a) if f else a for a,s,f in zip(x.src[1:], r.shape, r.arg)]))), + (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="x"), + lambda r,x: UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+ + tuple([a.const_like(0) if resolve(x!=y, False) else a for a,x,y in zip(x.src[1:], r.src[0].shape, r.shape)]))), + + (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, name="r"),), allow_any_len=True, name="x"), map_reshape), + (UPat(Ops.INDEX, src=(UPat(Ops.PAD, name="r"),), allow_any_len=True, name="x"), map_pad), # bring where to the front #(UPat(GroupOp.Binary, name="base", src=(UPat.var("c").where(UPat.var("x"), UPat(Ops.INVALID, name="inv")), UPat.var("a"))), @@ -115,6 +94,12 @@ pm_rangeify = PatternMatcher([ # move MAP through elementwise ALU (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.LOAD})),), allow_any_len=True, name="x"), lambda x: x.src[0].replace(src=tuple([UOp(Ops.INDEX, dtype=s.dtype, src=(s,)+x.src[1:]) for s in x.src[0].src]))), + + # CONST can't have axes + (UPat(Ops.INDEX, src=(UPat(Ops.CONST,name="c"),)), lambda c: c), + + # unbind...but this is too late + (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR, name="v"), UPat(Ops.CONST))), lambda v: v), ]) def name_the_sink(x:UOp): diff --git a/tinygrad/schedule/kernelize.py b/tinygrad/schedule/kernelize.py index fb38734dbb..f3c5b17571 100644 --- a/tinygrad/schedule/kernelize.py +++ b/tinygrad/schedule/kernelize.py @@ -419,6 +419,8 @@ remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace( new_fixups = PatternMatcher([ (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)), + # TODO: this should be BUFFER_VIEW + (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)), ]) @track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}") diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8627c35313..342c77e900 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -156,7 +156,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): # otherwise we get the shape from sources if not (src_sts := [x.st for x in self.src if x.st is not None]): return None - assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}" + if not all_same([x.shape for x in src_sts]): raise RuntimeError(f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}") match self.op: case Ops.MULTI: shape = tuple(self.src[0].shape[a]*len(self.device) if a == self.axis else s for a,s in enumerate(self.src[0].shape)) case Ops.BITCAST: @@ -220,7 +220,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): # constants can optionally have a DEVICE source try: st = self.st - except AssertionError: + except RuntimeError: st = None return UOp.const(self.dtype, b, device=self._device, shape=st.shape if st is not None else None) def broadcast(self, count:int):