s/lb_refcount/uop_refcount [pr] (#10865)

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__':

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}"

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()

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:

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