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)

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[tuple[BufferStorage, dict[Compiled, int]]]]) –
• dev (DeviceType) –
• lru –

cache instance-attribute

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

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__ –
• _load –

Attributes:

• buf –
• fxn (Any) –
• lvp –

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]
  else:
    self.buf = dev.allocator.alloc(len(obj.lib), BufferSpec(nolru=True))
    loaded = self._load(obj.lib, self.buf.buf)
    unwrap(self.buf.host)[:len(loaded)] = loaded

    # RWX pages are unrepresentable on apple silicon (https://blog.svenpeter.dev/posts/m1_sprr_gxf/), so flip the written pages to R-X
    libc.mprotect(ctypes.c_void_p(self.buf.buf), len(obj.lib), libc.PROT_READ | libc.PROT_EXEC)

    # __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.buf.buf), ctypes.c_void_p(self.buf.buf + len(loaded)))
    else:
      # msync should be a universal POSIX way to do this
      libc.msync(ctypes.c_void_p(self.buf.buf), len(loaded), libc.MS_SYNC | libc.MS_INVALIDATE)

    self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.buf.buf)

buf instance-attribute

buf = dev.allocator.alloc(
    len(obj.lib), BufferSpec(nolru=True)
)

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'

__call__

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

Source code in tinygrad/runtime/ops_cpu.py

def __call__(self, *args:int, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
             wait:bool=False, timeout:int|None=None) -> float|None:
  hdr = 12 if self.lvp else 0
  kernargs = ctypes.create_string_buffer(hdr + len(args) * 8)
  if self.lvp: struct.pack_into('<3I', kernargs, 0, *data64_le(ctypes.addressof(kernargs) + hdr), len(args) * 2)
  for a, (off, dt) in zip(args, TinyELF.iter_sig(self.signature)): struct.pack_into(f'<{dt.fmt}', kernargs, hdr + off, a)
  if (remote:=self.dev.remote) is not None:
    ret = (remote._rpc if wait else remote._post)(remote.sock, RemoteCmd.EXEC_PROG, self.fxn, len(kernargs), int(wait), payload=bytes(kernargs))
    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: self.fxn(kernargs)
  return float(unwrap(prof.en) - prof.st) * 1e-6 if wait else None

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