now that needs rangeify 2 also

* fix bad range merges

* fix rng

* fix uop gc

test/external/external_uop_gc.py

@@ -63,6 +63,7 @@ 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,6 +123,7 @@ 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_rangeify.py

@@ -40,6 +40,11 @@ class TestRangeifyOpt(unittest.TestCase):
 
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
 class TestRangeify(unittest.TestCase):
+  def test_groupnorm(self):
+    # ranges 1 and 3 are merging
+    x = nn.GroupNorm(32, 128)
+    x(Tensor.empty(1, 128, 64, 64)).realize()
+
   def test_expand_children(self):
     A = Tensor.empty(N, N).sum(axis=1)
     ba = A.expand(N, N)

test/test_softmax_fusion.py

@@ -44,6 +44,7 @@ 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,86 +544,6 @@ 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
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
@@ -51,7 +51,11 @@ 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)
-    self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
+    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
 
   def test_self_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

tinygrad/codegen/opt/postrange.py

@@ -71,6 +71,13 @@ class Scheduler:
                         or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
     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
+    """
+    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])
+    """
+
     return [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE and x.arg[1] == AxisType.LOOP] if store_rngs else []
 
   def convert_loop_to_global(self):

tinygrad/codegen/simplify.py

@@ -17,20 +17,24 @@ pm_flatten_range = PatternMatcher([
 
 def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  reduce_ranges = [x.ranges for x in u.sparents if x.op is Ops.REDUCE]
   i = range_start[u.op]
   while i < len(u.src)-1:
     r0, r1 = u.src[i], u.src[i+1]
     # check same type
     if r0.arg[-1] == r1.arg[-1]:
-      s0, s1 = r0.src[0], r1.src[0]
-      # 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]}")
-      # check if it simplifies
-      if count_divmod(nidx) <= count_divmod(u):
-        u = nidx
-        continue
+      # check if the ranges to merge are in the same reduces
+      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
+        s0, s1 = r0.src[0], r1.src[0]
+        # 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]}")
+
+        # check if it simplifies
+        if count_divmod(nidx) <= count_divmod(u):
+          u = nidx
+          continue
     i += 1
   return u
 

tinygrad/schedule/rangeify.py

@@ -7,6 +7,7 @@ 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
@@ -566,6 +567,9 @@ 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),
 ])
@@ -608,7 +612,7 @@ def split_store(ctx:list[UOp], x:UOp):
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
+  ret = graph_rewrite(x, to_define_global+pm_flatten_range+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,8 +217,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # determine what ranges this is in
   @functools.cached_property
-  def ranges(self) -> dict[UOp, None]:
-    if self.op is Ops.RANGE: return {self:None}
+  def _ranges(self) -> dict[UOp, 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)
@@ -228,6 +227,11 @@ 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):