From 6fb74beb7951596839283db22a28f392632b9281 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Thu, 9 Oct 2025 14:14:25 +0800
Subject: [PATCH] remove on_stack from graph_rewrite

---
 tinygrad/uop/ops.py | 43 +++++++++++++------------------------------
 1 file changed, 13 insertions(+), 30 deletions(-)

diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 20b0b8bb82..b25f3b1faa 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -1009,7 +1009,6 @@ class RewriteContext:
 
   def unified_rewrite(self, root:UOp) -> UOp:
     stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
-    on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
     REWRITE_STACK_LIMIT = getenv("REWRITE_STACK_LIMIT", 250000)
     while stack:
       if len(stack) > REWRITE_STACK_LIMIT: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
@@ -1031,39 +1030,23 @@ class RewriteContext:
             self.replace[n] = new_n
             continue
         stack.append((n, 1, new_n))
-        for x in reversed(new_n.src):
-          if x in on_stack: continue
-          stack.append((x, 0, x))
-          on_stack.add(x)
+        for x in reversed(new_n.src): stack.append((x, 0, x))
       elif stage == 1:
-        tmp = []
-        for x in new_n.src:
-          if (rx:=self.replace.get(x, SENTINEL)) is SENTINEL:
-            # if some new sources aren't ready, we try this again later. happens with on_stack, maybe should remove?
-            stack.appendleft((n, 1, new_n))
-            break
-          tmp.append(rx)
+        # in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
+        if (new_src:=tuple([self.replace[x] for x in new_n.src])) == new_n.src:
+          # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
+          if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
+            self.replace[n] = new_n
+            continue
         else:
-          # in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
-          if (new_src:=tuple(tmp)) == new_n.src:
-            # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
-            if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
-              self.replace[n] = new_n
-              continue
-          else:
-            # if srcs changed from rewrites, construct a new UOp with the new srcs
-            new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag)
-          # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
-          stack.append((n, 2, new_src_n))
-          stack.append((new_src_n, 0, new_src_n))
+          # if srcs changed from rewrites, construct a new UOp with the new srcs
+          new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag)
+        # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
+        stack.append((n, 2, new_src_n))
+        stack.append((new_src_n, 0, new_src_n))
       else:
         # in stage 2, we link the result of new_n to the result of n
-        if (replaced_new_n:=self.replace.get(new_n, SENTINEL)) is SENTINEL:
-          # not ready, try the link later
-          stack.appendleft((n, 2, new_n))
-        else:
-          # otherwise we are done
-          self.replace[n] = replaced_new_n
+        self.replace[n] = self.replace[new_n]
     return self.replace[root]
 
 @track_matches