数据截至 2026-08-27(上游 commit 118e812b316f)

tinygrad — 架构与原理

30 秒导读: tinygrad 是一个小到能读完的深度学习框架(本 commit 全包约 2.4 万行 Python)。它对外是 PyTorch 风格的 eager Tensor API,对内是一台完整的编译器:你写的每个算子只是在建一张 UOp 图,.realize() 时调度器把图切成若干 GPU kernel,再逐 kernel 生成 PTX/Metal/LLVM 等源码并编译执行。它站在 micrograd(教学玩具)与 PyTorch(工业巨兽)之间——能真的训练 LLaMA,也真的能从第一个文件读到最后一个。

1. 这是什么(零基础也能懂)

一句话定义

tinygrad 是一个带自动微分的懒求值张量库 + 面向多种加速器(NVIDIA/AMD/Metal/CPU/WebGPU…)的 kernel 编译器,外加跑真实训练所需的 nn/optim/数据集。

它要解决谁的什么问题

设想你想搞清楚「一个深度学习框架到底由什么组成」:

PyTorch 有数百万行 C++/CUDA,读不完;
micrograd 只有一百多行,但只支持标量,离真实框架太远;
TVM/XLA 只是编译器,不管前端;JAX 的 IR 又藏得很深。

tinygrad 的回答是:用行数预算把「完整」压缩到「可读」。sz.py 是它自带的行数统计脚本,仓库文化就是给框架本体设行数上限——所以它被迫把每个机制都做到本质。

它能做什么

能力	具体支持
Tensor + autograd	eager API、广播、符号化 shape(`Variable`)、符号化反向传播(`tinygrad/tensor.py`、`tinygrad/mixin/`)
编译器	UOp IR、图重写、kernel 融合、BEAM 自动调优、TensorCore(`tinygrad/uop/`、`tinygrad/schedule/`、`tinygrad/codegen/`)
后端	NV(自研用户态驱动)、AMD、CUDA、Metal、CPU(clang/LLVM)、OpenCL、WebGPU、Qualcomm DSP(`tinygrad/runtime/`)
训练栈	`nn.Linear`/`Conv2d`/`BatchNorm`、SGD/Adam、fused optimizer、MNIST 数据集(`tinygrad/nn/`)
JIT/图执行	`@TinyJit` 捕获-重放、设备图(类 CUDA Graph)(`tinygrad/engine/jit.py`)
多设备	`Tensor.shard` 张量并行、allreduce(`tinygrad/schedule/multi.py`、`tinygrad/schedule/allreduce.py`)
模型	examples 里有 GPT-2、LLaMA、Stable Diffusion、Mixtral、olmoe 等(`examples/`)

用起来什么样

README 里的官方最小示例——一个手写 matmul,靠「懒」融成一个 kernel:

from tinygrad import Tensor

N = 1024
a, b = Tensor.empty(N, N), Tensor.empty(N, N)
# reshape + broadcast 乘 + sum,三次调用都不计算,只是在建图
c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2)
c.realize()   # 到这一行才调度、编译、执行;DEBUG=4 能看到生成的 kernel 源码

训练体验则是 PyTorch 味(摘自 examples/beautiful_mnist.py:18-29,它把 __call__ 标成 @function、把训练步标成 @TinyJit):

class Model:
  @function                                  # 整个前向变成一个可参数化的 CALL
  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)

  @TinyJit                                   # 第二次调用起直接重放 kernel 序列
  @Context(TRAINING=1)
  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
    opt.zero_grad()
    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
    return loss.realize(*opt.schedule_step())

一句话直觉

把 tinygrad 想成「张量界的 Lisp」: 一切(张量、算子、梯度、kernel、甚至模式规则本身)都是同一种 S-表达式式的 UOp 树;框架做的每一件事——建图、求导、融合、优化、生成代码——都是对这棵树做模式匹配重写。

2. 顶层全景(它大概怎么转)

2.1 主线结构

