openpilot v0.9.6 release

date: 2024-01-12T10:13:37
master commit: ba792d576a49a0899b88a753fa1c52956bedf9e6

tinygrad_repo/tinygrad/runtime/lib.py

@@ -0,0 +1,111 @@
+import ctypes
+import numpy as np
+from collections import defaultdict, deque
+from typing import TypeVar, Type, Any, Dict, Deque, Tuple
+from tinygrad.helpers import DType, dtypes, prod, GlobalCounters, ImageDType
+
+_T = TypeVar("_T")
+class RawBuffer:  # pylint: disable=abstract-method
+  def __init__(self, size:int, dtype:DType, buf:Any=None, allocator:Any=None, **kwargs):
+    self.size: int = size
+    self.dtype: DType = dtype
+    self._buf = buf if buf is not None else (allocator.alloc(size, dtype, **kwargs) if allocator else None) # If buf is provided, use it. Otherwise try to allocate from the allocator.
+    self._memsz: int = size*dtype.itemsize
+    self._allocator = allocator
+    self._device = kwargs.get('device', None)
+    GlobalCounters.mem_used += self._memsz
+  def __del__(self):  # NOTE: if it fails on init (bad dtype), it won't have a _memsz
+    if hasattr(self, '_memsz'): GlobalCounters.mem_used -= self._memsz
+    if hasattr(self, '_allocator') and self._allocator: self._allocator.free(self._buf)
+  def __repr__(self): return f"buffer<{self.size}, {self.dtype}, {id(self)}>"
+  @property
+  def key(self): return (self.size, self.dtype)
+
+  # NOTE: this interface allows for 0 copy
+  @classmethod
+  def fromCPU(cls:Type[_T], x:np.ndarray) -> _T: raise NotImplementedError("must be implemented")
+  def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
+
+class RawBufferCopyIn(RawBuffer):
+  def _copyin(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented")
+
+  @classmethod
+  def fromCPU(cls, x:np.ndarray, **kwargs):
+    ret = cls(prod(x.shape), dtypes.from_np(x.dtype), **kwargs)
+    if x.size > 0: ret._copyin(x)
+    return ret
+
+class RawBufferMapped(RawBufferCopyIn):
+  def _buffer(self) -> memoryview: raise NotImplementedError("must be implemented")
+  # NOTE: this metadata prevents the backing buffer from being freed. hack can be removed with PEP688
+  def buffer_view(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=np.dtype(self.dtype.np, metadata={"backing": self}), count=self.size)  # type: ignore
+  def toCPU(self) -> np.ndarray: return self.buffer_view().copy() # Need a copy, since jit will write to the same buffer.
+  def _copyin(self, x:np.ndarray) -> None: np.copyto(self.buffer_view(), x.reshape(-1))
+
+# this one is simple enough that i moved it out of the runtimes
+class RawMallocBuffer(RawBufferMapped):
+  def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float64:ctypes.c_double, dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.bfloat16: ctypes.c_int16, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int32: ctypes.c_int32, dtypes.uint32: ctypes.c_uint32, dtypes.int64: ctypes.c_int64, dtypes.uint64: ctypes.c_uint64, dtypes.int16: ctypes.c_int16, dtypes.uint16: ctypes.c_uint16}[dtype] * size)())
+  def _buffer(self): return memoryview(self._buf)
+
+class RawBufferCopyInOut(RawBufferCopyIn):
+  def _copyout(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented")
+
+  def toCPU(self) -> np.ndarray:
+    x: np.ndarray = np.empty(self.size, dtype=self.dtype.np)
+    if x.size > 0: self._copyout(x)
+    return x
+
+class RawBufferTransfer(RawBuffer):
+  def _transfer(self, x) -> None: raise NotImplementedError("must be implemented")
+
+  @classmethod
+  def transfer(cls, x, shape, dtype, **kwargs):
+    ret = cls(prod(shape), dtype, **kwargs)
+    ret._transfer(x)
+    return ret
+
+class LRUAllocator:
+  def __init__(self, dev_memsz=(4<<30)):
+    self.epoch = 0
+    self.free_space: Dict[Any, int] = defaultdict(lambda: dev_memsz)
+    self.buffer_info: Dict[Any, Tuple[int, DType, str]] = dict()
+    self.cached_buffers: Dict[Tuple[int, ...], Deque[Tuple[Any, int]]] = defaultdict(deque) # Cached buffer storage, splitted by type and size, newest first.
+    self.aging_order: Dict[Any, Deque[Tuple[Tuple[int, ...], int]]] = defaultdict(deque) # Keys of cached_buffers, ordered from oldest to newest updates.
+
+  def _cache_reuse_buffer(self, rawbufs: Deque[Tuple[Any, int]]): # The newest cached buffer is reused.
+    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.
+      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.free_space[device] -= size*dtype.itemsize
+    newbuf = self._do_alloc(max(1, size), dtype, device, **kwargs)
+    self.buffer_info[newbuf] = (size, dtype, device)
+    return newbuf
+
+  def _free_buffer(self, buf_to_free):
+    self.free_space[self.buffer_info[buf_to_free][2]] += self._underlying_buf_memsz(buf_to_free)
+    GlobalCounters.mem_cached -= self._underlying_buf_memsz(buf_to_free)
+    self.buffer_info.pop(buf_to_free)
+    self._do_free(buf_to_free)
+
+  def alloc(self, size, dtype, device='0', **kwargs):
+    rawbufs = self.cached_buffers.get(self._cached_bufkey(size, dtype, device), None)
+    return self._cache_reuse_buffer(rawbufs) if rawbufs else self._alloc_buffer(size, dtype, device, **kwargs)
+
+  def free(self, buf): # free() just caches buffer. It might be freed later when OOM during allocation.
+    self.epoch += 1
+    size, dtype, device = self.buffer_info[buf]
+    self.cached_buffers[self._cached_bufkey(size, dtype, device)].appendleft((buf, self.epoch))
+    self.aging_order[device].append((self._cached_bufkey(size, dtype, device), self.epoch))
+    GlobalCounters.mem_cached += self._underlying_buf_memsz(buf)
+
+  def _underlying_buf_memsz(self, buf): return self.buffer_info[buf][0] * self.buffer_info[buf][1].itemsize
+  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

tinygrad_repo/tinygrad/runtime/ops_cpu.py

@@ -0,0 +1,52 @@
+import numpy as np
+import operator
+from typing import Callable, Dict, Tuple, Optional
+from tinygrad.helpers import dtypes, DType
+from tinygrad.ops import BufferOps, UnaryOps, BinaryOps, MovementOps, ReduceOps, TernaryOps, Op, Interpreted
+from tinygrad.runtime.lib import RawBuffer
+
+def shape_to_axis(old_shape:Tuple[int, ...], new_shape:Tuple[int, ...]) -> Tuple[int, ...]:
+  assert len(old_shape) == len(new_shape), "reduce shapes must have same dimensions"
+  return tuple(i for i,(a,b) in enumerate(zip(old_shape, new_shape)) if a != b)
+
+base_fxn_for_op: Dict[Op, Callable] = {
+  BufferOps.MEM: lambda x: x._buf, UnaryOps.NEG: operator.neg, BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, BinaryOps.MUL: operator.mul, BinaryOps.DIV: operator.truediv,
+  ReduceOps.SUM: lambda x, new_shape: x.sum(shape_to_axis(x.shape, new_shape), keepdims=True) if tuple(x.shape) != tuple(new_shape) else x[:],
+  ReduceOps.MAX: lambda x, new_shape: (x.amax if hasattr(x, 'amax') else x.max)(shape_to_axis(x.shape, new_shape), keepdims=True) if tuple(x.shape) != tuple(new_shape) else x[:],
+  MovementOps.RESHAPE: lambda x, arg: x.reshape(arg), MovementOps.SHRINK: lambda x, arg: x[tuple(slice(p[0], p[1], None) for p in arg)],
+}
+
+def match_types(x, y):
+  up = x.dtype if dtypes.from_np(x.dtype).priority > dtypes.from_np(y.dtype).priority else y.dtype
+  return x.astype(up, copy=False), y.astype(up, copy=False)
+
+def einsum_mulacc(einsum, get_strides, expand):
+  def einscripts(x): return ''.join(["abcdefghijklmnopqrstuvwxyz"[i] for i in x])
+  def axes_slice(strides): return [i for i,s in enumerate(strides) if s != 0], tuple([slice(None) if s != 0 else 0 for i,s in enumerate(strides)])
+  def mulacc(a, b, new_shape):
+    (a_axes, a_slices), (b_axes, b_slices) = axes_slice(get_strides(a)), axes_slice(get_strides(b))
+    out = [i for i in range(len(new_shape)) if a.shape[i] == new_shape[i] and (i in a_axes or i in b_axes)]
+    ret = einsum(f"{einscripts(a_axes)}, {einscripts(b_axes)} -> {einscripts(out)}", a[a_slices], b[b_slices])
+    return expand(ret.reshape([(1 if i not in a_axes and i not in b_axes else s) for i,s in enumerate(new_shape)]), new_shape)
+  return mulacc
+
+numpy_fxn_for_op: Dict[Op, Callable] = {**base_fxn_for_op, **{
+  BufferOps.CONST: lambda val, dtype: np.array(val, dtype=dtype.np),
+  UnaryOps.NOOP: lambda x: np.require(x, requirements='C'), UnaryOps.EXP2: np.exp2, UnaryOps.LOG2: np.log2, UnaryOps.SIN: np.sin,
+  UnaryOps.CAST: lambda x,y: x.view(y[0].np) if y[1] else x.astype(y[0].np, copy=False),
+  BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: lambda x,y: (x<y).astype(np.promote_types(x.dtype,y.dtype)), BinaryOps.ADD: lambda x, y: np.add(*match_types(x, y)),
+  BinaryOps.SUB: lambda x, y: np.subtract(*match_types(x, y)), BinaryOps.MUL: lambda x, y: np.multiply(*match_types(x, y)),
+  BinaryOps.DIV: lambda x, y: np.divide(*match_types(x, y)), UnaryOps.SQRT: np.sqrt,
+  MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to,
+  MovementOps.STRIDE: lambda x, arg: x[tuple(slice(None, None, i) for i in arg)],
+  MovementOps.AS_STRIDED: lambda x, arg: np.ndarray(arg[0], buffer=np.require(x, requirements='C'), dtype=x.dtype, offset=arg[2]*x.dtype.itemsize, strides=tuple(y*x.dtype.itemsize for y in arg[1])),
+  TernaryOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, *match_types(a.copy(), b.copy()), optimize=True), lambda x: x.strides, np.broadcast_to),
+  TernaryOps.WHERE: np.where,
+}}
+
+class RawNumpyBuffer(RawBuffer):
+  def __init__(self, size:int, dtype:DType, buf:Optional[np.ndarray]=None): super().__init__(size, dtype, buf if buf is not None else np.empty([size], dtype.np))
+  @classmethod
+  def fromCPU(cls, x): return cls(x.size, dtypes.from_np(x.dtype), x)
+  def toCPU(self): return self._buf
+CPUBuffer = Interpreted(RawNumpyBuffer, numpy_fxn_for_op, from_underlying=RawNumpyBuffer.fromCPU)

