add cache collector (#1595)

* init cache collector

* add test_cache_collector.py

* switch GlobalCounters.cache to CacheCollector

* init jit models test

* jitted SD

* add debug msg to print loaded bufs count

* moved cache collctor to jit

* clearer SD

* no double device import

docs/abstractions.py

@@ -299,11 +299,10 @@ result = Tensor(2) + Tensor(3)
 
 # we have a global cache used by the JIT
 # from there, we can see the generated clang code
-from tinygrad.helpers import GlobalCounters
-GlobalCounters.cache = []    # enables the cache
+from tinygrad.jit import CacheCollector
+CacheCollector.start()       # enables the cache
 result.realize()             # create the program and runs it
-cache_saved = GlobalCounters.cache
-GlobalCounters.cache = None  # disable the cache
+cache_saved = CacheCollector.finish()  # disable the cache
 
 # there's one ASTRunner in the cache
 assert len(cache_saved) == 1

examples/benchmark_train_efficientnet.py

@@ -8,6 +8,7 @@ from tinygrad.nn import optim
 from tinygrad.tensor import Tensor
 from tinygrad.ops import GlobalCounters
 from tinygrad.helpers import getenv
+from tinygrad.jit import CacheCollector
 
 def tensors_allocated():
   return sum(isinstance(x, Tensor) for x in gc.get_objects())
@@ -41,7 +42,7 @@ if __name__ == "__main__":
       st = time.monotonic()
       out = model.forward(x_train)
       loss = out.log_softmax().mul(y_train).mean()
-      if i == 2 and CLCACHE: GlobalCounters.cache = []
+      if i == 2 and CLCACHE: CacheCollector.start()
       if BACKWARD:
         optimizer.zero_grad()
         loss.backward()
@@ -57,8 +58,7 @@ if __name__ == "__main__":
       et = time.monotonic()
 
     if i == 2 and CLCACHE:
-      cl_cache = GlobalCounters.cache
-      GlobalCounters.cache = None
+      cl_cache = CacheCollector.finish()
 
     mem_used = GlobalCounters.mem_used
     loss_cpu = loss.detach().numpy()

examples/stable_diffusion.py

@@ -13,6 +13,7 @@ from tinygrad.helpers import dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, Linear, GroupNorm, LayerNorm, Embedding
 from extra.utils import download_file
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
+from tinygrad.jit import TinyJit
 
 class AttnBlock:
   def __init__(self, in_channels):
@@ -621,6 +622,15 @@ if __name__ == "__main__":
     x_prev = math.sqrt(a_prev) * pred_x0 + dir_xt #+ noise
     return x_prev, pred_x0
 
+  @TinyJit
+  def do_step(latent, timestep):
+    e_t = get_model_output(latent, timestep)
+    x_prev, _ = get_x_prev_and_pred_x0(latent, e_t, index)
+    #e_t_next = get_model_output(x_prev)
+    #e_t_prime = (e_t + e_t_next) / 2
+    #x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t_prime, index)
+    return x_prev.realize()
+
   # start with random noise
   latent = Tensor.randn(1,4,64,64)
 
@@ -628,13 +638,7 @@ if __name__ == "__main__":
   for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
     GlobalCounters.reset()
     t.set_description("%3d %3d" % (index, timestep))
-    e_t = get_model_output(latent, Tensor([timestep]))
-    x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t, index)
-    #e_t_next = get_model_output(x_prev)
-    #e_t_prime = (e_t + e_t_next) / 2
-    #x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t_prime, index)
-    latent = x_prev
-    latent.realize()
+    latent = do_step(latent, Tensor([timestep]))
 
   # upsample latent space to image with autoencoder
   x = model.first_stage_model.post_quant_conv(1/0.18215 * latent)

extra/dist/world.py

@@ -1,10 +1,11 @@
 from typing import Any, Optional, Tuple
 from extra import dist
 from multiprocessing import shared_memory
-from tinygrad.helpers import DEBUG, GlobalCounters, colored
+from tinygrad.helpers import DEBUG, colored
 from tinygrad.lazy import LazyBuffer
 from tinygrad.runtime.lib import RawBufferCopyIn, RawBufferCopyInOut
 from tinygrad.runtime.ops_shm import RawShmBuffer
+from tinygrad.jit import CacheCollector
 from tinygrad.tensor import Tensor, Function
 import numpy as np
 
@@ -35,7 +36,7 @@ def _send_rb(x:RawBufferCopyInOut, target_rank:int, cache_id:Optional[str]=None)
   __send_rb((x, rb, target_rank, (shm_name, cache_id)))
 
   # jit support
-  if GlobalCounters.cache is not None: GlobalCounters.cache.append((__send_rb, [x, rb, target_rank, None], {}))
+  CacheCollector.add(__send_rb, [x, rb, target_rank, None], {})
 setattr(_send_rb, "shared_memory_cache", {})
 
 # receive a rawbuffer from the target rank
@@ -52,7 +53,7 @@ def _recv_rb(x:RawBufferCopyIn, target_rank:int):
     s.unlink()
 
   # jit support
