From ec405b919f1ad723d960f4dca90ca6420891874a Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Mon, 31 Mar 2025 14:03:38 +0800
Subject: [PATCH] Revert "Revert "do not block gc in UOp.toposort (#9623)"
 (#9624)" (#9639)

This reverts commit 7ef02d0e1c7d35b536d935ad9da6ba5d8b619ecf.
---
 tinygrad/ops.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 2257d598fe..ae63a3c467 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -241,6 +241,15 @@ class UOpMetaClass(type):
 buffers:weakref.WeakKeyDictionary[UOp, Buffer] = weakref.WeakKeyDictionary() # this maps BUFFER uops to their device Buffers
 all_metadata:weakref.WeakKeyDictionary[UOp, Metadata] = weakref.WeakKeyDictionary()
 
+def _toposort(u:UOp, cache:set[UOp]):
+  if u in cache: return {}
+  nodes: dict[UOp, None] = {}
+  # NOTE: this is a lot faster than the comprehension in parents
+  for parent in u.src: nodes.update(_toposort(parent, cache))
+  nodes[u] = None
+  cache.add(u)
+  return nodes
+
 # NOTE: this should be frozen, but frozen is slower
 @dataclass(eq=False, slots=True)
 class UOp(MathTrait, metaclass=UOpMetaClass):
@@ -271,14 +280,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @property
   def toposort(self) -> dict[UOp, None]:
-    def _toposort(u:UOp, cache:set[UOp]):
-      if u in cache: return {}
-      nodes: dict[UOp, None] = {}
-      # NOTE: this is a lot faster than the comprehension in parents
-      for parent in u.src: nodes.update(_toposort(parent, cache))
-      nodes[u] = None
-      cache.add(u)
-      return nodes
     return _toposort(self, cache=set())
 
   # returns map of UOps to their children in the graph rooted by self