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,
    graph=None,
    arch=None,
)

Methods:

• synchronize –

  Synchronize all pending operations on the device.

synchronize

synchronize()

Synchronize all pending operations on the device.

This method ensures that all previously queued operations on the device have been completed before proceeding.

Allocator

The Allocator class is responsible for managing memory on the device. There is also a version called the LRUAllocator, which caches allocated buffers to optimize performance.

Allocator

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

Bases: Generic[DeviceType]

Methods:

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

Attributes:

• default_buffer_spec (BufferSpec) –
• dev (DeviceType) –

default_buffer_spec instance-attribute

default_buffer_spec: BufferSpec = BufferSpec()

dev instance-attribute

dev: DeviceType = dev

_alloc

_alloc(size: int, options: BufferSpec)

_copyin

_copyin(dest, src: memoryview)

_copyout

_copyout(dest: memoryview, src)

_encode_decode

_encode_decode(
    bufout,
    bufin,
    desc,
    hist: list,
    shape: tuple[int, ...],
    frame_pos: int,
)

_free

_free(opaque, options: BufferSpec)

_map

_map(buf)

_offset

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

_unmap

_unmap(mb)

alloc

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

free

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

map

map(buf: Buffer)

LRUAllocator

LRUAllocator(dev: DeviceType, **kwargs)

Bases: Allocator, Generic[DeviceType]

The LRU Allocator is responsible for caching buffers. It ensures that buffers are not freed until it is absolutely necessary, optimizing performance.

Methods:

• alloc –
• free –
• free_cache –

Attributes:

• cache (dict[tuple[int, BufferSpec | None], Any]) –

cache instance-attribute

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

alloc

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

free

free(
    opaque: Any,
    size: int,
    options: BufferSpec | None = None,
)

free_cache

free_cache()

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 –
• lvp –
• mem –

Source code in tinygrad/runtime/ops_cpu.py

def __init__(self, dev:CPUDevice, obj:TinyELF):
  self.dev, self.name, self.signature = dev, obj.name, obj.signature
  self.lvp = obj.target.renderer == "LVP"

  if 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 = (
    ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr)
    if self.lvp
    else ctypes.CFUNCTYPE(None)(self.addr)
)

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: HCQBuffer,
    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:HCQBuffer, 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:
  st = time.perf_counter()
  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, *[b.va_addr for b in 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:
    args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
    self.fxn(*[ctypes.c_uint64(x) for x in args])
  return time.perf_counter() - st if wait else None

__del__

__del__()

Source code in tinygrad/runtime/ops_cpu.py

@suppress_finalizing
def __del__(self):
  if 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=[self.libm, self.rt_lib])

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)