Runtime Overview

Overview

A typical runtime consists of the following parts:

• Compiled
• Allocator
• Program
• Compiler

Compiled

The Compiled class is responsible for initializing and managing a device.

Compiled

Compiled(
    device: str,
    allocator: Allocator,
    renderers: list[type[Renderer]],
    runtime: type[Program[Self]] | None,
    arch=None,
)

Methods:

• synchronize –

synchronize

synchronize(timeout: int | None = None)

Allocator

The Allocator class manages memory on the device and caches allocated buffers for reuse.

Allocator

Allocator(
    dev: DeviceType,
    supports_copy_from_disk: bool = True,
    supports_transfer: bool = True,
)

Bases: Generic[DeviceType]

Methods:

• _alloc –
• _copyin –
• _copyout –
• _free –
• _map –
• _offset –
• _unmap –
• alloc –
• do_free –
• free –
• free_cache –
• map –

Attributes:

• cache (dict[tuple[int, BufferSpec | None], list[BufferStorage]]) –
• default_buffer_spec (BufferSpec) –
• dev (DeviceType) –
• lru –

cache instance-attribute

cache: dict[
    tuple[int, BufferSpec | None], list[BufferStorage]
] = defaultdict(list)

default_buffer_spec instance-attribute

default_buffer_spec: BufferSpec = BufferSpec()

dev instance-attribute

dev: DeviceType = dev

lru class-attribute instance-attribute

lru = True

_alloc

_alloc(size: int, options: BufferSpec) -> BufferStorage

_copyin

_copyin(dest, src: memoryview)

_copyout

_copyout(dest: memoryview, src)

_free

_free(storage: BufferStorage, options: BufferSpec)

_map

_map(buf) -> BufferStorage

_offset

_offset(buf, size: int, offset: int)

_unmap

_unmap(mb)

alloc

alloc(
    size: int, options: BufferSpec | None = None
) -> BufferStorage

do_free

do_free(storage: BufferStorage, options: BufferSpec)

free

free(
    storage: BufferStorage,
    size: int,
    options: BufferSpec | None = None,
)

free_cache

free_cache()

map

map(buf: Buffer) -> BufferStorage

Program

The Program class is created for each loaded program. It is responsible for executing the program on the device. As an example, here is a CPUProgram implementation which loads program and runs it.

CPUProgram

CPUProgram(dev: CPUDevice, obj: TinyELF)

Bases: Program['CPUDevice']

Methods:

• __call__ –
• __del__ –
• _load –

Attributes:

• addr –
• fxn (Any) –
• lvp –
• mem –

Source code in tinygrad/runtime/ops_cpu.py

def __init__(self, dev:CPUDevice, obj:TinyELF):
  self.dev, self.name, self.signature, self.profile_key = dev, obj.name, obj.signature, obj.profile_key
  self.lvp = obj.target.renderer == "LVP"
  if dev.remote is not None: self.fxn:Any = dev.remote.rpc(RemoteCmd.LOAD_PROG, len(obj.lib), payload=obj.lib)[0]
  elif sys.platform == "win32": # mypy doesn't understand when WIN is used here
    PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
    ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
    self.addr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(obj.lib)), MEM_COMMIT | MEM_RESERVE,
                                                    PAGE_EXECUTE_READWRITE)
    ctypes.memmove(self.addr, (loaded:=self._load(obj.lib, self.addr)), len(loaded))
    ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
    proc = ctypes.windll.kernel32.GetCurrentProcess()
    ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(loaded)))
    self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
  else:
    # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
    # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
    self.mem = mmap.mmap(-1, len(obj.lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
    self.addr = mv_address(self.mem)

    if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
    self.mem.write(loaded:=self._load(obj.lib, mv_address(self.mem)))
    if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)

    # __clear_cache isn't a normal libc function, but a compiler support routine found in libgcc_s for gcc and compiler-rt for clang.
    # libgcc_s comes as shared library but compiler-rt is only a bunch of static library archives which we can't directly load, but fortunately
    # it somehow found its way into libSystem on macos (likely because it used __builtin_clear_cache) and libgcc_s is ~always present on linux
    # Using ["name"] instead of .name because otherwise name is getting mangled: https://docs.python.org/3.12/reference/expressions.html#index-5
    if 'rt' in DLL._loaded_: CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(self.addr), ctypes.c_void_p(self.addr + len(loaded)))
    else:
      # msync should be a universal POSIX way to do this
      libc.msync(ctypes.c_void_p(self.addr), len(loaded), libc.MS_SYNC | libc.MS_INVALIDATE)

    self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)

addr instance-attribute

addr = ctypes.windll.kernel32.VirtualAlloc(
    ctypes.c_void_p(0),
    ctypes.c_size_t(len(obj.lib)),
    MEM_COMMIT | MEM_RESERVE,
    PAGE_EXECUTE_READWRITE,
)

fxn instance-attribute

fxn: Any = dev.remote.rpc(
    RemoteCmd.LOAD_PROG, len(obj.lib), payload=obj.lib
)[0]

lvp instance-attribute

lvp = obj.target.renderer == 'LVP'

mem instance-attribute

mem = mmap.mmap(
    -1,
    len(obj.lib),
    mmap.MAP_ANON
    | mmap.MAP_PRIVATE
    | (MAP_JIT if OSX else 0),
    mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC,
)

__call__

__call__(
    *bufs: int,
    global_size: tuple[int, int, int] = (1, 1, 1),
    local_size: tuple[int, int, int] = (1, 1, 1),
    vals: tuple[int | None, ...] = (),
    wait: bool = False,
    timeout: int | None = None
) -> float | None

Source code in tinygrad/runtime/ops_cpu.py

def __call__(self, *bufs:int, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
             vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
  args = [*bufs, *cast(tuple[int, ...], vals)]
  if (remote:=self.dev.remote) is not None:
    data = struct.pack(f'<{len(args)}Q', *(a & 0xffffffffffffffff for a in args))
    ret = (remote._rpc if wait else remote._post)(remote.sock, RemoteCmd.EXEC_PROG, self.fxn, len(args), int(wait), payload=data)
    return ret[0] / 1e9 if ret is not None else None
  with cpu_profile(self.name, self.dev.device, profile_key=self.profile_key) as prof:
    if self.lvp:
      lvp_args = bytearray(12 + (len(bufs) + len(vals)) * 8)
      addr = mv_address(lvp_args)
      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *bufs)
      for v,(off,dt) in zip(vals, TinyELF.iter_sig(self.signature[-len(vals):], len(bufs)*8)): struct.pack_into(f'<{dt.fmt}', lvp_args, 12+off, v)
      self.fxn(addr)
    else: self.fxn(*[ctypes.c_uint64(x) for x in args])
  return float(unwrap(prof.en) - prof.st) * 1e-6 if wait else None

__del__

__del__()

Source code in tinygrad/runtime/ops_cpu.py

@suppress_finalizing
def __del__(self):
  if self.dev.remote is None and sys.platform == 'win32':
    ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.addr), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE

_load

_load(lib, base=0)

Source code in tinygrad/runtime/ops_cpu.py

def _load(self, lib, base=0): return lib if lib[:4] != libc.ELFMAG.encode() else jit_loader(lib, base=base, link_libs=list(DLL._loaded_.values()))

Compiler

The Compiler class compiles the output from the Renderer and produces it in a device-specific format.

Compiler

Compiler(cachekey: str | None = None)

Methods:

• compile –
• compile_cached –
• compile_server –
• disassemble –
• server –

Attributes:

• cachekey –

Source code in tinygrad/device.py

def __init__(self, cachekey:str|None=None): self.cachekey = cachekey if CCACHE else None

cachekey instance-attribute

cachekey = cachekey if CCACHE else None

compile

compile(src: str) -> bytes

Source code in tinygrad/device.py

def compile(self, src:str) -> bytes: return src.encode()   # NOTE: empty compiler is the default

compile_cached

compile_cached(src: str) -> bytes

Source code in tinygrad/device.py

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

compile_server

compile_server(src: str, proc: Popen) -> bytes

Source code in tinygrad/device.py

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

disassemble

disassemble(lib: bytes)

Source code in tinygrad/device.py

def disassemble(self, lib:bytes): pass

server

server(cmd: str, arch: str, *args) -> Popen

Source code in tinygrad/device.py

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)