-  if GlobalCounters.cache is not None: GlobalCounters.cache.append((__recv_rb, [x, rb, target_rank], {}))
+  CacheCollector.add(__recv_rb, [x, rb, target_rank], {})
 
 # sends a lazybuffer from our rank to the target rank
 def _send_lb(x:LazyBuffer, target_rank:int, cache_id:Optional[str]=None) -> None: _send_rb(x.contiguous().realize().realized, target_rank, cache_id=cache_id)

extra/thneed.py

@@ -5,7 +5,7 @@ import json
 import traceback
 import numpy as np
 from tinygrad.runtime.ops_gpu import CLProgram
-from tinygrad.helpers import prod, getenv
+from tinygrad.helpers import DEBUG, getenv
 from collections import defaultdict
 import pyopencl as cl
 from tinygrad.runtime.ops_gpu import CL, OSX_TIMING_RATIO
@@ -140,6 +140,8 @@ class Thneed:
         aaa.append(aa)
       self.cl_cache.append((kernel, [k['global_work_size'], k['local_work_size'], *aaa]))
 
+    if DEBUG >= 1: print(f"thneed: total bufs loaded: {len(bufs.keys())}")
+
     # load inputs
     for k in jdat['inputs']:
       self.inputs[k['name']] = bufs[k['buffer_id']]

openpilot/compile.py

@@ -18,6 +18,7 @@ import numpy as np
 
 import tinygrad.graph as graph
 from tinygrad.ops import GlobalCounters
+from tinygrad.jit import TinyJit, CacheCollector
 
 import pyopencl as cl
 from tinygrad.runtime.ops_gpu import CL
@@ -34,13 +35,11 @@ def get_random_input_tensors(input_shapes):
   np_inputs = {k:v.realize().numpy() for k,v in inputs.items()}
   return inputs, np_inputs
 
-from tinygrad.jit import TinyJit
-
 @TinyJit
 def model_exec(run_onnx, using_graph, **inputs):
   ret = next(iter(run_onnx(inputs).values())).cast(dtypes.float32)
   GlobalCounters.reset()
-  GlobalCounters.cache = []  # don't cache pre-realize
+  CacheCollector.start() # don't cache pre-realize
   if using_graph: graph.GRAPH = True
   print("realizing")
   return ret.realize()

test/external/external_copy_benchmark.py

@@ -3,15 +3,16 @@ from tinygrad.helpers import prod
 from tinygrad.ops import Device
 from tinygrad.tensor import Tensor
 from tinygrad.ops import GlobalCounters
+from tinygrad.jit import CacheCollector
 
 class TestCopy(unittest.TestCase):
   def test_add1(self):
     pts = []
     for i in range(16384, 16384*256, 16384):
       t = Tensor.randn(i).realize()
-      GlobalCounters.cache = []
+      CacheCollector.start()
       t.assign(t+1).realize()
-      fxn, args, _ = GlobalCounters.cache[0]
+      fxn, args, _ = CacheCollector.finish()[0]
       GlobalCounters.reset()
       def run(): return fxn(args, force_wait=True)
       ct = min([run() for _ in range(10)])

