From 1d2165182392df111f40f856fce5e42250567f6f Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sun, 15 Dec 2024 19:50:42 +0200
Subject: [PATCH] free toposort cache after it goes out of scope [pr] (#8264)

---
 test/test_gc.py | 1 -
 tinygrad/ops.py | 9 +++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/test/test_gc.py b/test/test_gc.py
index 94953cfc3831..010f59039f56 100644
--- a/test/test_gc.py
+++ b/test/test_gc.py
@@ -66,7 +66,6 @@ def test_schedule_gc_with_inputs(self):
     del y
     self.assertEqual(bufs_allocated()-init, 0)
 
-  @unittest.expectedFailure
   def test_toposort_blocks_gc(self):
     init = bufs_allocated()
     x = Tensor.ones(4,4).contiguous().realize()+1
diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 1974e9a09628..247ad6ea2901 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -252,14 +252,15 @@ def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is Ops.
 
   @property
   def toposort(self) -> Dict[UOp, None]:
-    @functools.lru_cache(None)
-    def _toposort(u:UOp):
+    def _toposort(u:UOp, cache:Dict[UOp, Dict[UOp, None]]):
+      if (cret:=cache.get(u)) is not None: return cret
       nodes: Dict[UOp, None] = {}
       # NOTE: this is a lot faster than the comprehension in parents
-      for parent in u.src: nodes.update(_toposort(parent))
+      for parent in u.src: nodes.update(_toposort(parent, cache))
       nodes[u] = None
+      cache[u] = nodes
       return nodes
-    return _toposort(self)
+    return _toposort(self, cache={})
 
   @functools.cached_property
   def tuplize(self:UOp) -> Tuple[int, Any, Optional[DType], Tuple]: