Merge branch 'master' into fix_range_merging_sd

test/external/external_uop_gc.py

@@ -63,7 +63,6 @@ if __name__ == "__main__":
     views_to_valid_uop.cache_clear()
 
     new_uops = uops_allocated()
-    print_uops()
     gc.collect()
     new_uops_gc = uops_allocated()
     print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")

test/test_linearizer.py

@@ -123,7 +123,6 @@ class TestLinearizer(unittest.TestCase):
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
 
-  @unittest.skip("this is handled at higher level now")
   def test_upcast_cse(self):
     # when upcasting, within a subtree, there may be common expressions.
 

test/test_softmax_fusion.py

@@ -44,7 +44,6 @@ class TestFuse(unittest.TestCase):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a / a.mean(axis=1), a)
 
-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
   def test_fuse_argmax(self):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a.argmax(axis=-1), a)

test/test_uops.py

@@ -544,6 +544,86 @@ class TestUopsObject(unittest.TestCase):
     with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
     assert len(ret) == 10000
 
+
+class TestShapeSpec(unittest.TestCase):
+  # ** CONST is CONST(VIEW(DEVICE)) -> RESHPAE -> EXPAND
+
+  def test_expanded_const(self):
+    a = Tensor(1).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+    a = Tensor.ones((4, 4)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))
+
+  # NOTE: CONST ShapeTracker comes from its source
+  def test_scalar_const(self):
+    a = Tensor(0).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+
+  def test_scalar_var(self):
+    vv = UOp.variable("a", 1, 4).bind(2)
+    t = Tensor(vv).uop
+    self.assertEqual(t.st, ShapeTracker.from_shape(()))
+
+  # ** ASSIGN is ASSIGN(VIEW(BUFFER), new_val)
+
+  def test_assign_flat(self):
+    buffer = Tensor.arange(4).realize()
+    a = buffer.assign(Tensor.zeros((4,), dtype=dtypes.int))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.BUFFER), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  def test_assign_permuted(self):
+    buffer = Tensor.arange(4).reshape(2, 1, 2).contiguous().realize()
+    a = buffer.permute((1, 2, 0)).assign(Tensor.arange(4).reshape(1, 2, 2).contiguous())
+    a.realize()
+    self.assertEqual(buffer.tolist(), [[[0, 2]], [[1, 3]]])
+
+  def test_assign_reshaped(self):
+    buffer = Tensor.ones((4,)).contiguous().realize()
+    a = buffer.reshape((2, 2)).assign(Tensor.zeros((2, 2)))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER))), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  # setitem is a partial assign
+  def test_setitem(self):
+    a = Tensor.ones((4,)).contiguous().realize()
+    assign = a.shrink(((1, 2),)).assign(Tensor.zeros((1,)))
+    # the ASSIGN UOp has size=1
+    self.assertEqual(assign.uop.size, 1)
+    # the ASSIGN views the buffer with a shrunk st
+    self.assertEqual(assign.uop.src[0].st, ShapeTracker.from_shape((4,)).shrink(((1, 2),)))
+    # the underlying BUFFER has a size=4
+    self.assertEqual(assign.uop.buf_uop.size, 4)
+    # NOTE: output shape is different from the BUFFER shape
+    self.assertNotEqual(assign.uop.shape, a.uop.shape)
+    assign.realize()
+    self.assertEqual(a.tolist(), [1, 0, 1, 1])
+
+  def test_buffer_st(self):
+    a = UOp.new_buffer(Device.DEFAULT, 10, dtypes.float)
+    self.assertEqual(a.st, ShapeTracker.from_shape((10,)))
+
+  def test_ops_st(self):
+    # view / mop
+    a = Tensor.empty(4, 2, 1).permute((1, 2, 0)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape((4, 2, 1)).permute((1, 2, 0)))
+    # alu / reduce
+    alu = a*2
+    self.assertEqual(alu.st, ShapeTracker.from_shape((2, 1, 4)))
+    r = Tensor.empty(4, 4).sum(axis=1)
+    self.assertEqual(r.uop.st, ShapeTracker.from_shape((4,)))
+
+  def test_st_wmma_none(self):
+    A = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('a', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 1)))
+    B = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('b', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 2)))
+    C = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('c', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 3)))
+    wmma = UOp(Ops.WMMA, dtypes.float.vec(16), (A, B, C))
+    assert wmma.st is None
+
 class TestUOpChildren(unittest.TestCase):
   def test_children_exist(self):
     a = UOp.variable("weird_name_234", 0, 10)

test/test_uops_stats.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
@@ -51,11 +51,7 @@ class TestMemoryCount(unittest.TestCase):
     a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     _, mem = get_stats(a+b)
-    if RANGEIFY:
-      # rangeify is smart!
-      self.assertEqual(mem, 1024 + 2*1024)  # 2 lil reads + 1 lil write
-    else:
-      self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
+    self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
 
   def test_self_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

tinygrad/codegen/opt/postrange.py

@@ -72,7 +72,7 @@ class Scheduler:
     for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
 
     # filter any not in reduces
-    # TODO: enable this
+    # TODO: reenable this
     """
     reduce_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.REDUCE]
     for ls in reduce_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])

tinygrad/codegen/simplify.py

@@ -29,7 +29,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
         # do the merge
         new_range = r0.replace(src=(s0*s1,))
         nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+                            name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
 
         # check if it simplifies
         if count_divmod(nidx) <= count_divmod(u):

tinygrad/schedule/rangeify.py

@@ -7,7 +7,6 @@ from tinygrad.uop.symbolic import sym, symbolic_simple
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup
 from tinygrad.schedule.kernelize import Kernel
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
-from tinygrad.codegen.simplify import pm_flatten_range
 
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
@@ -567,9 +566,6 @@ to_define_global = PatternMatcher([
   # this is only needed if you are using symbolic
   (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 
-  # remove RANGE with 0 size
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.const(dtypes.index, 0) if r.vmax == 0 else None),
-
   # renumber the ranges starting with 0 so that kernel deduping works
   (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
@@ -612,7 +608,7 @@ def split_store(ctx:list[UOp], x:UOp):
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
+  ret = graph_rewrite(x, to_define_global+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
 
   # gather the metadata
   metadatas = [ctx[y].metadata for y in lctx.parent_tags]

tinygrad/uop/ops.py

@@ -217,7 +217,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # determine what ranges this is in
   @functools.cached_property
-  def _ranges(self) -> dict[UOp, None]:
+  def ranges(self) -> dict[UOp, None]:
+    if self.op is Ops.RANGE: return {self:None}
     ret: dict[UOp, None] = {}
     if self.op in range_start.keys():
       for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
@@ -227,11 +228,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       for s in self.src: ret.update(s.ranges)
     return ret
 
-  @property
-  def ranges(self) -> dict[UOp, None]:
-    if self.op is Ops.RANGE: return {self:None}
-    return self._ranges
-
   # *** uop evaluation ***
 
   def simplify(self, tracked=False):