test/external/external_test_jit_on_models.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.state import get_parameters
+from tinygrad.ops import LazyOp, LoadOps
+from tinygrad.jit import TinyJit, JIT_SUPPORTED_DEVICE
+from tinygrad.helpers import dtypes, CI
+from tinygrad.lazy import Device
+
+from examples.llama import Transformer
+
+# for speed
+def derandomize(x):
+  if isinstance(x, LazyOp):
+    if x.op == LoadOps.RAND: x.op = LoadOps.EMPTY
+    x.src = [derandomize(s) for s in x.src]
+  else:
+    x.op = derandomize(x.op)
+  return x
+
+def derandomize_model(model):
+  for p in get_parameters(model):
+    p.lazydata = derandomize(p.lazydata)
+    p.realize()
+
+def helper_test_jitted_correctness(gen, train, train_jit):
+  nojit = train(*gen()).numpy()
+  for _ in range(5): jit = train_jit(*gen()).numpy()
+  np.testing.assert_allclose(nojit, jit, rtol=1e-3, atol=1e-5)
+
+@unittest.skipUnless(Device.DEFAULT in JIT_SUPPORTED_DEVICE, "needs JIT")
+class TestJittedModels(unittest.TestCase):
+  def test_jitted_tiny_llama(self):
+    old_type = Tensor.default_type
+    Tensor.default_type = dtypes.float16
+
+    args_tiny = {"dim": 1024, "multiple_of": 256, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
+    model = Transformer(**args_tiny)
+    derandomize_model(model)
+    def test(t): return model(t, 0).realize()
+
+    @TinyJit
+    def test_jit(t): return model(t, 0).realize()
+    helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
+    Tensor.default_type = old_type
+
+  @unittest.skipUnless(not CI, "huge for CI")
+  def test_jitted_stable_diffusion(self):
+    from examples.stable_diffusion import UNetModel
+    model = UNetModel()
+    derandomize_model(model)
+    def test(t, t2): return model(t, 801, t2).realize()
+
+    @TinyJit
+    def test_jit(t, t2): return model(t, 801, t2).realize()
+    helper_test_jitted_correctness(lambda: (Tensor.randn(1, 4, 16, 16),Tensor.randn(1, 77, 768)), test, test_jit)
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_test_opt.py

@@ -15,6 +15,7 @@ from tinygrad.helpers import getenv
 from tinygrad.nn import optim
 from tinygrad.ops import GlobalCounters, MovementOps, ReduceOps
 from tinygrad.lazy import PUSH_PERMUTES
+from tinygrad.jit import CacheCollector
 
 class CLCache():
   def __init__(self, allowed=None, strict=False, preclear=True): self.allowed, self.strict, self.preclear = allowed, strict, preclear
@@ -24,13 +25,13 @@ class CLCache():
       for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
         x.realize()
       GlobalCounters.reset()
-    GlobalCounters.cache = []
+    CacheCollector.start()
     print("cache: entering")
   def __exit__(self, type, value, traceback):
-    print(f"cache: exiting with size {len(GlobalCounters.cache)}", f"allowed {self.allowed}" if self.allowed is not None else "")
+    cache = CacheCollector.finish()
+    print(f"cache: exiting with size {len(cache)}", f"allowed {self.allowed}" if self.allowed is not None else "")
     if self.allowed is not None:
-      assert len(GlobalCounters.cache) <= self.allowed and (not self.strict or len(GlobalCounters.cache) == self.allowed), f"used too many kernels! {len(GlobalCounters.cache)} > {self.allowed}"
-    GlobalCounters.cache = None
+      assert len(cache) <= self.allowed and (not self.strict or len(cache) == self.allowed), f"used too many kernels! {len(cache)} > {self.allowed}"
 
 from models.convnext import ConvNeXt
 from models.efficientnet import EfficientNet
@@ -79,7 +80,7 @@ class TestInferenceMinKernels(unittest.TestCase):
     img = Tensor.randn(1, 3, 224, 224)
     with CLCache(223): # NOTE: this is way too high
       out = model.forward(img)
-      assert len(GlobalCounters.cache) == 0, "ViT prerealized?"
+      assert len(CacheCollector.cache) == 0, "ViT prerealized?"
       out.realize()
 
   def test_llama(self):
@@ -100,7 +101,7 @@ class TestOptBinOp(unittest.TestCase):
       if f2 is not None: d = f2(a, b)
       c.realize()
       if f2 is not None: d.realize()
-      assert len(GlobalCounters.cache) == allowed, "binop was rerun!"
+      assert len(CacheCollector.cache) == allowed, "binop was rerun!"
     if f2 is not None: np.testing.assert_allclose(c.numpy().ravel(), d.numpy().ravel(), rtol=1e-3, atol=1e-5)
 
   def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1))
@@ -124,7 +125,7 @@ class TestOptReduceLoop(unittest.TestCase):
       b = t.reshape(16,1).expand(16,16).sum(0)
       c = (t+b)
       c.realize()
-      assert len(GlobalCounters.cache) == 2, "loop left fusion broken"
+      assert len(CacheCollector.cache) == 2, "loop left fusion broken"
 
   def test_loop_right(self):
     a = Tensor.randn(16, 16)
@@ -134,7 +135,7 @@ class TestOptReduceLoop(unittest.TestCase):
       b = t.reshape(16,1).expand(16,16).sum(0)
       c = (b+t)
       c.realize()
-      assert len(GlobalCounters.cache) == 2, "loop right fusion broken"
+      assert len(CacheCollector.cache) == 2, "loop right fusion broken"
 
 @unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
 class TestOptWChild(unittest.TestCase):
@@ -146,7 +147,7 @@ class TestOptWChild(unittest.TestCase):
       d = c+1
       e = c+2
       d.realize()
-      assert len(GlobalCounters.cache) == 2, "don't fuse if you have children"
+      assert len(CacheCollector.cache) == 2, "don't fuse if you have children"
 
 @unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
 class TestOpt(unittest.TestCase):
@@ -155,7 +156,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       d = a * b + c
       d.realize()
-      assert len(GlobalCounters.cache) == 1, "optimizer didn't fold muladd"
+      assert len(CacheCollector.cache) == 1, "optimizer didn't fold muladd"
     np.testing.assert_allclose(d.numpy(), np.ones((2,2))*2, rtol=1e-5)
 
   def test_fold_reduce_elementwise(self):
@@ -164,7 +165,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       ret = img.sum() + addme
       ret.realize()
-      assert len(GlobalCounters.cache) == 1, "optimizer didn't fold reduce/elementwise"
+      assert len(CacheCollector.cache) == 1, "optimizer didn't fold reduce/elementwise"
     assert ret.numpy()[0] == 33
 
   def test_fold_batchnorm(self):
@@ -175,7 +176,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_bn = bn(img).realize()
       print(img_bn)
-      assert len(GlobalCounters.cache) == 3, f"optimizer didn't fold batchnorm, got {len(GlobalCounters.cache)}"
+      assert len(CacheCollector.cache) == 3, f"optimizer didn't fold batchnorm, got {len(CacheCollector.cache)}"
     Tensor.training = False
 
   def test_fold_conv_sgd(self):
