remove on_stack from graph_rewrite

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