Merge branch 'master' into fix_range_merging_sd

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)

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: 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])
+    """
+
     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