mirror of
https://jihulab.com/mr-one/openpilot.git
synced 2026-10-28 05:23:46 +08:00
Comma Device
540 lines
29 KiB
Python
540 lines
29 KiB
Python
from __future__ import annotations
|
|
from dataclasses import dataclass, replace, field
|
|
from collections import defaultdict
|
|
from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
|
|
import importlib, inspect, functools, itertools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct, mmap, time, statistics
|
|
from tinygrad.helpers import mv_address, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
|
|
from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
|
|
from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up, is_numpy_ndarray
|
|
from tinygrad.helpers import cpu_profile, perf_counter_us, to_name, HCQ_RUNTIME_DEV, to_tuple
|
|
from tinygrad.dtype import dtypes, DType
|
|
from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
|
|
if TYPE_CHECKING:
|
|
from tinygrad.renderer import Renderer
|
|
|
|
# **************** Device ****************
|
|
|
|
ALL_DEVICES = ["METAL", "AMD", "NV", "CUDA", "QCOM", "CL", "CPU", "DSP", "WEBGPU"]
|
|
class _Device:
|
|
def __init__(self) -> None:
|
|
self._devices = [x.stem[len("ops_"):].upper() for x in (pathlib.Path(__file__).parent/"runtime").iterdir() if x.stem.startswith("ops_")]
|
|
self._opened_devices:set[str] = set()
|
|
@functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none
|
|
def _canonicalize(self, device:str) -> str: return re.sub(r":0$", "", (d:=device.split(":", 1)[0].upper()) + device[len(d):])
|
|
# NOTE: you can't cache canonicalize in case Device.DEFAULT changes
|
|
def canonicalize(self, device:str|None) -> str: return self._canonicalize(device if device is not None else Device.DEFAULT)
|
|
def __getitem__(self, ix:str) -> Compiled:
|
|
ix = self.canonicalize(ix)
|
|
assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
|
|
return self.__get_canonicalized_item(ix)
|
|
@functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none
|
|
def get_class(self, ix:str):
|
|
base = (__package__ or __name__).split('.')[0] # tinygrad
|
|
x = ix.split(":")[0].lower()
|
|
return [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) if (cname.lower() == x + "device")][0]
|
|
@functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none
|
|
def __get_canonicalized_item(self, ix:str) -> Compiled:
|
|
ret = self.get_class(ix)(ix)
|
|
if DEBUG >= 1: print(f"opened device {ix} from pid:{os.getpid()}")
|
|
self._opened_devices.add(ix)
|
|
return ret
|
|
@property
|
|
def default(self) -> Compiled: return self[self.DEFAULT]
|
|
def get_available_devices(self) -> Iterator[str]:
|
|
for device in ALL_DEVICES:
|
|
with contextlib.suppress(Exception): yield self[device].device
|
|
@property
|
|
def DEFAULT(self) -> str: return DEV.device or self._select_device
|
|
@DEFAULT.setter
|
|
def DEFAULT(self, v): raise AttributeError(f'setting Device.DEFAULT is deprecated, use "with Context(DEV={v!r})" or "DEV.value = {v!r}"')
|
|
@functools.cached_property
|
|
def _select_device(self) -> str:
|
|
assert (dev:=next((d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1), None)) is None, \
|
|
f"{dev}=1 is deprecated, use DEV={dev} instead"
|
|
try:
|
|
device = next(self.get_available_devices())
|
|
os.environ["DEV"] = device # we set this in environment for spawned children
|
|
return device
|
|
except StopIteration as exc: raise RuntimeError("no usable devices") from exc
|
|
Device: _Device = _Device()
|
|
atexit.register(lambda: [Device[dn].finalize() for dn in tuple(Device._opened_devices)])
|
|
|
|
def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
|
|
if not isinstance(device, (tuple, list)): return Device.canonicalize(device)
|
|
return canonical[0] if len(canonical:=tuple(Device.canonicalize(d) for d in device)) == 1 else canonical
|
|
|
|
def is_disk_device(device:str|tuple[str, ...]) -> bool:
|
|
return any(d.split(":", 1)[0].upper() == "DISK" for d in to_tuple(device))
|
|
|
|
# **************** Profile ****************
|
|
|
|
@dataclass(frozen=True)
|
|
class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
|
|
|
|
@dataclass(frozen=True)
|
|
class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None; profile_key:bytes|None=None # noqa: E702
|
|
|
|
@dataclass(frozen=True)
|
|
class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int; profile_key:bytes|None=None # noqa: E702
|
|
|
|
@dataclass(frozen=True)
|
|
class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
|
|
|
|
# **************** Buffer + Allocators ****************
|
|
|
|
@dataclass(frozen=True, eq=True)
|
|
class BufferSpec:
|
|
uncached: bool = False
|
|
cpu_access: bool = False
|
|
host: bool = False
|
|
nolru: bool = False
|
|
external_ptr: int|None = None
|
|
|
|
class MultiBuffer:
|
|
def __init__(self, device:tuple[str, ...], nbytes:int):
|
|
self.bufs = [Buffer(d, nbytes) for d in device]
|
|
@property
|
|
def device(self): return tuple(x.device for x in self.bufs)
|
|
@property
|
|
def nbytes(self): return self.bufs[0].nbytes
|
|
def is_allocated(self): return all(x.is_allocated() for x in self.bufs)
|
|
def __repr__(self): return f"<multibuf real:{self.is_allocated()} device:{tuple(x.device for x in self.bufs)} nbytes:{self.nbytes}>"
|
|
|
|
@dataclass(frozen=True)
|
|
class BufferStorage: buf:Any; meta:Any=None; host:MMIOInterface|None=None; maps:dict[Compiled, BufferStorage]=field(default_factory=dict) # noqa: E702
|
|
|
|
class Buffer:
|
|
profile_events:list[ProfileEvent] = []
|
|
def __init__(self, device:str, nbytes:int, opaque:Any=None, options:BufferSpec|None=None,
|
|
initial_value:bytes|pickle.PickleBuffer|None=None, base:Buffer|None=None, offset:int=0, preallocate=False,
|
|
allocator:Allocator|None=None):
|
|
self.device, self.nbytes, self.offset, self._base = Device.canonicalize(device), nbytes, offset, base
|
|
if allocator is not None: self.allocator = allocator
|
|
self.options = options if options is not None else BufferSpec()
|
|
self._storage:BufferStorage|None = None
|
|
if base is None:
|
|
assert offset == 0, "base buffers can't have offset"
|
|
if opaque is not None: self.allocate(opaque)
|
|
if initial_value is not None:
|
|
self.allocate()
|
|
if (host:=self.get_storage().host) is not None: host[:] = memoryview(initial_value).cast('B')
|
|
else: # same-size chunks reuse one upload program instead of compiling one per buffer size
|
|
src = Buffer("PYTHON", self.nbytes, initial_value=initial_value)
|
|
for offset in range(0, self.nbytes, 1 << 20):
|
|
size = min(self.nbytes - offset, 1 << 20)
|
|
self.view(size, offset).ensure_allocated().copy_from(src.view(size, offset).ensure_allocated())
|
|
if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
|
|
else:
|
|
assert base._base is None, "base can't have a base"
|
|
assert self.device == base.device, "base must have the same device"
|
|
if preallocate: self.allocate()
|
|
|
|
@suppress_finalizing
|
|
def __del__(self): self._storage is None or self.deallocate()
|
|
|
|
def __repr__(self):
|
|
return f"<buf real:{self.is_allocated()} device:{self.device} nbytes:{self.nbytes}" + \
|
|
(f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options != BufferSpec() else "") + ">"
|
|
|
|
@property
|
|
def base(self) -> Buffer: return self._base if self._base is not None else self
|
|
@functools.cached_property
|
|
def allocator(self) -> Allocator: return self.base.allocator if self._base is not None else Device[self.device].allocator
|
|
@property
|
|
def _buf(self) -> Any: return self.get_storage().buf
|
|
@property
|
|
def host(self) -> MMIOInterface: return unwrap(self.get_storage().host)
|
|
@property
|
|
def meta(self) -> Any: return self.get_storage().meta
|
|
|
|
def get_storage(self, device:str|None=None) -> BufferStorage:
|
|
storage = unwrap(self.ensure_allocated()._storage)
|
|
device = Device.canonicalize(device) if device is not None else self.device
|
|
if device == self.device: return storage
|
|
if (dev:=Device[device]) not in storage.maps:
|
|
alloc = dev.allocator
|
|
storage.maps[dev] = BufferStorage(alloc._offset(self.base.get_buf(device), self.nbytes, self.offset)) if self._base else alloc.map(self)
|
|
if storage.maps[dev].host is not storage.host: storage.maps[dev] = replace(storage.maps[dev], host=storage.host)
|
|
return storage.maps[dev]
|
|
|
|
def get_buf(self, device:str) -> Any: return self.get_storage(device).buf
|
|
|
|
def is_allocated(self) -> bool: return self._storage is not None and (self._base is None or self._base_storage is self.base._storage)
|
|
def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self
|
|
def allocate(self, opaque=None, external_ptr=None) -> Buffer:
|
|
assert not self.is_allocated(), "can't allocate already allocated buffer"
|
|
if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
|
|
# MAX_BUFFER_SIZE limits raw storage bytes (not logical elements).
|
|
if not self.device.startswith("NULL") and self.nbytes > MAX_BUFFER_SIZE > 0 and self.options.external_ptr is None:
|
|
raise RuntimeError(f"buffer of size {self.nbytes/1e6:.2f}MB is too large")
|
|
if external_ptr is not None: self.options = replace(self.options, external_ptr=external_ptr)
|
|
if self._base is not None:
|
|
storage = replace(self.base.get_storage(), buf=self.allocator._offset(self.base._buf, self.nbytes, self.offset), maps={})
|
|
elif opaque is not None:
|
|
self.options = replace(self.options, nolru=True)
|
|
if is_numpy_ndarray(opaque):
|
|
if not opaque.flags.c_contiguous: opaque = opaque.copy(order='C')
|
|
opaque = BufferStorage(addr:=opaque.ctypes.data, memoryview(opaque), MMIOInterface(addr, self.nbytes))
|
|
elif isinstance(opaque, memoryview):
|
|
opaque = BufferStorage(addr:=mv_address(opaque) if self.nbytes else 0, opaque, MMIOInterface(addr, self.nbytes))
|
|
storage = opaque if isinstance(opaque, BufferStorage) else BufferStorage(opaque)
|
|
else: storage = self.allocator.alloc(self.nbytes, self.options)
|
|
storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None)
|
|
if self._base is None:
|
|
if not self.device.startswith("DISK") and self.options.external_ptr is None:
|
|
GlobalCounters.mem_used += self.nbytes
|
|
GlobalCounters.mem_used_per_device[self.device] += self.nbytes
|
|
if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"nbytes":self.nbytes}))
|
|
self._storage, self._base_storage = storage, self.base._storage if self._base else None
|
|
return self
|
|
|
|
def deallocate(self):
|
|
assert self._storage is not None, "buffer must be allocated to deallocate"
|
|
if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
|
|
if self._base is None:
|
|
if GlobalCounters is not None and not self.device.startswith("DISK") and self.options.external_ptr is None:
|
|
GlobalCounters.mem_used -= self.nbytes
|
|
GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
|
|
if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
|
|
self.allocator.free(self._storage, self.nbytes, self.options)
|
|
self._storage, self._base_storage = None, None
|
|
|
|
def __reduce_ex__(self, protocol):
|
|
buf:bytearray|pickle.PickleBuffer|None = None
|
|
if self._base is not None:
|
|
return self.__class__, (self.device, self.nbytes, None, None, None, self.base, self.offset, self.is_allocated())
|
|
if self.device == "NPY" and self.is_allocated():
|
|
import numpy as np
|
|
arr = np.frombuffer(self.meta, np.uint8, count=self.nbytes) # an out-of-band pickle buffer keeps the storage alive
|
|
return self.__class__, (self.device, self.nbytes, arr, self.options, None)
|
|
if self.is_allocated():
|
|
buf = pickle.PickleBuffer(self.as_memoryview()) if protocol >= 5 else bytearray(self.as_memoryview())
|
|
return self.__class__, (self.device, self.nbytes, None, self.options, buf)
|
|
|
|
@property
|
|
def trace_num(self) -> int:
|
|
if not hasattr(self, '_trace_num'): self._trace_num = len(Buffer.profile_events)
|
|
return self._trace_num
|
|
|
|
def _host_mv(self) -> memoryview|None:
|
|
if self.is_allocated() and hasattr(host:=self.get_storage().host, 'mv'): return unwrap(host).view(fmt='B').mv
|
|
if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
|
|
return None
|
|
|
|
def as_memoryview(self, allow_zero_copy=False) -> memoryview:
|
|
if (mv:=self._host_mv()) is not None:
|
|
self.allocator.dev.synchronize()
|
|
if allow_zero_copy: return mv
|
|
with cpu_profile(f"{self.device} -> TINY", f"{self.device}:COPY"): return memoryview(bytearray(mv))
|
|
Buffer("NPY", self.nbytes, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
|
|
return mv
|
|
|
|
def copy_from(self, src:Buffer) -> Buffer:
|
|
assert self.nbytes == src.nbytes, f"copy size mismatch, {self.nbytes} != {src.nbytes}"
|
|
assert self.is_allocated() and src.is_allocated(), "copy requires allocated buffers"
|
|
from tinygrad.engine.realize import run_linear
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
du, su = UOp.from_buffer(self, dtypes.uint8), UOp.from_buffer(src, dtypes.uint8)
|
|
run_linear(UOp(Ops.LINEAR, src=(du.store_call(su),)), update_stats=False)
|
|
return self
|
|
|
|
def view(self, nbytes:int, offset:int) -> Buffer:
|
|
assert offset < self.nbytes, "offset must be less than nbytes"
|
|
return Buffer(self.device, nbytes, base=self.base, offset=self.offset+offset)
|
|
|
|
DeviceType = TypeVar('DeviceType', bound='Compiled')
|
|
|
|
# TODO: size, dest, src are the same type. can we enforce this?
|
|
class Allocator(Generic[DeviceType]):
|
|
lru = True
|
|
|
|
def __init__(self, dev:DeviceType):
|
|
self.dev: DeviceType = dev
|
|
self.default_buffer_spec: BufferSpec = BufferSpec()
|
|
self.cache:dict[tuple[int, BufferSpec|None], list[BufferStorage]] = defaultdict(list)
|
|
|
|
def alloc(self, size:int, options:BufferSpec|None=None) -> BufferStorage:
|
|
assert size > 0, f"alloc size must be positive, getting {size}"
|
|
if len(c:=self.cache[(size, options)]): return c.pop()
|
|
spec = options if options is not None else self.default_buffer_spec
|
|
try: return self._alloc(size, spec)
|
|
except (RuntimeError, MemoryError): self.free_cache()
|
|
try: return self._alloc(size, 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, storage:BufferStorage, size:int, options:BufferSpec|None=None):
|
|
spec = options if options is not None else self.default_buffer_spec
|
|
if LRU and self.lru and not spec.nolru and spec.external_ptr is None: self.cache[(size, options)].append(storage)
|
|
else: self.do_free(storage, spec)
|
|
|
|
def free_cache(self):
|
|
for (_, options), storages in self.cache.items():
|
|
for storage in storages: self.do_free(storage, options if options is not None else self.default_buffer_spec)
|
|
storages.clear()
|
|
|
|
def do_free(self, storage:BufferStorage, options:BufferSpec):
|
|
for dev in storage.maps: dev.synchronize()
|
|
for dev, mb in storage.maps.items(): dev.allocator._unmap(mb)
|
|
if options.external_ptr is None: self._free(storage, options)
|
|
|
|
def map(self, buf:Buffer) -> BufferStorage: return self._map(buf.ensure_allocated())
|
|
|
|
# implemented by the runtime
|
|
def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: raise NotImplementedError("need alloc")
|
|
def _free(self, storage:BufferStorage, options:BufferSpec): pass # if opaque is a Python object, you don't need a free
|
|
def _copyin(self, dest, src:memoryview): raise NotImplementedError("need copyin")
|
|
def _copyout(self, dest:memoryview, src): raise NotImplementedError("need copyout")
|
|
def _map(self, buf) -> BufferStorage: raise NotImplementedError("need map")
|
|
def _unmap(self, mb): pass # default no-op; override if _map allocates iface-side state
|
|
def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
|
|
|
|
class HostAllocator(Allocator):
|
|
def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
|
|
if options.external_ptr is not None: view, meta = self._view(options.external_ptr, size), None
|
|
elif (remote:=getattr(self.dev, "remote", None)) is not None: view, meta = remote.alloc_sysmem(round_up(size, mmap.PAGESIZE))
|
|
else: view = self._view(mv_address(meta:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE)), size)
|
|
return BufferStorage(view.addr, meta, view)
|
|
|
|
def _free(self, storage:BufferStorage, options:BufferSpec):
|
|
if (remote:=getattr(self.dev, "remote", None)) is not None:
|
|
self.dev.synchronize()
|
|
remote.free_sysmem(storage.host)
|
|
|
|
def _view(self, addr:int, size:int) -> MMIOInterface:
|
|
return remote.cpu_view(addr, size) if (remote:=getattr(self.dev, "remote", None)) is not None else MMIOInterface(addr, size, fmt='B')
|
|
|
|
def _copyin(self, dest:int, src:memoryview):
|
|
self.dev.synchronize()
|
|
with cpu_profile(f"TINY -> {self.dev.device}", f"{self.dev.device}:COPY"): self._view(dest, src.nbytes)[:] = src.cast('B')
|
|
def _copyout(self, dest:memoryview, src:int):
|
|
self.dev.synchronize()
|
|
with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): dest[:] = self._view(src, dest.nbytes)[:]
|
|
def _map(self, buf:Buffer) -> BufferStorage:
|
|
if Device[buf.device].host != self.dev.host: raise RuntimeError(f"host memory is not on the node of {self.dev.device}")
|
|
return BufferStorage(buf.host.addr)
|
|
def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
|
|
|
|
# **************** for Compiled Devices ****************
|
|
|
|
class CompileError(Exception): pass
|
|
|
|
class Compiler:
|
|
def __init__(self, cachekey:str|None=None): self.cachekey = cachekey if CCACHE else None
|
|
def compile(self, src:str) -> bytes: return src.encode() # NOTE: empty compiler is the default
|
|
def compile_cached(self, src:str) -> bytes:
|
|
if self.cachekey is None or (lib := diskcache_get(self.cachekey, src)) is None:
|
|
assert not getenv("ASSERT_COMPILE"), f"tried to compile with ASSERT_COMPILE set\n{src}"
|
|
lib = self.compile(src)
|
|
if self.cachekey is not None: diskcache_put(self.cachekey, src, lib)
|
|
return lib
|
|
def disassemble(self, lib:bytes): pass
|
|
def server(self, cmd:str, arch:str, *args) -> subprocess.Popen:
|
|
argv = f"{cmd} {pathlib.Path(__file__).parent}/runtime/support/compileserver.py {type(self).__module__}:{type(self).__name__} {arch}"
|
|
return subprocess.Popen(argv.split() + [str(a) for a in args], stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
|
|
def compile_server(self, src:str, proc:subprocess.Popen) -> bytes:
|
|
unwrap(proc.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
|
|
if (lib:=unwrap(proc.stdout).read(struct.unpack("I", unwrap(proc.stdout).read(4))[0])): return lib
|
|
raise CompileError("Compilation Error")
|
|
|
|
|
|
@dataclass
|
|
class TinyELF:
|
|
lib: bytes
|
|
name: str
|
|
target: Target
|
|
# tuple of (name, slot, dtype, shape)
|
|
signature: tuple[tuple[str|None, int, DType, tuple], ...]
|
|
profile_key: bytes|None = None
|
|
|
|
@staticmethod
|
|
def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:
|
|
for _,_,dt,_ in signature:
|
|
yield (offset:=round_up(offset, dt.itemsize)), dt
|
|
offset += dt.itemsize
|
|
|
|
class Program(Generic[DeviceType]):
|
|
def __init__(self, dev:DeviceType, obj:TinyELF): pass
|
|
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(),
|
|
wait=False) -> float|None: pass
|
|
|
|
class Compiled:
|
|
ifaces:list[Callable] = []
|
|
profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
|
|
|
|
timestamp_divider: float = 1000.0
|
|
wait_timeout_ms: float = 30000.0
|
|
sleep_timeout_ms: int|None = None
|
|
can_recover:bool = False
|
|
rtalloc_size:int = 64<<20 # the pool every per-linear buffer is carved out of
|
|
var_vals: dict[str, int] = {}
|
|
|
|
# hcq2
|
|
pm_batch:Any = None
|
|
pm_encode:Any = None
|
|
pm_lower:Any = None
|
|
pm_bufferize:Any = None # one for all devices: each adds the rules of the placeholders it owns, its tags start with its name
|
|
|
|
def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, arch=None):
|
|
from tinygrad.renderer import Renderer
|
|
|
|
self.device, self.allocator, self.runtime_t, self.renderers = device, allocator, runtime, renderers or [Renderer]
|
|
self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
|
|
self.peer_group = getattr(getattr(self, 'iface', None), 'peer_group', device.split(":")[0])
|
|
self.cached_renderer:dict[Any, Renderer] = {}
|
|
|
|
# profiling
|
|
self.prof_ents:dict[tuple[Buffer, int], ProfileGraphEntry] = {} # (a batch's timestamps, start slot) -> entry, read at synchronize
|
|
|
|
@property
|
|
def has_copy_queue(self) -> bool: return True
|
|
|
|
@property
|
|
def host(self) -> str: return f"CPU:{self.peer_group[7:]}" if self.peer_group.startswith("remote:") else HCQ_RUNTIME_DEV.device
|
|
|
|
@property
|
|
def renderer(self) -> Renderer: return self._select_renderer()
|
|
|
|
@property
|
|
def compiler(self) -> Compiler: return self.renderer.compiler
|
|
|
|
def runtime(self, obj:TinyELF) -> Program[Self]: return unwrap(self.runtime_t)(self, obj)
|
|
|
|
@functools.cache
|
|
def rt_allocator(self, spec:BufferSpec) -> BumpAllocator: return BumpAllocator(self.rtalloc_size)
|
|
|
|
@functools.cache
|
|
def rt_buffer(self, spec:BufferSpec) -> Buffer:
|
|
return Buffer(self.device, self.rt_allocator(spec).size, options=spec, preallocate=True)
|
|
|
|
def tag(self, *parts:str) -> str: return to_name(self.device, *parts) # of the placeholders it owns
|
|
|
|
@functools.cached_property
|
|
def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]
|
|
return Buffer(self.device, 16, options=BufferSpec(host=True, uncached=True, cpu_access=True), initial_value=bytes(16))
|
|
|
|
@functools.cached_property
|
|
def error_state(self) -> Buffer: return Buffer(self.host, 8, options=BufferSpec(nolru=True), preallocate=True)
|
|
|
|
def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
|
|
timeout = timeout if timeout is not None and self.can_recover else None
|
|
st, done, err = time.perf_counter(), sig[0], self.error_state.host.view(fmt='q')
|
|
while done < value and not err[0]:
|
|
if done != (done:=sig[0]): st = time.perf_counter()
|
|
elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: raise RuntimeError(f"{self.device} signal wait timed out")
|
|
elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
|
|
if err[0]: raise RuntimeError(f"{self.device} failed with {err[0]}")
|
|
|
|
def synchronize(self, timeout:int|None=None):
|
|
try: self._wait_signal(tl:=self.timeline.host.view(fmt='Q'), tl[1], timeout)
|
|
except RuntimeError:
|
|
self.on_device_hang()
|
|
raise
|
|
if self.prof_ents: self.collect_prof()
|
|
|
|
def count(self) -> int:
|
|
"""
|
|
Returns the number of physical accelerators available to the runtime.
|
|
"""
|
|
return self.iface.count if hasattr(self, 'iface') else 1
|
|
|
|
def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
|
|
|
|
def on_sleep(self):
|
|
if (iface:=getattr(self, "iface", None)) is not None and hasattr(iface, "sleep"): iface.sleep(self.sleep_timeout_ms)
|
|
|
|
def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
|
|
|
|
def finalize(self):
|
|
"""
|
|
Called at the end of process lifetime to allow the device to finalize.
|
|
"""
|
|
try: self.synchronize() # try to finalize the device in any case
|
|
except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
|
|
if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
|
|
|
|
# helpers
|
|
|
|
def _renderer_name(self, r:type[Renderer]) -> str:
|
|
return r.__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
|
|
|
|
def _select_renderer(self, cv=DEV) -> Renderer:
|
|
assert (rn:=next((self._renderer_name(r) for r in self.renderers if getenv(f"{self.device}_{self._renderer_name(r)}")), None)) is None, \
|
|
f"{self.device}_{rn}=1 is deprecated, use DEV={self.device}:{rn} instead"
|
|
t = cv.target(self.device.split(':')[0], **({"arch":self.arch} if self.arch else {}))
|
|
return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
|
|
f"No renderer for {self.device} is available", self.cached_renderer, t)
|
|
|
|
def _select_iface(self, device:str):
|
|
self.device_id = int(device.split(":")[1]) if ":" in device else 0
|
|
assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "", \
|
|
f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
|
|
t = DEV.target(dev:=type(self).__name__[:-6])
|
|
filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
|
|
filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
|
|
return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
|
|
f"No interface for {dev}:{self.device_id} is available")
|
|
|
|
# profiling
|
|
|
|
def collect_prof(self):
|
|
div = decimal.Decimal(self.timestamp_divider)
|
|
for buf, g in itertools.groupby(self.prof_ents.items(), key=lambda x: x[0][0]) if PROFILE else []:
|
|
Compiled.profile_events.append(ProfileGraphEvent([e for _, e in g], [], [t/div for t in buf.host.view(fmt='Q')[:]]))
|
|
self.prof_ents.clear()
|
|
|
|
def _at_profile_finalize(self):
|
|
if not self.pm_encode.patterns: return
|
|
from tinygrad.tensor import Tensor
|
|
tdiffs = []
|
|
for _ in range(5):
|
|
with Context(DEBUG=0, BEAM=0, TRACK_MATCH_STATS=0): Tensor.ones(1, device=self.device).contiguous().realize()
|
|
if not (ents:=list(self.prof_ents.items())): return
|
|
self.prof_ents.clear()
|
|
st = perf_counter_us()
|
|
self.synchronize()
|
|
gpu = max(buf.host.view(fmt='Q')[e.en_id] for (buf, _), e in ents)/decimal.Decimal(self.timestamp_divider)
|
|
tdiffs.append((st+perf_counter_us())/2 - gpu)
|
|
Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))
|
|
|
|
if PROFILE:
|
|
@atexit.register
|
|
def finalize_profile():
|
|
devs = [Device[d] for d in Device._opened_devices]
|
|
for dev in devs: dev.synchronize()
|
|
for dev in devs: dev._at_profile_finalize()
|
|
|
|
with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
|
|
|
|
PROFILE.value = 0
|
|
from tinygrad.uop.ops import launch_viz
|
|
launch_viz("PROFILE", fn)
|
|
|
|
def enumerate_devices_str() -> Generator[str, None, None]:
|
|
from tinygrad import Tensor, Device
|
|
|
|
for device in ALL_DEVICES:
|
|
ren_results, iface_results = [], []
|
|
try:
|
|
d = Device[device]
|
|
for iface in [i for i in d.ifaces if not i.__name__.startswith("MOCK")]:
|
|
try:
|
|
name = iface.__name__[:-5]
|
|
default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore
|
|
iface_results.append(f"{colored('+', 'green')} {name}: {pluralize('device', count)} {default_text}")
|
|
except Exception as e: iface_results.append(f"{colored('-', 'red')} {iface.__name__[:-5]}: {e}")
|
|
for r in d.renderers:
|
|
try:
|
|
with Context(CACHELEVEL=0, DEV=f"{device}:{d._renderer_name(r)}"): test = (Tensor([1,2,3], device=device) * 2).tolist()
|
|
if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
|
|
default_text = '(default)' if type(d.renderer) is r else f'(DEV={device}:{d._renderer_name(r)} to make default)'
|
|
ren_results.append(f"{colored('+', 'green')} {d._renderer_name(r)} {default_text}")
|
|
except Exception as e: ren_results.append(f"{colored('-', 'red')} {d._renderer_name(r)}: {e}")
|
|
result = (colored('PASS', 'green') + ("\n"+" "*12+"interfaces:\n" if iface_results else "") + '\n'.join([" "*13+x for x in iface_results]) +
|
|
(("\n"+" "*12+"renderers:\n") + '\n'.join([" "*13+x for x in ren_results]) if len(ren_results) > 1 else ""))
|
|
except Exception as e: result = f"{colored('FAIL', 'red')} {e}"
|
|
yield f"{'*' if device == Device.DEFAULT else ' '} {device:8s}: {result}"
|
|
|
|
if __name__ == "__main__":
|
|
for s in enumerate_devices_str(): print(s)
|