@@ -191,7 +192,7 @@ class TestOpt(unittest.TestCase):
       # TODO: this should be 4, but the sum output child stays around
       # with pushing_permutes it can be 3
       # TODO: broken with optim fixes
-      assert len(GlobalCounters.cache) in [4,5,6], f"optimizer didn't fold conv-backward SGD, got {len(GlobalCounters.cache)}"
+      assert len(CacheCollector.cache) in [4,5,6], f"optimizer didn't fold conv-backward SGD, got {len(CacheCollector.cache)}"
     Tensor.training = False
 
   def test_fold_2convs_sgd(self):
@@ -244,7 +245,7 @@ class TestOpt(unittest.TestCase):
     img_conv = bn(c1(img)).relu().realize()
     with CLCache():
       img_conv = bn(c1(img)).relu().realize()
-      assert len(GlobalCounters.cache) == 1, f"optimizer didn't fold conv-batchnorm at test time, got {len(GlobalCounters.cache)}"
+      assert len(CacheCollector.cache) == 1, f"optimizer didn't fold conv-batchnorm at test time, got {len(CacheCollector.cache)}"
 
   def test_fold_conv_batchnorm(self):
     Tensor.training = True
@@ -254,7 +255,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_conv = bn(c1(img)).relu().realize()
       print(img_conv)
-      assert len(GlobalCounters.cache) == 4, f"optimizer didn't fold conv-batchnorm, got {len(GlobalCounters.cache)}"
+      assert len(CacheCollector.cache) == 4, f"optimizer didn't fold conv-batchnorm, got {len(CacheCollector.cache)}"
     Tensor.training = False
 
   def test_fold_conv_elu(self):
@@ -264,7 +265,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_conv = img.sequential([c1, Tensor.elu, c2, Tensor.elu]).realize()
       print(img_conv)
-      assert len(GlobalCounters.cache) == 2, "optimizer didn't fold conv/elu"
+      assert len(CacheCollector.cache) == 2, "optimizer didn't fold conv/elu"
 
   def test_fold_conv_relu(self):
     img = Tensor.ones(1,4,8,8)
@@ -273,7 +274,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_conv = img.sequential([c1, Tensor.relu, c2, Tensor.relu]).realize()
       print(img_conv)
-      assert len(GlobalCounters.cache) == 2, "optimizer didn't fold conv/relu"
+      assert len(CacheCollector.cache) == 2, "optimizer didn't fold conv/relu"
 
   def test_fold_conv_relu_nobias(self):
     img = Tensor.ones(1,4,8,8)
@@ -282,7 +283,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_conv = img.sequential([c1, Tensor.relu, c2, Tensor.relu]).realize()
       print(img_conv)
-      assert len(GlobalCounters.cache) == 2, "optimizer didn't fold conv/relu"
+      assert len(CacheCollector.cache) == 2, "optimizer didn't fold conv/relu"
 
   def test_permute_was_pushed(self):
     a = Tensor.randn(16, 16, 16)
@@ -290,7 +291,7 @@ class TestOpt(unittest.TestCase):
       c = a.sum(2)
       d = c.permute(1,0).contiguous()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
     if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
@@ -300,7 +301,7 @@ class TestOpt(unittest.TestCase):
       c = a.sum(-1)
       d = c.reshape(16,16).permute(1,0).contiguous()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,16).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
     if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
@@ -310,7 +311,7 @@ class TestOpt(unittest.TestCase):
       c = a.sum(-1)
       d = c.reshape(16,1,16).permute(2,1,0).contiguous()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,1,16).transpose(2,1,0), d.numpy(), rtol=1e-3, atol=1e-5)
     if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
@@ -323,7 +324,7 @@ class TestOpt(unittest.TestCase):
       c = a.sum(2)
       d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0).reshape(4,4,4,4), d.numpy(), rtol=1e-3, atol=1e-5)
     if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