怎么读这张图:自上而下是一次 .realize() 的数据流;每个方框是框架的一层,标注了对应的源码目录。

 Tensor API(mixin 算子,PyTorch 味)
      │  每个算子只建图,不算
      ▼
 UOp 懒计算图 ──── realize() ────► callify:整张图包成一个 CALL
      │                             (输入 buffer 换成 PARAM)
      ▼ schedule/(rangeify:融合 + 切 kernel + 依赖排序)
 LINEAR:一串 kernel CALL
      │  一个 CALL = 一个 kernel = SINK(AST) + buffer 列表
      ▼ codegen/(逐 kernel:优化 → lowering → 渲染 → 编译)
 PROGRAM(源码 SOURCE + 二进制 BINARY)
      │
      ▼ engine/realize.py run_linear:逐个 CALL 派发
 设备执行层(runtime/ops_*.py:GPU/CPU/…)

2.2 部件职责

部件	干什么	在哪个文件
`Tensor`	用户 API,只是 `uop/grad/is_param` 三个字段的薄壳	`tinygrad/tensor.py:270`
mixin 算子库	elementwise/movement/reduce/创建等全部算子,实现为建 UOp	`tinygrad/mixin/op.py`、`tinygrad/mixin/elementwise.py:12`
`UOp` / `Ops`	统一 IR:一个 op 码 + dtype + src 子节点 + arg	`tinygrad/uop/ops.py:238`、`tinygrad/uop/__init__.py:13`
`PatternMatcher` / `UPat`	声明式重写规则与模式语言,全框架的驱动方式	`tinygrad/uop/ops.py:1474`、`:1326`
`compute_gradient`	符号化 autograd:对 UOp 图逐节点套梯度规则	`tinygrad/mixin/gradient.py:100`
`transform_to_call`	「callify」:把张量图包成一个参数化 CALL	`tinygrad/tensor.py:220`
`run_rangeify` / `get_kernel_graph`	融合决策:movement op 变索引算术,决定哪些中间结果落地成 buffer,切成 kernel	`tinygrad/schedule/indexing.py:188`、`tinygrad/schedule/rangeify.py:378`
`create_schedule`	按 RAW/WAR 依赖把 kernel 拓扑排序成 LINEAR	`tinygrad/schedule/__init__.py:28`
`full_rewrite_to_sink`	单 kernel 的约二十步 lowering(展开 reduce、加 local 内存、加 GPU 维度、反量化分解…)	`tinygrad/codegen/__init__.py:286`
`Scheduler` / `beam_search`	kernel 级优化:轴类型变换(UPCAST/LOCAL/UNROLL/TC)与自动搜索	`tinygrad/codegen/opt/postrange.py:14`、`tinygrad/codegen/opt/search.py:111`
`Renderer` 家族	把线性化后的 UOp 渲染成 C 系源码 / PTX / LLVM IR / WGSL / NIR	`tinygrad/renderer/cstyle.py:120`、`tinygrad/renderer/ptx.py:137`
`Device` / `Buffer` / `Allocator` / `Compiled`	设备发现、显存分配与缓存、编译与加载	`tinygrad/device.py:15`、`:102`、`:250`、`:343`
`run_linear` / `pm_exec`	执行引擎:编译 LINEAR 并逐 CALL 派发	`tinygrad/engine/realize.py:320`、`:299`
`TinyJit` / `GraphRunner`	捕获 kernel 序列,之后按 key 重放;可拼成设备图	`tinygrad/engine/jit.py:214`、`:89`
`Optimizer`	参数分组、fused optimizer(所有参数拼成一根向量一次更新)	`tinygrad/nn/optim.py:7`

2.3 主线走一遍:`(a @ b).sum().realize()` 发生了什么

建图(不算)。 a @ b 走 Tensor.matmul → dot(tinygrad/mixin/op.py:391、:367):reshape/转置/乘/sum 全经 Tensor._apply_uop(tinygrad/tensor.py:330)拼成一张 UOp 图。此刻没有任何计算。
callify。 .realize()(tinygrad/tensor.py:416)先调 transform_to_call(tinygrad/tensor.py:220):给图里要落地的值分配输出 buffer,把输入 buffer 换成 PARAM,整张图变成一个 CALL。
调度。 create_linear_with_vars(tinygrad/schedule/__init__.py:181)对每个图跑 prepare_rangeify → get_kernel_graph(tinygrad/schedule/rangeify.py:378)→ create_schedule(tinygrad/schedule/__init__.py:28):reshape/广播等 movement op 被消成索引算术(RANGE),matmul 与 sum 融合,切成一个(或几个)kernel,按依赖排序成 LINEAR。
编译。 run_linear(tinygrad/engine/realize.py:320)先 lower_and_compile(tinygrad/engine/realize.py:268):每个 kernel 的 SINK 经 to_program(tinygrad/codegen/__init__.py:506)走完 lowering、轴优化(可能 BEAM 搜索)、渲染,得到带源码和二进制的 PROGRAM;多个 kernel 用进程池并行编译。
执行。 pm_exec(tinygrad/engine/realize.py:299)按 CALL 的 ast 类型派发:PROGRAM 走 exec_kernel(绑定 buffer、发射 global/local 尺寸),跨设备 COPY 走 exec_copy。
回填。 输出 buffer 的 UOp 写回所有存活 Tensor(_apply_map_to_tensors,tinygrad/tensor.py:244),之后 c.numpy() 就能直接读数。

这条线最值得记住的一点: 除了第 5 步,一切都是纯 Python 的图变换;设备只在最后被碰到。这就是 tinygrad「前端、调度器、lowering、执行」四段论(它自己的 docs/developer/developer.md 也这么分)。

3. 阅读地图(建议顺序)

五章由浅入深。时间有限读 01 → 03 → 04 就能抓住 tinygrad 的全部要害:它怎么表示计算、怎么把计算切成 kernel、怎么把 kernel 变成机器码。

顺序	章节	讲什么	适合谁
1	01-tensor-and-autograd.md	Tensor 薄壳、算子建图、符号化 autograd、@function	所有人必读,框架的「形状」
2	02-uop-and-rewrites.md	UOp/Ops、interning、UPat 模式语言、graph_rewrite 引擎	想读懂任何一处源码的人(重写无处不在)
3	03-schedule-kernels.md	callify、rangeify 融合、切 kernel、依赖排序、内存规划	关心「融合是怎么决定的」的人
4	04-codegen-renderer.md	lowering 流水线、轴优化/BEAM/TensorCore、渲染器	关心性能与 GPU 代码生成的人
5	05-device-runtime-jit.md	Device/Buffer/Allocator、后端、TinyJit、多设备	要接新硬件或部署的人

4. 巧妙之处(可借鉴的技术)

自我设限的行数预算。 仓库自带行数统计 sz.py,本 commit 的 tinygrad/ 包合计约 2.4 万行。预算逼着作者把每个概念做到最小:autograd 没有 Function 子类体系,渲染器没有 visitor 类体系,全没了。
一切皆图重写。 建图、求导、融合、lowering、渲染,全用同一个 PatternMatcher + graph_rewrite(tinygrad/uop/ops.py:1474、:1753)。学会一套机制,读懂整个框架。
UOp 全局 interning。 UOpMetaClass.ucache(tinygrad/uop/ops.py:199-200)保证「同构子图即同一对象」:相等比较是指针比较,缓存键是 O(1) 的;UOp.key(tinygrad/uop/ops.py:269-270)再给出内容哈希,调度缓存(tinygrad/schedule/__init__.py:118)和编译缓存(tinygrad/codegen/__init__.py:505)都建在它的上面。
符号化 autograd。 梯度不是预先写死的 backward 函数,而是对 UOp 图现场套 pm_gradient 规则生成的另一张 UOp 图(tinygrad/mixin/gradient.py:55)——反向图自动享受和前向完全相同的融合、优化与编译。这是 JAX 式「IR 上的 AD」。
movement op 永不搬数据。 reshape/permute/expand/pad/shrink/flip 只是给「坐标」做变换;rangeify 时它们被消成对 RANGE 变量的算术(apply_movement_op,tinygrad/schedule/indexing.py:169),连 reshape 合并都成了符号化简问题。
优化 = 改轴的类型标签。 GPU 优化不改写代码,而是把某个 RANGE 的 AxisType 从 WEAK 改成 GLOBAL/LOCAL/UPCAST/UNROLL/GROUP_REDUCE(Scheduler.shift_to,tinygrad/codegen/opt/postrange.py:90)。BEAM 搜索的空间就是「轴类型赋值」的空间。
权重文件即 Tensor。 Tensor(pathlib.Path) 直接建一个 DISK tensor(tinygrad/tensor.py:286),加载模型就是 mmap + 按需拷贝,不需要单独的「反序列化」层。
可视化内建。 VIZ=1 会把每一步 graph_rewrite 的前后对比存下来,用 tinygrad/viz/serve.py 起服务逐步回放——调试编译器的一等工具。

