From 1b1752b6f8c0685ca2bf8753a52a24622334918d Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Wed, 8 Oct 2025 21:11:36 +0800
Subject: [PATCH] Revert "remove the on_stack stuff and any retries"

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

diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index c535dc6c37..2622997923 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -1019,6 +1019,7 @@ 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)")
@@ -1040,23 +1041,39 @@ class RewriteContext:
             self.replace[n] = new_n
             continue
         stack.append((n, 1, new_n))
-        for x in reversed(new_n.src): stack.append((x, 0, x))
+        for x in reversed(new_n.src):
+          if x in on_stack: continue
+          stack.append((x, 0, x))
+          on_stack.add(x)
       elif stage == 1:
-        # 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
+        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)
         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))
+          # 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))
       else:
         # in stage 2, we link the result of new_n to the result of n
-        self.replace[n] = self.replace[new_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
     return self.replace[root]
 
 @track_matches