@@ -335,7 +336,7 @@ class TestOpt(unittest.TestCase):
       d = a.sum(2).permute(1,0)
       c.realize()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(c.numpy().transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
     assert cache_len == 1, "reduceop was rerun!"
 
@@ -347,7 +348,7 @@ class TestOpt(unittest.TestCase):
       d = a.sum(2)
       c.realize()
       d.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(c.numpy(), d.numpy().transpose(1,0), rtol=1e-3, atol=1e-5)
     assert cache_len == 1, "reduceop was rerun!"
 
@@ -357,7 +358,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       c = (a.sum(2).contiguous() + b).contiguous()
       c.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     assert cache_len == 1, "contiguous wasn't folded"
 
   def _test_fold_expand_reduce_helper(self, n, m, axis, allowed):
@@ -365,7 +366,7 @@ class TestOpt(unittest.TestCase):
     with CLCache(allowed=allowed):
       a = Tensor.ones(n, m).sum(axis).reshape(n, 1).expand(n, m).sum(axis)
       a.realize()
-      cache_len = len(GlobalCounters.cache)
+      cache_len = len(CacheCollector.cache)
     np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
     return cache_len
 
@@ -376,7 +377,7 @@ class TestOpt(unittest.TestCase):
         with CLCache(allowed=2):
           a = Tensor.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
           a.realize()
-          cache_len = len(GlobalCounters.cache)
+          cache_len = len(CacheCollector.cache)
         np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
         return cache_len
 
@@ -387,7 +388,7 @@ class TestOpt(unittest.TestCase):
       with CLCache(allowed=3):
         a = Tensor.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
         a.realize()
-        cache_len = len(GlobalCounters.cache)
+        cache_len = len(CacheCollector.cache)
       np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
       return cache_len
 

test/test_cache_collector.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python
+import unittest
+from tinygrad.runtime.lib import RawBuffer, LRUAllocator
+from tinygrad.helpers import dtypes
+from tinygrad.jit import CacheCollector
+from weakref import ref
+
+class FakeDeviceBuffer():
+  def __init__(self, sz, dt, device):
+    self.size = sz
+    self.dtype = dt
+    self.device = device
+
+class FakeAllocator(LRUAllocator):
+  def _do_alloc(self, size, dtype, device, **kwargs): return FakeDeviceBuffer(size, dtype, device)
+
+FAKE_GLOBAL_ALLOCATOR = None
+class FakeBuffer(RawBuffer):
+  def __init__(self, size, dtype, device='0'):
+    global FAKE_GLOBAL_ALLOCATOR
+    super().__init__(size, dtype, allocator=FAKE_GLOBAL_ALLOCATOR, **{'device': device})
+    assert self._buf.size == size and self._buf.dtype == dtype and self._buf.device == device, "This allocator requires 100% match of dtype and size."
+
+def alloc(allocator, size, dtype, **kwargs):
+  global FAKE_GLOBAL_ALLOCATOR
+  FAKE_GLOBAL_ALLOCATOR = allocator
+  buf = FakeBuffer(size, dtype, **kwargs)
+  assert buf.dtype == dtype and buf.size == size
+  FAKE_GLOBAL_ALLOCATOR = None
+  return buf
+
+def anybuf(size, dtype):
+  return FakeBuffer(size, dtype)
+
+def add_to_cache(bufs):
+  CacheCollector.add(None, bufs, None)
+  return bufs[0]
+
+def add_to_cache_refed(bufs):
+  CacheCollector.add(None, bufs, None)
+  return bufs[0], [ref(buf) for buf in bufs]
+
+def get_bufs_count(cache):
+  ss = set()
+  for (_,bufs,_) in cache:
+    for buf in bufs: ss.add(buf)
+  return len(ss)
+
+class TestCacheCollector(unittest.TestCase):
+  def test_cache_collector_optimization(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32) for _ in range(2)]
+    CacheCollector.start()
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), inps[0]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out, inps[1]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    assert cache[-1][1][0] == out, "Output does not match."
+    assert get_bufs_count(cache) == 4, "Should have 4 buffers in total"
+    assert cache[-1][1][0] == cache[0][1][0], "Should reuse final output buffer as output in 1st kernel"
+    FAKE_GLOBAL_ALLOCATOR = None
+
+  def test_cache_collector_cycle_avoidance(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32) for _ in range(2)]
+    CacheCollector.start()
+    # Output buffer here cannot be shared with final output buffer, since we could get a cycle the next step as inps[1] has the same shape and dtype.
+    out = add_to_cache([FakeBuffer(64, dtypes.float32), inps[0]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out, inps[1]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    out = add_to_cache([FakeBuffer(64, dtypes.float32), out])
+    out = add_to_cache([FakeBuffer(64, dtypes.float32), out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    assert cache[-1][1][0] == out, "Output does not match."
+    assert get_bufs_count(cache) == 6, "Should have 6 buffers in total"
+    assert cache[-1][1][0] != cache[0][1][0] and cache[0][1][0] == cache[3][1][0], "Output buffers from 1st and 4th kernel could not be the same as the 5th."
+    FAKE_GLOBAL_ALLOCATOR = None
+
+  def test_cache_collector_all_alive(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32) for _ in range(2)]
+    outs = [FakeBuffer(128, dtypes.float32) for _ in range(4)]
+    CacheCollector.start()
+    out = add_to_cache([outs[0], inps[0]])
+    out = add_to_cache([outs[1], out, inps[1]])
+    out = add_to_cache([outs[2], out])
+    out = add_to_cache([outs[3], out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    assert cache[0][1][0] == outs[0], "Output0 should be on its place."
+    assert cache[1][1][0] == outs[1], "Output1 should be on its place."
+    assert cache[2][1][0] == outs[2], "Output2 should be on its place."
+    assert cache[3][1][0] == outs[3], "Output3 should be on its place."
+    assert cache[-1][1][0] == out, "Output does not match."
+    assert get_bufs_count(cache) == len(outs) + len(inps), "Nothing to optimize, since buffers are alive and might be used as outputs"
+    FAKE_GLOBAL_ALLOCATOR = None
+
+  def test_cache_collector_middle_input(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32) for _ in range(2)]
+    outs = [FakeBuffer(32, dtypes.float32) for _ in range(1)]
+    CacheCollector.start()
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), inps[0]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out, inps[1]])
+    out,refs2 = add_to_cache_refed([outs[0], out, FakeBuffer(32, dtypes.float32)])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    assert cache[2][1][2] == refs2[2](), "Input should be captured."
+    assert cache[0][1][0] != cache[2][1][2], "None of outputs buffer should reuse new_input."
+    assert cache[1][1][0] != cache[2][1][2], "None of outputs buffer should reuse new_input."
+    assert cache[3][1][0] != cache[2][1][2], "None of outputs buffer should reuse new_input."
+    assert cache[4][1][0] != cache[2][1][2], "None of outputs buffer should reuse new_input."
+    assert cache[-1][1][0] == out, "Output does not match."
+    assert get_bufs_count(cache) == 7
+    FAKE_GLOBAL_ALLOCATOR = None
+
+  def test_cache_collector_multidev(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32, '1') for _ in range(2)]
+    CacheCollector.start()
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '1'), inps[0]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '1'), out, inps[1]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '1'), out])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '2'), out])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '2'), out])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32, '2'), out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    for i in range(3):
+      assert cache[i][1][0]._device == '1', f"Device does not match {i}, has {cache[i][1][0]._device}."
+    for i in range(3, 6):
+      assert cache[i][1][0]._device == '2', f"Device does not match {i}, has {cache[i][1][0]._device}."
+    assert get_bufs_count(cache) == 6
+    FAKE_GLOBAL_ALLOCATOR = None
+
+  def test_cache_collector_anybufs_inputs(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256 << 30)
+    inps = [FakeBuffer(64, dtypes.float32, '1') for _ in range(2)]
+    CacheCollector.start()
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), inps[0]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out, inps[1]])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), 32, None])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out, 58, None])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    out = add_to_cache([FakeBuffer(32, dtypes.float32), out])
+    cache = CacheCollector.finish()
+    assert cache[0][1][1] == inps[0], "Input should be on its place."
+    assert cache[1][1][2] == inps[1], "Input should be on its place."
+    assert get_bufs_count(cache) == 7
+    FAKE_GLOBAL_ALLOCATOR = None
+
+if __name__ == "__main__":
+  unittest.main()

