add gc tests [pr] (#9718)

* add gc tests [pr]

* del

* more gc tests

* add NullGraph

.github/workflows/test.yml

@@ -337,6 +337,8 @@ jobs:
       run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
     - name: Run unit tests
       run: PYTHONPATH="." python -m pytest -n=auto test/unit/
+    - name: Run GC tests
+      run: PYTHONPATH="." python test/external/external_uop_gc.py
     - name: Repo line count < 12500 lines
       run: MAX_LINE_COUNT=12500 python sz.py
 

test/external/external_uop_gc.py

@@ -0,0 +1,72 @@
+import gc
+from tinygrad import Tensor, UOp, Device
+from tinygrad.shape.shapetracker import views_to_indexed_uops
+from tinygrad.engine.realize import method_cache, get_kernel
+
+def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
+def print_uops():
+  for x in gc.get_objects():
+    if isinstance(x, UOp): print(x)
+
+def start(): pass
+def single_tensor(): Tensor([2])
+def two_plus_two(): Tensor([2])+Tensor([2])
+def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
+def two_plus_two_kernel():
+  si = (Tensor([2])+Tensor([2])).schedule()[-1]
+  get_kernel(Device.default.renderer, si.ast)
+def two_plus_two_linearize():
+  si = (Tensor([2])+Tensor([2])).schedule()[-1]
+  k = get_kernel(Device.default.renderer, si.ast)
+  k.get_optimized_ast()
+  #k.linearize()
+def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
+def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
+def gradient_test():
+  x = Tensor.eye(3, requires_grad=True)
+  y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+  z = y.matmul(x).sum()
+  z.backward()
+def realized_eye():
+  Tensor.eye(3, requires_grad=True).realize()
+def realized_list():
+  Tensor([[2.0,0,-2.0]], requires_grad=True).realize()
+def kernel_matmul():
+  x = Tensor.eye(3, requires_grad=True)
+  y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+  z = y.matmul(x)
+  si = z.schedule()[-1]
+  get_kernel(Device.default.renderer, si.ast)
+def realized_matmul():
+  x = Tensor.eye(3, requires_grad=True)
+  y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+  z = y.matmul(x)
+  Tensor.realize(z)
+def realized_gradient():
+  x = Tensor.eye(3, requires_grad=True)
+  y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+  z = y.matmul(x).sum()
+  z.backward()
+  Tensor.realize(x, y, z, x.grad, y.grad)
+tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
+         two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
+         realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient]
+
+if __name__ == "__main__":
+  gc.disable()
+  start_uops = uops_allocated()
+  # there's a few consts created as default values
+  print_uops()
+  for t in tests:
+    t()
+
+    # these caches will keep uops alive
+    method_cache.clear()
+    views_to_indexed_uops.cache_clear()
+
+    new_uops = uops_allocated()
+    gc.collect()
+    new_uops_gc = uops_allocated()
+    print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
+    assert new_uops == start_uops
+  #print_uops()

tinygrad/codegen/kernel.py

@@ -577,7 +577,7 @@ class Kernel:
   def get_optimized_ast(self, name_override:Optional[str]=None) -> UOp:
     @functools.cache
     def fixup_ast(op:UOp) -> UOp:
-      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src))
+      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
       if op.op in GroupOp.Buffer and op in self.bufs:
         st_uop = self.sts[self.bufs.index(op)].to_uop()
         # NOTE: if CONST got masked after applying opts, we create a new VALID
@@ -652,8 +652,9 @@ class Kernel:
           return UOp(Ops.LOAD, op.dtype, (local_buffer, st_uop, UOp.store(local_buffer, st_uop, grouped_reduce)))
 
       return ret
-
-    return graph_rewrite(fixup_ast(self.ast), view_left)
+    fixed_ast = fixup_ast(self.ast)
+    del fixup_ast
+    return graph_rewrite(fixed_ast, view_left)
 
   # **** this is the lowerer ****
 

tinygrad/gradient.py

@@ -48,15 +48,17 @@ pm_gradient = PatternMatcher([
 # copied from tensor.py, get relevant toposort of gradients
 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
   @functools.cache
-  def is_in_target_path(x:UOp) -> bool: return any(u in targets or is_in_target_path(u) for u in x.src)
+  def is_in_target_path(x:UOp) -> bool: return any(u in targets or is_in_target_path(u) for u in x.src) # noqa: F821
   def _walk(node:UOp, visited:set[UOp]) -> Iterator[UOp]:
     visited.add(node)
     if node.op is Ops.DETACH: return
-    if is_in_target_path(node):
+    if is_in_target_path(node): # noqa: F821
       for i in node.src:
-        if i not in visited: yield from _walk(i, visited)
+        if i not in visited: yield from _walk(i, visited) # noqa: F821
       yield node
-  return list(_walk(root, set()))
+  ret = list(_walk(root, set()))
+  del is_in_target_path, _walk
+  return ret
 
 def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp]:
   grads = {root: root_grad}

tinygrad/runtime/ops_null.py

@@ -1,4 +1,5 @@
 from tinygrad.device import Compiled, Compiler, Renderer, Allocator
+from tinygrad.engine.jit import MultiGraphRunner
 
 class NullRenderer(Renderer):
   def render(self, uops:list) -> str: return ""
@@ -13,5 +14,8 @@ class NullAllocator(Allocator):
   def _copyin(self, dest, src:memoryview): pass
   def _copyout(self, dest:memoryview, src): pass
 
+class NullGraph(MultiGraphRunner):
+  def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3
+
 class NullDevice(Compiled):
-  def __init__(self, device:str): super().__init__(device, NullAllocator(), NullRenderer(), Compiler(), NullProgram)
+  def __init__(self, device:str): super().__init__(device, NullAllocator(), NullRenderer(), Compiler(), NullProgram, NullGraph)