diff --git a/test/test_gc.py b/test/test_gc.py index 78827a0fdc..faa25a741f 100644 --- a/test/test_gc.py +++ b/test/test_gc.py @@ -86,10 +86,10 @@ class TestGC(unittest.TestCase): a.realize() real_buf = a.uop.buffer # after the Tensor UOp is deleted there shouldn't be any references on the Buffer - self.assertEqual(real_buf.lb_refcount, 1) + self.assertEqual(real_buf.uop_refcount, 1) self.assertEqual(bufs_allocated()-init, 1) del a.uop - self.assertEqual(real_buf.lb_refcount, 0) + self.assertEqual(real_buf.uop_refcount, 0) self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive del real_buf self.assertEqual(bufs_allocated()-init, 0) @@ -99,14 +99,14 @@ class TestGC(unittest.TestCase): a = Tensor.full((4,), 1.).contiguous() a.realize() real_buf = a.uop.buffer - self.assertEqual(real_buf.lb_refcount, 1) + self.assertEqual(real_buf.uop_refcount, 1) a.assign(Tensor.full((4,), 2.)) self.assertIs(a.uop.src[0].buffer, real_buf) # NOTE: this is still 1, we don't count the ASSIGN - self.assertEqual(real_buf.lb_refcount, 1) + self.assertEqual(real_buf.uop_refcount, 1) a.realize() del a - self.assertEqual(real_buf.lb_refcount, 0) # no UOps for this Buffer + self.assertEqual(real_buf.uop_refcount, 0) # no UOps for this Buffer self.assertEqual(bufs_allocated()-init, 1) # Buffer is alive if __name__ == '__main__': diff --git a/test/test_memory_planner.py b/test/test_memory_planner.py index dcd91569bd..0ba3f90404 100644 --- a/test/test_memory_planner.py +++ b/test/test_memory_planner.py @@ -18,13 +18,13 @@ def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]]): first_appearance, last_appearance = {}, {} for i,u in enumerate(buffers): for buf in u: - if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue + if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue if buf.base not in first_appearance: first_appearance[buf.base] = i last_appearance[buf.base] = i for i,u in enumerate(buffers): for buf in u: - if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue + if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue cur, base = assigned.get(buf, buf), assigned.get(buf.base, buf.base) if buf._base is not None: assert cur.base == base.base and cur.offset == buf.offset + base.offset, f"failed: {buf} {cur} {base} {buf.offset} {base.offset}" diff --git a/tinygrad/device.py b/tinygrad/device.py index 921376667b..21a0c29bfb 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -106,14 +106,14 @@ class MultiBuffer: class Buffer: def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:Optional[BufferSpec]=None, initial_value:Optional[bytes]=None, - lb_refcount=0, base:Optional[Buffer]=None, offset:int=0, preallocate=False): + uop_refcount=0, base:Optional[Buffer]=None, offset:int=0, preallocate=False): if isinstance(dtype, ImageDType): options = BufferSpec(image=dtype) # TODO: image hack shouldn't be here. where should it be? else: assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType) self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0 if base is None: assert offset == 0, "base buffers can't have offset" self._base = None - self._lb_refcount = lb_refcount + self._uop_refcount = uop_refcount if opaque is not None: self.allocate(opaque) if initial_value is not None: self.allocate() @@ -126,9 +126,9 @@ class Buffer: @property def base(self) -> Buffer: return self._base if self._base is not None else self @property - def lb_refcount(self): return self.base._lb_refcount + def uop_refcount(self): return self.base._uop_refcount def ref(self, cnt): - self.base._lb_refcount += cnt + self.base._uop_refcount += cnt return self def is_allocated(self) -> bool: return hasattr(self, '_buf') def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self @@ -160,11 +160,11 @@ class Buffer: buf = None if self._base is not None: return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated()) - if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.lb_refcount) + if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount) if self.is_allocated(): buf = bytearray(self.nbytes) self.copyout(memoryview(buf)) - return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.lb_refcount) + return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount) @property def nbytes(self): return self.size*self.dtype.itemsize def __del__(self): (not self.is_allocated()) or self.deallocate() diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index c63810f101..a093ba05b9 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -232,7 +232,7 @@ class TinyJit(Generic[ReturnType]): def add_buffer(self, b:Buffer) -> Buffer: if found:=self._buffer_replace.get(b, None): return found - if b.is_allocated() or b.lb_refcount > 0: return b + if b.is_allocated() or b.uop_refcount > 0: return b if b._base is not None: self._buffer_replace[b] = ret = Buffer(b.device, b.size, b.dtype, base=self.add_buffer(b._base), offset=b.offset) else: diff --git a/tinygrad/engine/memory.py b/tinygrad/engine/memory.py index e925972d1d..96fd6a0678 100644 --- a/tinygrad/engine/memory.py +++ b/tinygrad/engine/memory.py @@ -14,7 +14,7 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign first_appearance, last_appearance, buf_to_opt = {}, {}, set() for i,u in enumerate(buffers): for buf in u: - should_skip = buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers) + should_skip = buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers) if not ignore_checks and should_skip: continue if buf.base not in first_appearance: first_appearance[buf.base] = i last_appearance[buf.base] = i