test/test_conv_shapetracker.py

@@ -2,7 +2,7 @@
 import unittest
 from tinygrad.tensor import Tensor, Device
 from tinygrad.nn import Conv2d
-from tinygrad.ops import GlobalCounters
+from tinygrad.jit import CacheCollector
 import pytest
 
 pytestmark = pytest.mark.webgpu
@@ -14,10 +14,9 @@ class TestConvShapetracker(unittest.TestCase):
     inp = Tensor.randn(1,16,10,10).realize()
     conv = Conv2d(16, 32, (3,3))
     conv(inp).realize()
-    GlobalCounters.cache = []
+    CacheCollector.start()
     conv(inp).realize()
-    test = GlobalCounters.cache
-    GlobalCounters.cache = None
+    test = CacheCollector.finish()
     assert len(test) == 1, f"conv should only have one kernel {[x[0].name for x in test]}"
     print(test[0][0].prg)
     for arg in test[0][1]:

test/test_lazybuffer.py

@@ -4,7 +4,7 @@ import unittest
 from tinygrad.lazy import LazyBuffer, Device
 from tinygrad.tensor import Tensor
 from tinygrad.shape.symbolic import Variable
-from tinygrad.ops import GlobalCounters
+from tinygrad.jit import CacheCollector
 
 class TestLazyBuffer(unittest.TestCase):
   def test_fromcpu_buffer_sharing(self):
@@ -55,18 +55,18 @@ class TestLazyBuffer(unittest.TestCase):
     d3 = in1.reshape(8,8)
     assert len(d3.lazydata.op.src[0].children) == 2
 
-    GlobalCounters.cache = []
+    CacheCollector.start()
     l = Tensor.rand(8,8)
     r = Tensor.rand(8,8)
     dd = d1 + l
     dd.realize()
     de = d3 + r
     de.realize()
-    assert len(GlobalCounters.cache) == 3
-    assert GlobalCounters.cache[0][0].name.startswith("r_") # Reduce should not merged 2 times.
-    assert GlobalCounters.cache[1][0].name.startswith("E_")
-    assert GlobalCounters.cache[2][0].name.startswith("E_")
-    GlobalCounters.cache = None
+    cache = CacheCollector.finish()
+    assert len(cache) == 3
+    assert cache[0][0].name.startswith("r_") # Reduce should not merged 2 times.
+    assert cache[1][0].name.startswith("E_")
+    assert cache[2][0].name.startswith("E_")
 
 if __name__ == "__main__":
   unittest.main()

test/test_linearizer.py

@@ -2,9 +2,9 @@ import numpy as np
 import unittest
 
 from tinygrad.codegen.linearizer import Linearizer, UOps
-from tinygrad.ops import Device
-from tinygrad.ops import GlobalCounters, Compiled
+from tinygrad.ops import Compiled, Device
 from tinygrad.tensor import Tensor
+from tinygrad.jit import CacheCollector
 
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
@@ -12,10 +12,9 @@ class TestLinearizer(unittest.TestCase):
       self.skipTest("Only Compiled supports cache")
     a, b = Tensor.randn(4), Tensor.randn(4)
     np_a, np_b = a.numpy(), b.numpy()