tinygrad_repo/tinygrad/runtime/ops_disk.py

@@ -0,0 +1,41 @@
+import os, mmap
+from typing import Optional
+from typing import Callable, Dict, Tuple
+from tinygrad.helpers import prod, DType
+from tinygrad.runtime.lib import RawBufferMapped
+from tinygrad.ops import Interpreted, Op, MovementOps, UnaryOps, BufferOps
+
+class RawDiskBuffer(RawBufferMapped):
+  def __init__(self, size, dtype:DType, device:Optional[str]=None, buf=None, shape=None, offset=0):  # pylint: disable=super-init-not-called
+    self.shape = (size, ) if shape is None else shape
+    self.offset = offset  # this is an offset in bytes
+    assert device is not None or buf is not None, "disk tensor needs a path or a buf"
+    if device is not None:
+      f = open(device, "a+b")
+      if os.path.getsize(device) < size * dtype.itemsize: os.ftruncate(f.fileno(), size * dtype.itemsize)
+      buf = [f, mmap.mmap(f.fileno(), size * dtype.itemsize), 1]
+    else:
+      buf[2] += 1
+    # NOTE: we don't call super since disk tensors don't use RAM
+    self.size, self.dtype, self._buf = size, dtype, buf
+  def __del__(self):
+    self._buf[2] -= 1
+    if self._buf[2] == 0: self._buf[0].close()
+  def cast(self, arg:Tuple[DType, bool]): return RawDiskBuffer(self.size, arg[0], buf=self._buf, shape=self.shape, offset=self.offset)
+  def reshape(self, arg): return RawDiskBuffer(self.size, self.dtype, buf=self._buf, shape=arg, offset=self.offset)
+  def shrink(self, arg):
+    assert arg[1:] == tuple([(0,x) for x in self.shape[1:]]), f"can only slice the first dim of disk tensor {arg}"
+    offset = arg[0][0]*prod(self.shape[1:])*self.dtype.itemsize
+    size = (arg[0][1]-arg[0][0]) * prod(self.shape[1:])
+    return RawDiskBuffer(size, self.dtype, buf=self._buf, offset=self.offset+offset, shape=(arg[0][1]-arg[0][0],)+self.shape[1:])
+
+  def as_strided(self, arg):
+    return RawDiskBuffer(prod(arg[0]), self.dtype, buf=self._buf, offset=self.offset+arg[2]*self.dtype.itemsize, shape=arg[0])
+
+  def _buffer(self): return memoryview(self._buf[1])[self.offset:self.offset+self.size*self.dtype.itemsize]
+  def readinto(self, buf):
+    self._buf[0].seek(self.offset)
+    self._buf[0].readinto(buf)
+
+disk_fxn_for_op: Dict[Op, Callable] = { BufferOps.MEM: lambda x: x, UnaryOps.NOOP: lambda x: x, UnaryOps.CAST: RawDiskBuffer.cast, MovementOps.AS_STRIDED: RawDiskBuffer.as_strided }
+DiskBuffer = Interpreted(RawDiskBuffer, disk_fxn_for_op, from_underlying=lambda x:x)

tinygrad_repo/tinygrad/runtime/ops_gpu.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+import os
+os.environ['PYOPENCL_NO_CACHE'] = '1'
+import pathlib
+import numpy as np
+import pyopencl as cl  # type: ignore
+from typing import Optional, List, Tuple
+from tinygrad.helpers import DEBUG, getenv, prod, ImageDType, OSX, fromimport, diskcache
+from tinygrad.ops import Compiled
+from tinygrad.renderer.opencl import OpenCLRenderer
+from tinygrad.runtime.lib import RawBufferCopyInOut, LRUAllocator, RawBufferTransfer
+from tinygrad.codegen.kernel import LinearizerOptions
+
+OSX_TIMING_RATIO = (125/3) if OSX else 1.0   # see test/external/external_osx_profiling.py to determine this ratio. it's in like GPU clocks or something
+
+# TODO: if you fork and exit the child process after creating anything with cl on AMD, it hangs on e.wait()
+ROCM_LLVM_PATH = pathlib.Path("/opt/rocm/llvm/bin")
+#ROCM_LLVM_PATH = pathlib.Path(__file__).parents[3] / "extra/rocm/build/llvm-project/bin"
+if DEBUG >= 5:
+  early_exec = fromimport("extra.helpers", "enable_early_exec")()
+
+class CLAllocator(LRUAllocator):
+  def _do_alloc(self, size, dtype, device, **kwargs):
+    if isinstance(dtype, ImageDType):
+      # NOTE: the memory is a bit off here due to padding, it's buf.row_pitch * buf.height * 4 * dtype.itemsize
+      assert size == prod(dtype.shape), f"image size mismatch {size} != {dtype.shape}"
+      fmt = cl.ImageFormat(cl.channel_order.RGBA, {2: cl.channel_type.HALF_FLOAT, 4: cl.channel_type.FLOAT}[dtype.itemsize])
+      buf = cl.Image(CL.cl_ctxs[int(device)], cl.mem_flags.READ_WRITE, fmt, shape=(dtype.shape[1], dtype.shape[0]))
+    else:
+      buf = cl.Buffer(CL.cl_ctxs[int(device)], cl.mem_flags.READ_WRITE, size * dtype.itemsize)
+    setattr(buf, 'device', int(device)) # device is tracked on the underlying buffer
+    return buf
+
+class _CL:
+  def __init__(self):
+    cl_platforms = cl.get_platforms()
+    platform_devices: List[List[cl.Device]] = [y for y in ([x.get_devices(device_type=cl.device_type.GPU) for x in cl_platforms] + [x.get_devices(device_type=cl.device_type.CPU) for x in cl_platforms]) if y]
+    self.devices = [device for device in platform_devices[getenv('CL_PLATFORM', 0)] if device.name not in getenv('CL_EXCLUDE', "").split(",")]
+    self.cl_platform = self.devices[0].platform
+  def post_init(self, device=None):
+    self.cl_ctxs: List[cl.Context] = [cl.Context(devices=[x]) for x in self.devices] if device is None else [cl.Context(devices=[self.devices[device]])]
+    if DEBUG >= 1: print(f"using devices: {[ctx.devices[0].hashable_model_and_version_identifier for ctx in self.cl_ctxs]}")
+    self.cl_queue: List[cl.CommandQueue] = [cl.CommandQueue(ctx, device=ctx.devices[0], properties=cl.command_queue_properties.PROFILING_ENABLE) for ctx in self.cl_ctxs]
+    self.cl_allocator = CLAllocator(CL.cl_ctxs[0].devices[0].get_info(cl.device_info.GLOBAL_MEM_SIZE))
+  def synchronize(self):
+    for q in self.cl_queue: q.finish()
+CL = _CL()
+if not getenv("DELAYED_RUNTIME_INIT", False): CL.post_init()
+
+class CLBuffer(RawBufferCopyInOut, RawBufferTransfer):
+  def __init__(self, size, dtype, device='0'): super().__init__(size, dtype, allocator=CL.cl_allocator, **{'device': device})
+  def _copyin(self, x:np.ndarray):
+    assert not self.dtype.name.startswith("image"), f"can't copyin images {self.dtype}"
+    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)
+    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 []))
+  def _transfer(self, x):
+    if "gfx" in CL.cl_ctxs[x._buf.device].devices[0].name:
+      cl.enqueue_copy_buffer_p2p_amd(CL.cl_platform, CL.cl_queue[x._buf.device], x._buf, self._buf, x.size * x.dtype.itemsize).wait()
+    else: raise NotImplementedError("p2p transfer between devices not implemented on non-amd")
+
+@diskcache
+def compile_gpu(prg:str) -> bytes:
+  clprg = cl.Program(CL.cl_ctxs[0], prg)
+  clprg.build()
+  return clprg.get_info(cl.program_info.BINARIES)[0]
+
+class CLProgram:
+  def __init__(self, name:str, prg:bytes, argdtypes=None, options=None):
+    self.name, self.clprograms = name, [cl.Program(ctx, ctx.devices, [prg]*len(ctx.devices)) for ctx in CL.cl_ctxs]  # type: ignore
+    self._clprgs = [clprogram.build(options=options) for clprogram in self.clprograms]
+    self.clprgs = [clprg.__getattr__(name) for clprg in self._clprgs]
+    if DEBUG >= 5 and not OSX:
+      if 'Adreno' in CL.cl_ctxs[0].devices[0].name:
+        fromimport('disassemblers.adreno', 'disasm')(prg)
+      elif CL.cl_ctxs[0].devices[0].name.startswith('gfx'):
+        asm = early_exec(([ROCM_LLVM_PATH / "llvm-objdump", '-d', '-'], prg))
+        print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
+      else:
+        # print the PTX for NVIDIA. TODO: probably broken for everything else
+        print(prg.decode('utf-8'))
+    if argdtypes is not None: self.set_argdtypes(argdtypes)
+
+  def set_argdtypes(self, argdtypes): self.argdtypes, _ = argdtypes, [clprg.set_scalar_arg_dtypes(argdtypes) for clprg in self.clprgs]
+
+  @staticmethod
+  def max_work_group_size(): return CL.cl_ctxs[0].devices[0].max_work_group_size
+
+  def __call__(self, *bufs, global_size:Tuple[int,int,int], local_size:Optional[Tuple[int,int,int]]=None, wait=False) -> Optional[float]:
+    if not hasattr(self, 'argdtypes'): self.set_argdtypes(tuple(None if x.__class__ is CLBuffer else np.int32 for x in bufs))
+    cl_bufs, wait_for = [], []
+    for x in bufs:
+      if x.__class__ is CLBuffer:
+        cl_bufs.append(x._buf)
+        if hasattr(x, "event"): wait_for.append(x.event)
+      else: cl_bufs.append(x)
+    e = self.clprgs[cl_bufs[0].device](CL.cl_queue[cl_bufs[0].device], [int(g*l) for g,l in zip(global_size, local_size)] if local_size is not None else global_size, local_size, *cl_bufs, wait_for=wait_for)
+    if wait:
+      e.wait()
+      try:
+        return ((e.profile.end - e.profile.start) * OSX_TIMING_RATIO) * 1e-9
+      except cl.RuntimeError:   # no profiling info available
+        return None
+    return None
+
+GPUBuffer = Compiled(CLBuffer, LinearizerOptions(), OpenCLRenderer, compile_gpu, CLProgram, CL.synchronize)