tinygrad 可扩展流水线实战如何像加中间件一样给编译器装一个 pass【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradtinygrad 是一套读得下去的深度学习栈Tensor、自动求导、调度、代码生成、渲染、运行时一层层叠起来。它的核心特点是可插拔——流水线里的每一段都可以单独换掉而不必动其余部分。下面用一个真实任务带你把它拆开看先跑起来再讲清每一段为什么能替换然后动手给流水线装一个自己的处理规则最后说说换后端和排错。一、一分钟跑起来先装。tinygrad 推荐从源码装仓库地址是 https://gitcode.com/GitHub_Trending/tiny/tinygrad 拉下来之后git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad cd tinygrad python3 -m pip install -e .然后写最小的张量计算from tinygrad import Tensor x Tensor([1, 2, 3, 4]) print((x * x).sum().item()) # 30.item()会强制把结果真正算出来也叫 realize。这就是 tinygrad 的懒执行你写的x*x和.sum()此刻只搭了一张图还没动手算直到取结果才触发整条流水线。想看它怎么算把DEBUG调高就行——DEBUG3会打印融合后的 kernelDEBUG4还会把生成出来的源码直接甩给你。DEBUG3 python3 -c from tinygrad import Tensor; N1024; a,bTensor.empty(N,N),Tensor.empty(N,N); (a.reshape(N,1,N)*b.T.reshape(1,N,N)).sum(axis2).realize()这一行的价值在于你看到的是一个融合好的 kernel而不是一串碎操作。懒执行 融合是后面所有可扩展故事的地基。二、它是怎么可插拔的一条五段式流水线把一次realize从后往前拆开数据在五个可替换的工位之间流动靠一个共同的中间结构UOp传递Tensor 前端你写的x*y、.sum()等被拆成 UOp 树同时记录自动求导信息。调度schedule把 UOp 树整理成先算哪个、后算哪个的 kernel 依赖图再做拓扑排序。代码生成codegen一串重写规则依次改写 UOp做融合、简化、展开、分配 GPU 维度——这一段最像中间件链。渲染renderer把改写完的 UOp 变成目标架构的源码文本WGSL、C/Clang、PTX、LLVM IR……。运行时runtime真正分配显存、编译、把 kernel 丢到设备上跑。# 概念示意同一条计算按 DEV 变量切到不同后端 from tinygrad import Tensor, Device print(Device.DEFAULT) # 自动选一个可用设备 Tensor([1,2,3]).sum().realize()可插拔具体体现在两个地方也是本仓库最容易下手改的地方运行时按文件自动发现。Device 单例 启动时扫一遍tinygrad/runtime/下所有ops_*.py看到ops_metal.py就支持 METAL看到ops_cpu.py就支持 CPU。你想接新硬件丢一个ops_xxx.py进去它就会被发现。渲染器是一个列表。每个设备在初始化时交出一个 Renderer 列表框架按名字挑一个。比如 CPU 后端交出[ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer]。换目标架构 换列表里的渲染器或往里加一个新的。Renderer 基类 定义了几件事把 UOp 列表渲染成源码字符串render、把源码汇编成二进制asm、每个算子怎么落地code_for_op字典、支持哪些数据类型。一个具体的 C 风格渲染器只要填一张UOp 子树 → 源码片段的匹配表就能把同一份 UOp 变成不同方言。三、实战给流水线加一个自己的 pass写第一个中间件前面提到代码生成是一串重写规则的链这就是 tinygrad 版的中间件每个规则是一个模式改写函数对叫PatternMatcher。它扫描 UOp 树凡匹配到某个模式就交给你的函数改写成别的样子然后下一个规则接着处理。先看仓库里一个真实、最小的例子来自 编译流水线def do_number_param(ctx, x): ctx[0] 1 return x.replace(argreplace(x.arg, slotctx[0]-1)) pm_number_params PatternMatcher([ (UPat(Ops.PARAM, namex), do_number_param), ])读法很直白UPat(Ops.PARAM, namex)是找所有 PARAM 节点do_number_param是给匹配到的节点重新编号。把它挂进链条只需一行graph_rewrite(sink, pm_number_params, ...)——相当于往中间件链里塞了个拦截器其余规则对它无感。自己动手写一个套路固定想清我要把哪种 UOp 子树变成什么。用UPat把目标模式画出来写成(模式, 回调)元组。包进PatternMatcher在合适的位置graph_rewrite一次。这就是 tinygrad 里中间件的完整形态无状态、按模式触发、只改自己匹配的那一小块树。因为规则之间不共享顺序状态你加一条不会破坏别的规则天然解耦、可复用——跟洋葱模型里每一层只关心自己前后两件事是一个道理。四、换后端、换渲染器换个目标架构装 pass 是改逻辑换后端是改落地。两者都靠列表和文件约定不用改框架选设备用环境变量DEV指定例如DEVCPU、DEVMETAL。同一设备还能选渲染器如DEVCPU:LLVM、DEVNV:PTX。加渲染器继承 cstyle 渲染器 那套基类填好code_for_op、type_map和一张字符串匹配表就能把 UOp 渲染成你的方言再把它加进对应设备的 Renderer 列表即可。加运行时在 各设备运行时目录 里加ops_xxx.py实现分配器管显存、程序加载、以及设备初始化时交出渲染器列表。设备发现机制会自动把它纳入。换句话说tinygrad 把可插拔做成了三件事的组合文件约定发现 列表装配 基类契约。你只在契约允许的位置填内容其余流水线照旧工作。五、卡住了怎么办算出来是 0 或没动静多半没触发 realize。给个.item()、.numpy()或.realize()强制落地再看输出。不知道在哪个设备上跑print(Device.DEFAULT)想看每个后端是否可用跑python3 tinygrad/device.py会列出 PASS/FAIL 清单。报错说没有 renderer通常是DEV写成了当前机器没有的组合换成列表里标了(default)的那个或确认对应编译器如 CPU 需要clang在 PATH 里。想亲眼追踪每一步DEBUG3/4看融合和源码开VIZ1能可视化整张 UOp 图在各阶段前后的样子对定位是哪条规则改坏的特别有用。改完逻辑结果对不上先确认你是不是在不该改的位置插了 pass。重写规则是有先后顺序的放错位置可能拿到还没简化完的树。排错口诀先看在哪段流水线出的问题调度 / 代码生成 / 渲染 / 运行时VIZDEBUG基本能把范围缩到某一个工位。六、接下来可以往哪走性能tinygrad 支持 BEAM 搜索自动在多种 kernel 配置里挑更快的理解代码生成那段优化链范围拆分、符号化、融合、内存合并之后你会知道从哪下手调。多设备调度层已经为跨卡 allreduce 做了准备往 schedule 方向读能看懂多卡怎么切。写测试改核心逻辑前后跑python3 -m pytest test/仓库还有 process replay 机制用来对比你的改动有没有悄悄改变生成的 kernel。接真实模型仓库examples/里有 mnist、llama、stable_diffusion、yolov8 等拿一个跑通再回头读流水线会比纯读代码好受很多。回到主线tinygrad 把可扩展落实成了一条五段式流水线段与段之间用一个公共结构传话每段都可以按文件约定和列表装配单独替换。你要做的加中间件落到这里就是——要么往重写链里塞一条 PatternMatcher 规则要么往设备里加一个渲染器或一个运行时。读懂契约剩下的就是填空。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考