Merge branch 'master' into fix_range_merging_sd

that
less topo
2025-09-30 18:50:13 +08:00 · 2025-09-30 18:50:04 +08:00 · 2025-09-30 17:47:33 +08:00 · 2025-09-30 17:35:25 +08:00 · 2025-09-30 17:28:23 +08:00 · 2025-09-30 17:05:32 +08:00
3 changed files with 25 additions and 9 deletions
@@ -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)
@@ -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):
@@ -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
Author	SHA1	Message	Date
George HotzandGitHub	c3f4ab1902	Merge branch 'master' into fix_range_merging_sd	2025-09-30 18:50:13 +08:00
geohot	c963e44ea0	that	2025-09-30 18:50:04 +08:00
geohot	036803f8e7	less topo	2025-09-30 17:47:33 +08:00
geohot	e73f0bbf98	no repeating work with globals	2025-09-30 17:35:25 +08:00
geohot	348188a0b5	ahh, that really fixes it	2025-09-30 17:28:23 +08:00
geohot	8904b6b7d7	fix range merging for stable diffusion	2025-09-30 17:05:32 +08:00