Repository: tinygrad/tinygrad
Pull request: fix jit buffer reuse when freed
Closes #1799

Problem statement:
Stable Diffusion Memory Usage on M1 8GB is 11.62 GB
Clip:
```
***    0 r_386_32_12352_4_4                arg   1 sz [386]              [32]         OPs   2441M/   0.00G  mem  4.27 GB tm  10836.92us/    10.84ms (  225.27 GFLOPS,    0.02 GB/s)
```

First step:
```
***    1 r_5_5                             arg   1 sz []                 []           OPs      0M/   0.00G  mem  3.77 GB tm      7.29us/     0.01ms (    0.00 GFLOPS,    0.00 GB/s)
```

Second step:
```
***    1 E_                                arg   2 sz []                 []           OPs      0M/   0.00G  mem  3.77 GB tm      6.92us/     0.02ms (    0.00 GFLOPS,    0.00 GB/s)
```

Third step:
```
***    1 E_                                arg   2 sz []                 []           OPs      0M/   0.00G  mem 11.62 GB tm      6.88us/     0.02ms (    0.00 GFLOPS,    0.00 GB/s)
```

Hints from issue discussion:

Commit message: fix jit buffer reuse when freed

Patch to review:
diff --git a/test/test_cache_collector.py b/test/test_cache_collector.py
index 2fd52f5a4b15f..9a8b2dc2c6387 100644
--- a/test/test_cache_collector.py
+++ b/test/test_cache_collector.py
@@ -164,5 +164,24 @@ def test_cache_collector_anybufs_inputs(self):
     assert get_bufs_count(cache) == 7
     FAKE_GLOBAL_ALLOCATOR = None
 
+  def test_cache_collector_optimize_when_not_cached_anymore(self):
+    global FAKE_GLOBAL_ALLOCATOR
+    FAKE_GLOBAL_ALLOCATOR = FakeAllocator(256)
+    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(240, 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(240, 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[0][1][0] == cache[3][1][0], "Output buffers from 1st and 4th should be the same"
+    FAKE_GLOBAL_ALLOCATOR = None
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tinygrad/jit.py b/tinygrad/jit.py
index 2f29bc4a19637..84db066b55da5 100644
--- a/tinygrad/jit.py
+++ b/tinygrad/jit.py
@@ -1,5 +1,6 @@
 from typing import Callable, List, Tuple, Any, Dict, cast, Union, Optional
 from weakref import ref
+from collections import defaultdict
 import functools, itertools
 from tinygrad.helpers import DEBUG, DType, merge_dicts
 from tinygrad.ops import Device
@@ -64,17 +65,18 @@ def __call__(self, *args, **kwargs) -> Any:
 
 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 __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 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.
+    self.freed_placeholders: Dict[Tuple[int,...], List[_CacheCollector._Placeholder]] = defaultdict(list)
   def start(self):
-    self.cache, self.placeholders, self.last_buftype, self.last_placeholder_index = [], {}, {}, {}
+    self.cache, self.placeholders, self.last_buftype, self.last_placeholder_index, self.freed_buffers = [], {}, {}, {}, defaultdict(list)
   def add(self, prg, rawbufs, var_vals):
     if self.cache is None: return
 
@@ -92,7 +94,11 @@ def get_signature(buf): return buf._buf if getattr(buf, '_buf', None) is not Non
 
     # 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])
+    if get_signature(rawbufs[0]) in self.placeholders: self.placeholders[get_signature(rawbufs[0])].ref = ref(rawbufs[0])
+    else:
+      # This is a new buffer, if we have any free placeholders to overwrite, let's do this, since this results into one buffer.
+      plh = self.freed_placeholders[self._buftype_key(rawbufs[0])].pop() if self.freed_placeholders[self._buftype_key(rawbufs[0])] else _CacheCollector._Placeholder(rawbufs[0])
+      self.placeholders.setdefault(get_signature(rawbufs[0]), plh).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):
@@ -111,7 +117,11 @@ def finish(self):
         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, {}, {}, {}
+    self.cache, self.placeholders, self.last_buftype, self.last_placeholder_index, self.freed_buffers = None, {}, {}, {}, defaultdict(list)
     return cache_result
-  def _buftype_key(self, buf): return (buf.size, buf.dtype, buf.dtype.shape if hasattr(buf.dtype, 'shape') else None)
+  def _on_buf_free(self, underlying_buf):
+    if underlying_buf not in self.placeholders: return
+    self.freed_placeholders[self._buftype_key(self.placeholders[underlying_buf])].append(self.placeholders[underlying_buf])
+    self.placeholders.pop(underlying_buf)
+  def _buftype_key(self, buf): return (buf.size, buf.dtype, buf._device, buf.dtype.shape if hasattr(buf.dtype, 'shape') else None)
 CacheCollector = _CacheCollector()
\ No newline at end of file
diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py
index 4107400909207..8a47cc4ec4191 100644
--- a/tinygrad/runtime/lib.py
+++ b/tinygrad/runtime/lib.py
@@ -95,6 +95,8 @@ def _alloc_buffer(self, size, dtype, device, **kwargs):
     self.buffer_info[newbuf] = (size, dtype, device)
     return newbuf
   def _free_buffer(self, buf_to_free):
+    from tinygrad.jit import CacheCollector
+    CacheCollector._on_buf_free(buf_to_free)
     self.free_space[self.buffer_info[buf_to_free][2]] += self._underlying_buf_memsz(buf_to_free)
     GlobalCounters.mem_cached -= self._underlying_buf_memsz(buf_to_free)
     self.buffer_info.pop(buf_to_free)


Code review comments:
diff_hunk: @@ -64,37 +65,41 @@ def __call__(self, *args, **kwargs) -> Any:
 
 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 __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 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.
+    self.freed_placeholders: Dict[Tuple[int,...], List[_CacheCollector._Placeholder]] = defaultdict(list)
+    self.cirular_signatures: Set[Any] = set()
path: tinygrad/jit.py
text: @user1:
typo