free toposort cache after it goes out of scope [pr] (tinygrad#8264)

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

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]: