beautiful mnist

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):

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):

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]))}")

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):