diff --git a/tinygrad/device.py b/tinygrad/device.py index 6f95b01a4c..4b69c41fa4 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -6,7 +6,7 @@ import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK -from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey +from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey, size_to_str from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype if TYPE_CHECKING: from tinygrad.renderer import Renderer @@ -212,7 +212,9 @@ class Allocator(Generic[DeviceType]): # overridden in LRUAllocator def alloc(self, size:int, options:BufferSpec|None=None): assert size > 0, f"alloc size must be positive, getting {size}" - return self._alloc(size, options if options is not None else self.default_buffer_spec) + try: return self._alloc(size, options if options is not None else self.default_buffer_spec) + except (RuntimeError, MemoryError) as e: raise MemoryError(f"Allocation of {size_to_str(size)} failed on {self.dev.device}. " + f"Used: {size_to_str(GlobalCounters.mem_used_per_device[self.dev.device])}") from e def free(self, opaque, size:int, options:BufferSpec|None=None): self._free(opaque, options if options is not None else self.default_buffer_spec) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 9af076e19c..b43ef6666a 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -36,6 +36,7 @@ def colored(st, color:str|None, background=False): # replace the termcolor libra return f"\u001b[{10*background+60*(color.upper() == color)+30+colors.index(color.lower())}m{st}\u001b[0m" if color is not None else st def colorize_float(x: float): return colored(f"{x:7.2f}x", 'green' if x < 0.75 else 'red' if x > 1.15 else 'yellow') def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,pr in [(1, "s "),(1e3, "ms")] if t > 10/d), f"{t * 1e6:{w}.2f}us") +def size_to_str(s:int) -> str: return next((f"{s / d:.2f} {pr}" for d,pr in [(1<<30, "GB"),(1<<20, "MB"),(1<<10, "KB")] if s >= d), f"{s} B") def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s) def ansilen(s:str): return len(ansistrip(s)) def make_tuple(x:int|Sequence[int], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x) diff --git a/tinygrad/runtime/support/memory.py b/tinygrad/runtime/support/memory.py index 349003323b..b3ef1a329e 100644 --- a/tinygrad/runtime/support/memory.py +++ b/tinygrad/runtime/support/memory.py @@ -248,7 +248,7 @@ class MemoryManager: nxt_range += 1 if nxt_range == len(self.palloc_ranges): for paddr, _ in paddrs: self.pfree(paddr) - raise MemoryError(f"Failed to allocate memory. (total allocation size={size:#x}, current try={self.palloc_ranges[nxt_range-1]})") + raise MemoryError(f"Failed to allocate memory (OOM). Request size={size:#x} ({self.palloc_ranges[nxt_range-1]})") continue rem_size -= self.palloc_ranges[nxt_range][0]