-    GlobalCounters.cache = []
+    CacheCollector.start()
     c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),)))).realize()
-    rawbufs = GlobalCounters.cache[0][1]
-    GlobalCounters.cache = None
+    rawbufs = CacheCollector.finish()[0][1]
     assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.lazydata.realized, b.lazydata.realized}
     np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
     np.testing.assert_allclose(np_c, c.numpy())

tinygrad/helpers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib
 import numpy as np
-from typing import Dict, Tuple, Union, List, NamedTuple, Final, Iterator, ClassVar, Optional, Callable, Any, Iterable
+from typing import Dict, Tuple, Union, List, NamedTuple, Final, Iterator, ClassVar, Optional, Iterable, Any
 from math import prod # noqa: F401 # pylint:disable=unused-import
 
 # NOTE: helpers is not allowed to import from anything else in tinygrad
@@ -129,6 +129,5 @@ class GlobalCounters:
   kernel_count: ClassVar[int] = 0
   mem_used: ClassVar[int] = 0   # NOTE: this is not reset
   mem_cached: ClassVar[int] = 0 # NOTE: this is not reset
-  cache: ClassVar[Optional[List[Tuple[Callable, Any, Dict[Any, int]]]]] = None  # List[Tuple[Callable, List[RawBuffer], Dict[Variable, int]]]
   @staticmethod
-  def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count, GlobalCounters.cache = 0,0,0.0,0,None
+  def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count = 0,0,0.0,0

tinygrad/jit.py

@@ -1,9 +1,10 @@
 from typing import Callable, List, Tuple, Any, Dict, cast, Union, Optional
+from weakref import ref
 import functools, itertools
 from tinygrad.helpers import DEBUG, DType, merge_dicts
 from tinygrad.ops import Device
 from tinygrad.tensor import Tensor
-from tinygrad.ops import GlobalCounters, RawBuffer
+from tinygrad.ops import RawBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.symbolic import Variable
 
@@ -27,7 +28,7 @@ class TinyJit:
     assert len(input_rawbuffers) != 0, "no inputs to JIT"
     assert len(set(input_rawbuffers.values())) == len(input_rawbuffers), "duplicate inputs to JIT"
     if self.cnt >= 2:
-      var_vals = dict(sorted(merge_dicts([arg.lazydata.st.var_vals for arg in args if isinstance(arg, Tensor)]).items(), key=lambda kv: kv[0].key))
+      var_vals = dict(sorted(merge_dicts([arg.lazydata.st.var_vals for arg in args if isinstance(arg, Tensor)]).items(), key=lambda kv: kv[0].key)) # type: ignore
       for (j,i),(input_name, expected_st, expected_type) in self.input_replace.items():
         assert input_rawbuffers[input_name][1].views == expected_st.views and input_rawbuffers[input_name][0].dtype == expected_type, f"ShapeTracker.views or type mismatch in JIT, <{input_rawbuffers[input_name][1].views}, {input_rawbuffers[input_name][0].dtype}> != <{expected_st.views}, {expected_type}>"
         self.jit_cache[j][1][i] = input_rawbuffers[input_name][0]
@@ -36,10 +37,9 @@ class TinyJit:
         prg(pargs, variables, jit=True)
       for (j,i) in self.input_replace.keys(): self.jit_cache[j][1][i] = None
     elif self.cnt == 1:
-      GlobalCounters.cache = []
+      CacheCollector.start()
       self.ret = self.fxn(*args, **kwargs)
-      self.jit_cache = GlobalCounters.cache
-      GlobalCounters.cache = None
+      self.jit_cache = CacheCollector.finish()
       assert len(self.jit_cache) != 0, "didn't JIT anything!"
       if DEBUG >= 1: print(f"JIT captured {len(self.jit_cache)} kernels with {len(input_rawbuffers)} inputs")
 
@@ -55,3 +55,57 @@ class TinyJit:
       self.ret = self.fxn(*args, **kwargs)
     self.cnt += 1
     return self.ret
