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.mdtinygrad 的运行时即每个具体硬件后端的实现由以下四大部分组成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, graphNone, archNone):device设备名字符串可含:N形式的设备序号会被解析为device_idallocator该设备的Allocator实例renderers该设备可用的渲染器列表Renderer 负责把 UOps 渲染成目标语言源码runtime对应的Program子类负责加载并执行程序arch目标架构字符串例如 CPU 后端传入x86_64,native或arm64,native。Compiled还通过属性把 Renderer 与 Compiler 串联起来renderer属性会从renderers列表中按当前环境配置选择渲染器见_select_renderer而compiler属性则返回renderer.compilerruntime(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 时为 1finalize()在进程生命周期结束时被调用由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。分配行为通过BufferSpectinygrad/device.py精细控制其字段包括字段含义uncached是否绕过分配器缓存cpu_access缓冲区是否需要 CPU 可访问host是否为主机端缓冲区nolru是否不进入 LRU 缓存zero是否要求清零分配external_ptr外部指针如从别的框架借用内存Buffertinygrad/device.py在allocate()时调用allocator.alloc(self.nbytes, self.options)获得设备侧 opaque 句柄并同步更新GlobalCounters.mem_used等统计其get_buf方法展示了跨设备访问路径同一设备直接ensure_allocated()子缓冲区base 机制走allocator._offset其他设备则走allocator.map。LRUAllocator缓存分配提升性能LRUAllocatortinygrad/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)其中TinyELFtinygrad/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与rtlibgcc_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__中若是 LVPLVPRenderer 的 NIR 前端目标则按签名把缓冲区地址与标量值打包进参数区再调用否则直接以ctypes函数指针方式把各缓冲区虚拟地址和标量值传入waitTrue时返回耗时time.perf_counter差值清理__del__在 Windows 上调用VirtualFree释放内存。CPUDevicetinygrad/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(fno compiler for {self.device})。完整链路从张量到程序执行的运行时视角把四个部件串起来一次典型计算在运行时层的完整路径是Device[CPU]触发CPUDevice.__init__创建CPUAllocator并注册渲染器与CPUProgram初始化与装配阶段对应Compiled张量分配内存时Buffer.allocate调用allocator.allocLRUAllocator优先复用缓存内存管理阶段对应Allocator内核源码经 Renderer 渲染后由Compiler.compile_cached编译为二进制可能走子进程编译服务器结果封装为TinyELFCompiled.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.devicedevice.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），仅供参考