Repository: tinygrad/tinygrad
Pull request: For cuda get current free space from device, and retry alloc failures
Instead of calculating free space use what gpu device says.
In rare case even with available space alloc may fail, due to excessive shaded space, handle that by free 10% of space more.

Problem statement:
CUDA LRU allocator runs OOM before freeing cache.
Observation: CUDA is using LRU allocator but free does not kick in and total memory gets exhausted(pycuda._driver.MemoryError: cuMemAlloc failed: out of memory), even when mem_used is very low.
Few observation:
* Initializer :  CUDAAlloc = CUDAAllocator(pycuda.driver.Context.get_device().total_memory()) 
    * But checking free and total memory, even before any gpu use, total is not same as free.
       free, total = cuda.mem_get_info()
      So maybe change it to: CUDAAlloc = CUDAAllocator(cuda.mem_get_info()[0]) 
* Code kicks in free only when current allocation will result in negative memory, which does not happen as maths to keep track of this is deterministic in code, but actual GPU allocations are >= requested bytes in parameter of mem_alloc() 
```

class CUDAAllocator(LRUAllocator):
    def __init__(self, dev_memsz):
      super().__init__(dev_memsz)
      free, total = cuda.mem_get_info()
      assert free < total
      self.sizes = 0
      self.total = dev_memsz

def _do_alloc(self, size, dtype, device, **kwargs):
      actual_rem = cuda.mem_get_info()[0]
      self.sizes += size * dtype.itemsize
      print(f'alloc {self.sizes} of {self.total} remaining {self.total - self.sizes}')
      a = cuda.mem_alloc(size * dtype.itemsize)  # type: ignore
      actual_rem = cuda.mem_get_info()[0]
      assert actual_rem == self.total - self.sizes, f'alloc {actual_rem}, {self.total}, {self.sizes}'       #FAILS
      return a

CUDAAlloc = CUDAAllocator(cuda.mem_get_info()[0])

```
This assert fails after first allocation.
I have observed it on 1x RTX A5000 on vast.ai cloud.

Hints from issue discussion:
pycuda also has memory pool in api https://documen.tician.de/pycuda/util.html#device-based-memory-pool, can be used directly as allocator.
This should be a problem with all GPU backends, maybe we can stop tracking free size, and free when alloc fails, I can send a change for this it it LG.

Commit message: For cuda get current free space from device, and rery alloc failures

Patch to review:
diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py
index 75e8b1b58e2a5..de5cef92b96a4 100644
--- a/tinygrad/runtime/lib.py
+++ b/tinygrad/runtime/lib.py
@@ -75,13 +75,13 @@ def _cache_reuse_buffer(self, rawbufs: Deque[Tuple[Any, int]]): # The newest cac
     GlobalCounters.mem_cached -= self._underlying_buf_memsz(rawbufs[0][0])
     return rawbufs.popleft()[0]
 
-  def ensure_has_free_space(self, size, dtype, device):
-    while len(self.aging_order[device]) and (self.free_space[device]-size*dtype.itemsize) < 0: # When OOM removing lru buffers.
+  def ensure_has_free_space(self, space_to_free, device):
+    while len(self.aging_order[device]) and self._get_free_space(device) < space_to_free: # When OOM removing lru buffers.
       bucket, epoch = self.aging_order[device].popleft()
       if self.cached_buffers[bucket] and self.cached_buffers[bucket][-1][1] == epoch: self._free_buffer(self.cached_buffers[bucket].pop()[0]) # Free cached buffer if it is still in cache.
 
   def _alloc_buffer(self, size, dtype, device, **kwargs):
-    self.ensure_has_free_space(size, dtype, device)
+    self.ensure_has_free_space(size*dtype.itemsize, device)
     self.free_space[device] -= size*dtype.itemsize
     newbuf = self._do_alloc(max(1, size), dtype, device, **kwargs)
     self.buffer_info[newbuf] = (size, dtype, device)
@@ -108,3 +108,4 @@ def _underlying_buf_memsz(self, buf): return self.buffer_info[buf][0] * self.buf
   def _cached_bufkey(self, size, dtype, device) -> Tuple[int, ...]: return (device, size, dtype, dtype.shape) if isinstance(dtype, ImageDType) else (device, size, dtype) # Provides a key for reusing device buffers with identical keys.
   def _do_alloc(self, size, dtype, device, **kwargs): raise NotImplementedError("must be implemented")
   def _do_free(self, buf): pass
+  def _get_free_space(self, device): return self.free_space[device]
diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py
index 71dc74c0a0104..dd43d5b333350 100644
--- a/tinygrad/runtime/ops_cuda.py
+++ b/tinygrad/runtime/ops_cuda.py
@@ -49,9 +49,16 @@ class device:
   import pycuda.autoprimaryctx # type: ignore # pylint: disable=unused-import # noqa: F401
   import pycuda.driver as cuda # type: ignore
   class CUDAAllocator(LRUAllocator):
-    def _do_alloc(self, size, dtype, device, **kwargs): return cuda.mem_alloc(size * dtype.itemsize) # type: ignore
+    def _do_alloc(self, size, dtype, device, **kwargs):
+      try: return cuda.mem_alloc(size * dtype.itemsize)
+      except cuda.MemoryError as e:
+        if len(self.aging_order[device]):
+          self.ensure_has_free_space(1.1*self._get_cur_free_spce(device), device) # increase free space by 10% and try again.
+          return self._do_alloc(size, dtype, device, **kwargs)
+        raise e
     def _cached_bufkey(self, size, dtype, device): return (device, size*dtype.itemsize) # Buffers of the same length could be reused, no matter what dtype.
