From ed90de6583fd3ce3e11f0fc0e1a12f90526955cf Mon Sep 17 00:00:00 2001
From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com>
Date: Sun, 28 Sep 2025 19:10:21 +0200
Subject: [PATCH] =?UTF-8?q?Revert=20"Bufferize=20early,=20fix=20"children?=
 =?UTF-8?q?=20not=20making=20progress"=20on=20big=20graphs=20(#1=E2=80=A6"?=
 =?UTF-8?q?=20(#12318)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This reverts commit 6f1cf717ded710aa8e2389c36731c83cfa42f50a.
---
 .github/workflows/test.yml    |  2 --
 test/test_schedule.py         |  2 +-
 tinygrad/schedule/rangeify.py | 21 +++------------------
 3 files changed, 4 insertions(+), 21 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 50e51b1dbf..23c075ca7d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -375,8 +375,6 @@ jobs:
       - name: Test openpilot model kernel count and gate usage
         run: |
           ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot model with rangeify
-        run: FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot alt model correctness (float32)
         run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot fastvits model correctness (float32)
diff --git a/test/test_schedule.py b/test/test_schedule.py
index 397136e147..3235d8b148 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -358,7 +358,7 @@ class TestSchedule(unittest.TestCase):
     out1 = r1 + y
     schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4)
     reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
-    assert len(reduceops) == 2
+    assert len(reduceops) == (3 if RANGEIFY else 2)
 
   def test_div_collapse_buffer(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py
index f0125528d3..743d17d54d 100644
--- a/tinygrad/schedule/rangeify.py
+++ b/tinygrad/schedule/rangeify.py
@@ -1,11 +1,10 @@
 from typing import Any, cast
 import functools, operator
 from dataclasses import dataclass, field
-from collections import defaultdict
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, graph_rewrite_map
 from tinygrad.uop.symbolic import sym, symbolic_simple
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup
 from tinygrad.schedule.multi import multi_pm
 
 from tinygrad.schedule.kernelize import Kernel
@@ -97,12 +96,9 @@ remove_contig_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.r
 # 2. mark all children
 
 @dataclass
-class ChildrenContext:
-  children: dict[UOp, list[UOp]]|None = None
-  realize_roots: defaultdict[UOp, list[UOp]] = field(default_factory=lambda: defaultdict(list))
+class ChildrenContext: children: dict[UOp, list[UOp]]|None = None
 def extract_children(ctx:ChildrenContext, x:UOp):
   if ctx.children is not None: return
-
   children_map = x.get_children_map()
   ctx.children = {}
   for k,v in children_map.items():
@@ -111,16 +107,6 @@ def extract_children(ctx:ChildrenContext, x:UOp):
     # NOTE: this gate shouldn't be here
     if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()) and any(x.op in {Ops.BUFFER, Ops.CONTIGUOUS} for x in k.toposort()):
       ctx.children[k] = non_sink_children
-  # if a node is in the toposort of multiple realizes, it will be indexed by different indices and we can bufferize early
-  # this prevents index_child: "children not making progress" error on big graphs
-  for r in [u for u in x.toposort() if u.op is Ops.REALIZE and (RANGEIFY<2 or u.arg is None)]:  # ignore partial realizes
-    for u in r.toposort(gate=lambda x: x is not Ops.REALIZE or (RANGEIFY>1 and u.arg is not None)):
-      ctx.realize_roots[u].append(r)
-
-def bufferize_early(ctx:ChildrenContext, x:UOp):
-  # this will also change the sources such that mark_children wont add the children/child uops anymore
-  new_srcs = [s.realize() if s in unwrap(ctx.children) and len(ctx.realize_roots[s])>1 else s for s in x.src]
-  return x.replace(src=tuple(new_srcs))
 
 def mark_children(ctx:ChildrenContext, x:UOp):
   assert ctx.children is not None
@@ -130,8 +116,7 @@ def mark_children(ctx:ChildrenContext, x:UOp):
 
 pm_children = PatternMatcher([
   (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), bufferize_early),
-  (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), mark_children),
+  (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN, Ops.SINK}, name="x"), mark_children),
 ])
 
 # *****************