tinygrad 运行时（Runtime）架构深度解析：Compiled、Allocator、Program 与 Compiler 四大组件-港品优选

tinygrad 运行时（Runtime）架构深度解析：Compiled、Allocator、Program 与 Compiler 四大组件

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

导读

本篇基于 docs/developer/runtime.md 展开，系统讲解 tinygrad 中"运行时（Runtime）"的完整架构：一个典型运行时由Compiled（设备初始化与管理）、Allocator/LRUAllocator（设备内存管理）、Program（程序加载与执行）、Compiler（代码编译为设备二进制）四大部分组成。读完本文，你将掌握 tinygrad 设备抽象层的核心接口设计、各组件之间的协作关系，并能结合 tinygrad/device.py 与 tinygrad/runtime/ops_cpu.py 中的真实实现，理解如何为一个新硬件后端编写运行时。

Runtime Overview：一个典型运行时由哪些部分组成

根据 docs/developer/runtime.md，tinygrad 的运行时（即每个具体硬件后端的实现）由以下四大部分组成：

Compiled—— 负责初始化和管理一个设备（device），是所有后端设备类的基类；
Allocator—— 负责设备上内存的分配与释放，另有缓存分配缓冲区的LRUAllocator变体用于性能优化；
Program—— 为每个已加载的程序创建，负责在设备上执行该程序；
Compiler—— 将Renderer输出的源码编译为设备特定的二进制格式。

这四个组件的抽象基类都定义在 tinygrad/device.py 中，而每个具体后端则在 tinygrad/runtime/ 目录下以ops_<设备名>.py的形式提供各自的实现（如ops_cpu.py、ops_metal.py、ops_amd.py、ops_nv.py、ops_cuda.py等）。docs/developer/developer.md 中对这一层也有呼应：运行时负责设备特定的交互，包括初始化设备、分配内存、加载/启动程序等，所有运行时实现都位于 runtime 目录。

值得一提的是，tinygrad 的设备分发入口是Device单例（见 tinygrad/device.py）：它通过扫描 runtime 目录下所有ops_*.py文件自动发现可用设备，Device["CPU"]、Device["METAL"]这样的索引操作会返回一个Compiled实例，并记录到_opened_devices集合中，程序退出时统一调用各设备的finalize()。

Compiled：设备的初始化与生命周期管理

Compiled类是每个设备后端的基类（定义于 tinygrad/device.py），它"负责初始化和管理一个设备"。其构造函数接收以下关键参数：

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

device：设备名字符串，可含:N形式的设备序号（会被解析为device_id）；
allocator：该设备的Allocator实例；
renderers：该设备可用的渲染器列表（Renderer 负责把 UOps 渲染成目标语言源码）；
runtime：对应的Program子类，负责加载并执行程序；
arch：目标架构字符串，例如 CPU 后端传入"x86_64,native"或"arm64,native"。

Compiled还通过属性把 Renderer 与 Compiler 串联起来：renderer属性会从renderers列表中按当前环境配置选择渲染器（见_select_renderer），而compiler属性则返回renderer.compiler；runtime(obj)方法则用runtime_t构造一个Program实例来装载给定的TinyELF。

synchronize：同步所有挂起操作

Compiled对外暴露的最核心方法之一就是synchronize()，其 docstring 明确说明其职责：

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

基类中它只是一个可被覆写的占位方法（注释# override this in your device implementation），真实行为由各后端实现。例如：

tinygrad/runtime/ops_metal.py 中的MetalDevice.synchronize()等待命令缓冲区（MTLCommandBuffer）完成；
tinygrad/runtime/ops_cpu.py 中的CPUDevice.synchronize()会遍历所有已打开的设备，对每个HCQ2Compiled实例调用其synchronize(timeout)，以保证主机读取数据时所有设备时间线都已追上。

count 与 finalize：设备探测与收尾

Compiled.count()返回运行时可用的物理加速器数量（默认取iface.count，无 interface 时为 1）；finalize()在进程生命周期结束时被调用（由Device单例在atexit时统一触发），用于释放设备资源，若 interface 提供device_fini则调用之。此外_at_profile_finalize()会在 profiling 结束时被调用，用于设备侧性能数据的收尾（见 tinygrad/device.py 中PROFILE模式下注册的finalize_profile）。