+
+class _CacheCollector:
+  class _Placeholder:
+    def __init__(self, buf): self.size, self.dtype, self.device, self.ref, self.buftype = buf.size, buf.dtype, getattr(buf, '_device', None), ref(buf), type(buf)
+    def alive(self): return self.ref() is not None
+    def alloc_rawbuf(self): return self.buftype(self.size, self.dtype, **({'device':self.device} if self.device is not None else dict()))
+
+  def __init__(self):
+    self.cache: Optional[List[Tuple[Callable, List[Any], Dict[Any,Any]]]] = None
+    self.placeholders: Dict[RawBuffer, _CacheCollector._Placeholder] = {} # Rawbuffers are replaced with placeholders to allow freeing of the real buffer while collecting cache.
+    self.last_buftype: Dict[Tuple[int,...], int] = {} # Last index of the cached entry where a buffer with the shape (shape is a key) is used as input to the prog.
+    self.last_placeholder_index: Dict[_CacheCollector._Placeholder, int] = {} # Last index where the placeholder is used as output. This allows tracking when we need to stick to the original buffer if it is still alive.
+  def start(self):
+    self.cache, self.placeholders, self.last_buftype, self.last_placeholder_index = [], {}, {}, {}
+  def add(self, prg, rawbufs, var_vals):
+    if self.cache is None: return
+
+    # When we got buffers with the same signature, we can use just 1(max 2, see cycle avoidance below) buffer insted of all of them.
+    # Current implementation of a signature is an underlying buffer, because if 2 or more different RawBuffers shares the same, all but the very last are dead.
+    def get_signature(buf): return buf._buf if getattr(buf, '_buf', None) is not None and getattr(buf, '_allocator', None) is not None else buf
+
+    for buf in rawbufs[1:]:
+      # Check if the input matches any of placeholder to determine if it's existing or newly created input.
+      # In case of newly created input remove placeholder and capture the whole buffer.
+      if isinstance(buf, RawBuffer) and get_signature(buf) in self.placeholders and self.placeholders[get_signature(buf)].ref != ref(buf):
+        self.placeholders.pop(get_signature(buf))
+      if isinstance(buf, RawBuffer) and get_signature(buf) not in self.placeholders:
+        self.last_buftype[self._buftype_key(buf)] = len(self.cache)
+
+    # Creating/updating a placeholder for the current output buffer. If we update output, set the ref to point to the new RawBuffer,
+    # since the previous RawBuffer is dead (overwise we won't get a new RawBuffer with the same signature). Do not care about dead buffers, they 100% could be replaced with any other buffer.
+    self.placeholders.setdefault(get_signature(rawbufs[0]), _CacheCollector._Placeholder(rawbufs[0])).ref = ref(rawbufs[0])
+    self.last_placeholder_index[self.placeholders[get_signature(rawbufs[0])]] = len(self.cache)
+    self.cache.append((prg,[self.placeholders.get(get_signature(x), x) for x in rawbufs],var_vals))
+  def finish(self):
+    if self.cache is None: return []
+    placeholder_mapper, cache_result = {}, []
+    for j,(p,cached_bufs,var_vals) in enumerate(self.cache):
+      if cached_bufs[0].__class__ is _CacheCollector._Placeholder:
+        if cached_bufs[0].alive():
+          # Since the placeholder is alive (someone holds refed RawBuffer) to avoid hazards when this output buffer could be used as input on the other launch (e.g., LSTM),
+          # we allocate a backing buffer and and use it until the penultimate entry (the last entry is 100% safe to use the original RawBuffer).
+          if self.last_buftype.get(self._buftype_key(cached_bufs[0]), -1) < j or self.last_placeholder_index[cached_bufs[0]] == j:
+            # Safe to use the original buffer when all inputs buffers of the same size and dtype are behind or this is the last usage of this buffer as output.
+            placeholder_mapper[cached_bufs[0]] = cached_bufs[0].ref()
+          elif cached_bufs[0] not in placeholder_mapper:
+            placeholder_mapper[cached_bufs[0]] = cached_bufs[0].alloc_rawbuf() # Allocating a backing buffer.
+        elif cached_bufs[0] not in placeholder_mapper:
+          placeholder_mapper[cached_bufs[0]] = cached_bufs[0].alloc_rawbuf()
+      cache_result.append((p, [placeholder_mapper.get(buf, buf) for buf in cached_bufs], var_vals))
+    self.cache, self.placeholders, self.last_buftype, self.last_placeholder_index = None, {}, {}, {}
+    return cache_result
+  def _buftype_key(self, buf): return (buf.size, buf.dtype, buf.dtype.shape if hasattr(buf.dtype, 'shape') else None)
+CacheCollector = _CacheCollector()

tinygrad/ops.py

@@ -144,8 +144,9 @@ class ASTRunner:
     return self
 
   def exec(self, bufs, var_vals:Optional[Dict[Variable, int]]=None, force_wait=False, optimizing=False) -> Optional[float]:
+    from tinygrad.jit import CacheCollector
     rawbufs = dedup([x.realized for x in bufs if buf_is_kernel_arg(x)])
-    if GlobalCounters.cache is not None and not optimizing: GlobalCounters.cache.append((self, rawbufs, var_vals if var_vals is not None else {}))
+    if not optimizing: CacheCollector.add(self, rawbufs, var_vals if var_vals is not None else {})
     return self(rawbufs, var_vals, force_wait=force_wait)
 
   def __call__(self, rawbufs:List[RawBuffer], var_vals:Optional[Dict[Variable, int]]=None, jit=False, force_wait=False) -> Optional[float]:

tinygrad/runtime/lib.py

@@ -12,6 +12,7 @@ class RawBuffer:  # pylint: disable=abstract-method
     self._buf = buf if buf is not None else (allocator.alloc(size, dtype, **kwargs) if allocator else None) # If buf is provided, use it. Otherwise try to allocate from the allocator.
     self._memsz: int = size*dtype.itemsize
     self._allocator = allocator
+    self._device = kwargs.get('device', None)
     GlobalCounters.mem_used += self._memsz
   def __del__(self):  # NOTE: if it fails on init (bad dtype), it won't have a _memsz
     if hasattr(self, '_memsz'): GlobalCounters.mem_used -= self._memsz