-  CUDAAlloc = CUDAAllocator(pycuda.driver.Context.get_device().total_memory())
+    def _get_cur_free_spce(self, deivice): return cuda.mem_get_info()[0]
+  CUDAAlloc = CUDAAllocator(cuda.memgetinfo()[0])
   class RawCUDABuffer(RawBufferCopyInOut): # type: ignore
     def __init__(self, size, dtype): super().__init__(size, dtype, allocator=CUDAAlloc)
     def _copyin(self, x:np.ndarray, stream:Optional[cuda.Stream]=None): cuda.memcpy_htod_async(self._buf, x.ravel(), stream) # type: ignore
diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py
index 8a9057b9762fc..5d31f3bb83c18 100644
--- a/tinygrad/runtime/ops_gpu.py
+++ b/tinygrad/runtime/ops_gpu.py
@@ -52,7 +52,7 @@ def _copyin(self, x:np.ndarray):
     self.event = cl.enqueue_copy(CL.cl_queue[self._buf.device], self._buf, np.require(x, requirements=['C', 'A']), is_blocking=False)
   def _copyout(self, x:np.ndarray):
     assert not self.dtype.name.startswith("image"), f"can't copyout images {self.dtype}"
-    CL.cl_allocator.ensure_has_free_space(self.size, self.dtype, self._device)
+    CL.cl_allocator.ensure_has_free_space(self.size*self.dtype.itemsize, self._device)
     buf = cl.Buffer(CL.cl_ctxs[self._buf.device], cl.mem_flags.WRITE_ONLY | cl.mem_flags.USE_HOST_PTR, 0, hostbuf=x.data)
     mapped, event = cl.enqueue_map_buffer(CL.cl_queue[self._buf.device], buf, cl.map_flags.WRITE, 0, self.size, dtype=self.dtype.np, is_blocking=False)
     with mapped.base: cl.enqueue_copy(CL.cl_queue[self._buf.device], mapped, self._buf, is_blocking=True, wait_for=[event] + ([self.event] if hasattr(self, "event") else []))


Code review comments:
diff_hunk: @@ -49,9 +49,16 @@ class device:
   import pycuda.autoprimaryctx # type: ignore # pylint: disable=unused-import # noqa: F401
   import pycuda.driver as cuda # type: ignore
   class CUDAAllocator(LRUAllocator):
-    def _do_alloc(self, size, dtype, device, **kwargs): return cuda.mem_alloc(size * dtype.itemsize) # type: ignore
+    def _do_alloc(self, size, dtype, device, **kwargs):
+      try: return cuda.mem_alloc(size * dtype.itemsize)
+      except cuda.MemoryError as e:
+        if len(self.aging_order[device]):
+          self.ensure_has_free_space(1.1*self._get_cur_free_spce(device), device) # increase free space by 10% and try again.
+          return self._do_alloc(size, dtype, device, **kwargs)
+        raise e
     def _cached_bufkey(self, size, dtype, device): return (device, size*dtype.itemsize) # Buffers of the same length could be reused, no matter what dtype.
-  CUDAAlloc = CUDAAllocator(pycuda.driver.Context.get_device().total_memory())
+    def _get_cur_free_spce(self, deivice): return cuda.mem_get_info()[0]
path: tinygrad/runtime/ops_cuda.py
text: @user1:
Can we use `_get_free_space` here as well and just overload it?

@author:
added __init__ to call this.

diff_hunk: @@ -49,9 +49,16 @@ class device:
   import pycuda.autoprimaryctx # type: ignore # pylint: disable=unused-import # noqa: F401
   import pycuda.driver as cuda # type: ignore
   class CUDAAllocator(LRUAllocator):
-    def _do_alloc(self, size, dtype, device, **kwargs): return cuda.mem_alloc(size * dtype.itemsize) # type: ignore
+    def _do_alloc(self, size, dtype, device, **kwargs):
+      try: return cuda.mem_alloc(size * dtype.itemsize)
+      except cuda.MemoryError as e:
+        if len(self.aging_order[device]):
+          self.ensure_has_free_space(1.1*self._get_cur_free_spce(device), device) # increase free space by 10% and try again.
path: tinygrad/runtime/ops_cuda.py
text: @user1:
Maybe move this logic into LRU to cover all backends? 

Btw, do we really need an overloaded `free_space` in this case? Yes our `self.free_space` is an approximation but in any case if we cannot alloc we will clear buffers till we can allocate.

Something like?
```python
def _alloc_buffer(self, size, dtype, device, **kwargs):
    newbuf, alloc_size = None, size*dtype.itemsize
    while newbuf is None:
        self.ensure_has_free_space(alloc_size, device)
        try: newbuf = self._do_alloc(max(1, size), dtype, device, **kwargs)
        except Exception as e:
            if not len(self.aging_order[device]): raise e
            alloc_size *= 1.1
    self.free_space[device] -= size*dtype.itemsize
    self.buffer_info[newbuf] = (size, dtype, device)
    return newbuf
```

@author:
Different backend will raise different exceptions, either we handle any exception, which can be unrelated to memory, or we do an 'or' of all known memory error from each backend, which will result in backend related imports in lib.py or we introduce our memory exception and wrap backend exception in it handle that in lib.py. To avoid all that I have just handled it cuda file.

Alloc is costly, finding free space on gpu is not. I am not sure if unsuccessful alloc is fail fast. Also I have observed it diverges significantly for free_space calculation.

@user1:
I don't think it belongs to cuda file, since it's related to all backends. I think it's fine to catch all exceptions.

@author:
Done.