Allocator 与 LRUAllocator：设备内存管理

Allocator 基类的接口契约

Allocator类（tinygrad/device.py）"负责管理设备上的内存"。它定义了两组方法：

对外公共接口：alloc(size, options)、free(opaque, size, options)、map(buf)。其中alloc断言 size 必须为正，并在分配失败（RuntimeError/MemoryError）时抛出带有设备名与已用内存统计（GlobalCounters.mem_used_per_device）的MemoryError，方便定位显存不足问题；
由运行时实现的内部钩子：_alloc、_free、_copyin、_copyout、_map、_unmap、_offset、_encode_decode。基类中这些方法要么抛出NotImplementedError("need alloc")等占位错误，要么是空操作（如_free注释说明：如果 opaque 是 Python 对象，就不需要显式 free）。

此外Allocator.__init__接收两个能力标志：supports_copy_from_disk（是否支持从磁盘拷贝）与supports_transfer（是否支持设备间传输），并持有default_buffer_spec（默认的BufferSpec）。

分配行为通过BufferSpec（tinygrad/device.py）精细控制，其字段包括：

字段	含义
`uncached`	是否绕过分配器缓存
`cpu_access`	缓冲区是否需要 CPU 可访问
`host`	是否为主机端缓冲区
`nolru`	是否不进入 LRU 缓存
`zero`	是否要求清零分配
`external_ptr`	外部指针（如从别的框架借用内存）

Buffer（tinygrad/device.py）在allocate()时调用allocator.alloc(self.nbytes, self.options)获得设备侧 opaque 句柄，并同步更新GlobalCounters.mem_used等统计；其get_buf方法展示了跨设备访问路径：同一设备直接ensure_allocated()，子缓冲区（base 机制）走allocator._offset，其他设备则走allocator.map。

LRUAllocator：缓存分配提升性能

LRUAllocator（tinygrad/device.py）是Allocator的缓存版，其 docstring 写道：

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

实现要点：

维护self.cache: dict[tuple[int, BufferSpec|None], list]，以(size, options)为键缓存空闲缓冲区；
alloc优先从缓存中pop复用；若分配失败（显存不足），先调用free_cache()清空整个缓存再重试一次，避免缓存占用导致分配失败；
free时并非真正释放：当环境变量LRU开启、且options允许（非nolru、非zero、无external_ptr）时，缓冲区被放入缓存而非释放；否则才调用super().free真正归还。

正是这套"能不释放就不释放"的策略，让频繁创建/销毁临时张量（如算子中间结果）时无需反复向驱动申请内存，从而显著降低运行时开销。各后端如 tinygrad/runtime/ops_metal.py 的MetalAllocator即继承自LRUAllocator。

CPUAllocator 实例：mmap 与 memmove

以 tinygrad/runtime/ops_cpu.py 的CPUAllocator为例，可以看到一个最小分配器的完整实现：

_alloc用mmap匿名映射分配可读写内存（Windows 走ACCESS_WRITE，类 Unix 走MAP_ANON | MAP_SHARED），并支持options.external_ptr直接借用外部指针；
_copyin/_copyout在同步后通过ctypes.memmove在主机内存与设备缓冲区之间搬运数据；
_do_map要求缓冲区带有MMIOInterface视图，返回一个包装句柄，_do_unmap则是空操作。

Program：程序加载与设备执行

Program是"为每个加载的程序创建的、负责在设备上执行程序"的类。其抽象基类（tinygrad/device.py）定义了两件事：构造函数接收(dev, obj: TinyELF)，其中TinyELF（tinygrad/device.py）封装了编译产物的lib字节、name、target与signature（签名描述了缓冲区与标量参数的布局）；__call__方法则按global_size、local_size、vals（标量参数）、wait执行程序。

文档中特别以CPUProgram作为示例实现，其源码位于 tinygrad/runtime/ops_cpu.py。这个实现颇具代表性，展示了"把编译产物装进可执行内存并调用"的完整链路：