5. 边界与局限

诚实清单,大部分能从代码或 README 直接读出:

API 不稳定。 框架处于快速重写期(本 commit 刚完成向 rangeify 调度器的迁移),内部接口几乎每个版本都变;源卡片也把它列为 gotcha。
纯 Python 前端开销。 调度与重写都在 Python 里跑,小 kernel 密集时 CPU 端会成为瓶颈;缓解手段是调度缓存 SCACHE(tinygrad/helpers.py:282)、并行编译进程池(tinygrad/engine/worker.py)、@TinyJit 重放、@function(precompile=True)。
功能取舍。 README 自述「fewer functional transforms (no full vmap/pmap yet)」;没有类似 autograd 的高阶梯度生态。
正确性靠 SPEC 而非类型系统。 UOp 图的合法性由 tinygrad/uop/spec.py 在 SPEC 开关下运行时校验(tinygrad/uop/ops.py:206-214),不是静态保证。
后端成熟度参差。 NV/AMD 是重点投入的自研用户态驱动;runtime/ 里其余后端(OpenCL、DSP、WEBGPU 等)能力不一,ops_null.py/ops_python.py 只是参照实现。
动态 shape 是「够用」级别。 用 Variable 符号化支持(如变长序列),JIT 重放时按 var_vals 重算发射尺寸(tinygrad/engine/jit.py:137-142),但不是 JAX 那样的一等抽象。
行数预算的代价。 代码有时「dense and clever rather than obvious」(源卡片原话)——单行里塞多个语义,初读需要适应。

6. 横向对比

维度	micrograd	tinygrad(本篇)	pytorch(inferred:并行写作中)
计算单元	标量 `Value`	任意形状 `Tensor`(UOp 图)	`Tensor`(ATen)
求值	立即(eager)	懒,`.realize()` 触发编译	eager + 可选 torch.compile
autograd	每节点存 `_backward` 闭包	符号化:对 IR 套规则生成反向图	每算子 C++ `BackwardFunction`
执行	Python 循环	融合成 kernel,编译到 GPU/CPU	预编译算子库 + 图编译器(Inductor)
规模	~150 行	~2.4 万行	数百万行

一句话:micrograd 教你「autograd 是什么」,tinygrad 教你「真实框架比它多出来的每一样东西——shape、广播、设备、融合、codegen——最少需要多少代码」。

7. 代码地图(入口级)

每章末尾有细粒度地图,这里只列从零开始读源码的六个入口:

主题	文件路径	符号名
Tensor 与建图	`tinygrad/tensor.py`	`Tensor`、`Tensor._apply_uop`、`Tensor.realize`、`transform_to_call`
autograd	`tinygrad/mixin/gradient.py`	`pm_gradient`、`compute_gradient`
IR 与重写引擎	`tinygrad/uop/ops.py`	`UOp`、`Ops`、`UPat`、`PatternMatcher`、`graph_rewrite`
调度器	`tinygrad/schedule/__init__.py`、`tinygrad/schedule/rangeify.py`	`create_linear_with_vars`、`get_kernel_graph`、`create_schedule`
编译器	`tinygrad/codegen/__init__.py`	`full_rewrite_to_sink`、`to_program`、`do_linearize`
设备与执行	`tinygrad/device.py`、`tinygrad/engine/realize.py`	`Device`、`Buffer`、`Compiled`、`run_linear`、`pm_exec`

1. 这是什么(零基础也能懂)​

一句话定义​

它要解决谁的什么问题​

它能做什么​

用起来什么样​

一句话直觉​

2. 顶层全景(它大概怎么转)​

2.1 主线结构​

2.2 部件职责​

2.3 主线走一遍:(a @ b).sum().realize() 发生了什么​

3. 阅读地图(建议顺序)​

4. 巧妙之处(可借鉴的技术)​

5. 边界与局限​

6. 横向对比​

7. 代码地图(入口级)​