加载与链接：_load方法识别 ELF 魔数（libc.ELFMAG），通过jit_loader加载并链接libm与rt（libgcc_s）两个运行时库；
可执行内存分配：Windows 上用VirtualAlloc(PAGE_EXECUTE_READWRITE)；macOS 上由于 SPRR（指针认证与写保护扩展）不允许 RWX 页，使用MAP_JIT标志分配内存，并在写入前调用pthread_jit_write_protect_np(False)关闭 JIT 写保护、写入后重新开启；
指令缓存刷新：通过__clear_cache（或msync(MS_SYNC | MS_INVALIDATE)兜底）确保 CPU 指令缓存看到新写入的机器码；
执行：__call__中若是 LVP（LVPRenderer 的 NIR 前端）目标，则按签名把缓冲区地址与标量值打包进参数区再调用；否则直接以ctypes函数指针方式把各缓冲区虚拟地址和标量值传入；wait=True时返回耗时（time.perf_counter差值）；
清理：__del__在 Windows 上调用VirtualFree释放内存。

CPUDevice（tinygrad/runtime/ops_cpu.py）则把这些部件组装起来：构造函数传入CPUAllocator、四个渲染器（ClangRenderer、CPULLVMRenderer、LVPRenderer、X86Renderer）与CPUProgram，并根据宿主机架构（amd64/aarch64 等）设置arch参数。

Compiler：把 Renderer 输出编译为设备二进制

Compiler类（tinygrad/device.py）"编译Renderer的输出，并产出设备特定格式"。它提供的成员包括：

cachekey：磁盘编译缓存键，仅当环境变量CCACHE开启时才启用；
compile(src)：把源码字符串编译为字节；基类默认实现是src.encode()（即"空编译器"，适用于不需要真正编译的渲染目标）；
compile_cached(src)：先查磁盘缓存（diskcache_get/diskcache_put），命中则直接返回，否则调用compile并写缓存；当设置了ASSERT_COMPILE环境变量时，任何实际编译尝试都会直接断言失败，便于测试确认缓存路径；
disassemble(lib)：反汇编产物（默认空操作）；
server(cmd, arch, *args)与compile_server(src, proc)：启动一个独立的编译服务器子进程（调用 tinygrad/runtime/support/compileserver.py），通过 stdin/stdout 传递源码、回收二进制——这是把重型编译器（如 LLVM/PTX 工具链）隔离到子进程、避免其污染主进程全局状态（例如 Metal 的 MTLCompiler 自带 LLVM，见 tinygrad/runtime/ops_metal.py）的关键机制。

在Compiled中，compiler属性即取自当前选中的renderer.compiler；若渲染器未声明编译器，会抛出RuntimeError(f"no compiler for {self.device}")。

完整链路：从张量到程序执行的运行时视角

把四个部件串起来，一次典型计算在运行时层的完整路径是：

Device["CPU"]触发CPUDevice.__init__，创建CPUAllocator并注册渲染器与CPUProgram（初始化与装配阶段，对应Compiled）；
张量分配内存时，Buffer.allocate调用allocator.alloc，LRUAllocator优先复用缓存（内存管理阶段，对应Allocator）；
内核源码经 Renderer 渲染后，由Compiler.compile_cached编译为二进制（可能走子进程编译服务器），结果封装为TinyELF；
Compiled.runtime(obj)构造CPUProgram，加载并链接二进制、映射可执行内存、刷新指令缓存；
调度执行时调用Program.__call__(*bufs, global_size, local_size, vals, wait)在设备上运行内核（执行阶段，对应Program）。

这一设计让每个新后端只需实现这四类组件即可接入 tinygrad：例如 tinygrad/runtime/ops_metal.py 就同时定义了MetalDevice(Compiled)、MetalCompiler(Compiler)、MetalProgram(Program)与MetalAllocator(LRUAllocator)。想快速验证当前环境的运行时能力，可以在仓库根目录直接运行python -m tinygrad.device（device.py末尾的__main__入口会打印每个设备的 interface 与 renderer 探测结果），或参考 test/backend/test_ops.py、test/device/ 下的设备专项测试来观察各运行时组件的实际行为。

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析