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

04 · codegen 与渲染器

这一章讲什么: 第 3 章切出的每个 kernel(一棵 SINK AST)怎么变成能在设备上跑的二进制。前半是「怎么优化一个 kernel」(轴类型变换 + BEAM 搜索),后半是「怎么把它打印成代码」(渲染器家族)。

1. 它要解决的小问题

调度器给的 AST 是「数学上正确」的:一堆 RANGE 循环套着 LOAD/ALU/STORE。但直接照抄成代码会很慢——GPU 要快,需要决定:

哪些循环维变成线程网格(global/local),哪些展开(upcast/unroll),哪些塞进 TensorCore;
访存怎么合并(coalescing)、局部内存(shared/local)怎么用;
SIN/EXP2 这类目标不支持的算子怎么分解(decomp)。

这些决定既依赖数学结构,也依赖具体设备——所以 codegen 分两层:与设备无关的 lowering + 设备专属的渲染。

2. 管线总览

入口 to_program(tinygrad/codegen/__init__.py:506-508,带磁盘级缓存),主体 full_rewrite_to_sink(tinygrad/codegen/__init__.py:286-405)。约二十步,按性质分四段:

 SINK(AST,来自调度器)
   │ A. 优化段(可被 BEAM 控制)
   │   load collapse → split/simplify ranges → apply_opts
   │   (hand_coded 启发式 或 BEAM 搜索,给 RANGE 改轴型)
   ▼
 B. 展开段
   │   expander(广播/REDUCE 展开)→ reduce→acc 寄存器累加
   │   → STAGE→local buffer → add_gpudims(GLOBAL/LOCAL→SPECIAL)
   ▼
 C. lowering 段
   │   广播展开 → devectorize → memory coalescing
   │   → 弱 dtype 落地 → decomp(超越函数/dtype 分解)
   │   → 隐式 barrier → 控制流(IF/END 结构化)
   ▼
 D. 成程序段(do_to_program)
   │   linearize(优先级拓扑排序成指令列表)
   │   → render(渲染器打印源码)→ compile(外部编译器出二进制)
   ▼
 PROGRAM(SINK + LINEAR + SOURCE + BINARY)

每一步都是一次 graph_rewrite,函数体里有逐步的 name= 标注(tinygrad/codegen/__init__.py:292-405),配 VIZ=1 可以逐步回放。下面挑灵魂步骤讲。

3. 优化段:优化 = 给循环轴换类型

3.1 AxisType:轴的「职业」

每个 RANGE 的 arg 末尾带一个 AxisType(tinygrad/uop/ops.py:17-19):

轴型	含义
`WEAK`	还没分配身份的循环轴
`GLOBAL`	grid 维(线程块索引)
`LOCAL` / `WARP` / `GROUP_REDUCE`	block 内维(线程/warp),local 类需要共享内存归约
`UPCAST`	尾端向量化展开(每个线程多算几个)
`UNROLL`	归约维的展开
`REDUCE`	顺序归约循环
`THREAD`	CPU 线程
`DEVICE`	多设备维

3.2 shift_to:一切优化的原子操作

Scheduler(tinygrad/codegen/opt/postrange.py:14)把 AST 的全体 RANGE 收出来,shift_to(tinygrad/codegen/opt/postrange.py:90-97)做唯一的那件事:

# 示意,非源码 —— shift_to(rng, 4, LOCAL) 做的事
new_rng = RANGE(4, type=LOCAL)      # 新轴:长度 4,身份 LOCAL
old_rng = rng.replace(size=old/4)   # 原轴变短
# 然后全图替换: rng → old_rng * 4 + new_rng(索引表达式换元)

「把这个轴按 4 切成两段,内层那段标记成 LOCAL」——UPCAST/UNROLL/THREAD/GROUP 全是这一个原语的不同参数(apply_opt,tinygrad/codegen/opt/postrange.py:121)。Opt 只是一个 (op, axis, arg) 三元组(tinygrad/codegen/opt/__init__.py:11-16)。

3.3 谁来选 Opt:启发式、BEAM、或外部指定

apply_opts(tinygrad/codegen/opt/postrange.py:334-351)按优先级三选一:

KernelInfo.opts_to_apply 非空 → 照单执行(schedule/上层指定);
beam >= 1 → beam_search 真机搜;
否则 → hand_coded_optimizations(tinygrad/codegen/opt/heuristic.py:8)。

启发式本身值得一读,它是 tinygrad 的「性能直觉」结晶:先尝试 TensorCore(tinygrad/codegen/opt/heuristic.py:28-41),然后是图像 upcast、matvec 特化(MV_BLOCKSIZE 参数化,tinygrad/codegen/opt/heuristic.py:61-70)等一串规则。

3.4 BEAM:把「编译器调参」变成「真机计时搜索」

beam_search(tinygrad/codegen/opt/search.py:111-178)的机制:

动作空间:一张大表 actions(tinygrad/codegen/opt/search.py:14-26)——各轴 × 各倍数的 UPCAST/UNROLL/LOCAL/GROUP/TC/SWAP/THREAD,约几百种;
每轮:对 beam 里每个候选,枚举所有合法下一步(get_kernel_actions,tinygrad/codegen/opt/search.py:87-108,会检查共享内存上限等约束);
评测:_try_compile 编译(进程池并行),_time_program 在真设备上跑 3 次取最优(tinygrad/codegen/opt/search.py:38-51),剪掉算力量异常和生成库重复的;
终止:最优时间不再进步(min_progress)就停(tinygrad/codegen/opt/search.py:165-167);
缓存:结果写磁盘(diskcache_put,tinygrad/codegen/opt/search.py:176),同 AST 同设备下次秒出。

注意一个工程约束:BEAM 的编译必须在父进程做,因为计时需要真设备(tinygrad/engine/realize.py:276-278 的注释)。

3.5 TensorCore:WMMA 节点

TensorCore(tinygrad/codegen/opt/tc.py:7-15)用声明式描述一种硬件 MMA:dims(N,M,K)、每 warp 线程数、输入/输出 dtype、以及一串 "ux"/"lx" 形式的 Opt 配方和 swizzle。渲染器的 tensor_cores 列表(tinygrad/renderer/__init__.py:77)声明硬件支持哪些;OptOps.TC 命中后,AST 里会出现 WMMA 节点,渲染成对应指令(如 PTX 的 mma.sync)。hand_coded_optimizations 的注释(tinygrad/codegen/opt/heuristic.py:10-26)把 USE_TC/TC_SELECT/TC_OPT 三个旋钮的语义写得清清楚楚。

4. lowering 段:从「数学」到「机器」

挑五个最能体现设计的步骤:

reduce → 寄存器累加器。 reduce_ranges_to_acc(tinygrad/codegen/__init__.py:211-221)把 REDUCE 变成一个 REG 地址空间的 placeholder:先存 identity(identity_element,tinygrad/uop/ops.py:55),循环体里累加。GROUP_REDUCE 还要先落 local buffer 再做第二层归约(fix_group_for_reduce,tinygrad/codegen/__init__.py:171-185)。
GPU 维度实体化。 add_gpudims(tinygrad/codegen/gpudims.py:41)把 GLOBAL/LOCAL 轴换成 SPECIAL 节点(即 blockIdx/threadIdx),并按设备 global_max/local_max 拆/并维(get_grouped_dims,tinygrad/codegen/gpudims.py:27-39)。
反量化向量化展开(devectorize)。 到这里之前,UPCAST 轴是「向量 shape」;do_devectorize(tinygrad/codegen/__init__.py:123-130)把向量节点按笛卡尔积摊成标量 STACK,渲染层只处理标量(外加显式的 float4 类 STACK)。
访存合并。 memory_coalescing(tinygrad/codegen/late/coalesce.py:104)把单位步长的相邻 LOAD/STORE 合并成向量访存——GPU 带宽的生命线。
decomp:目标没有的,现场展开。 get_transcendental_patterns(tinygrad/codegen/decomp/transcendental.py:268)把 SIN/LOG2/EXP2 展成多项式;pm_dtype_decomps(tinygrad/codegen/decomp/dtype.py:215)把目标不支持 dtype 的运算拆成支持的;x**y 变成 exp2(y*log2(x)) 等。renderer 的 code_for_op 决定哪些算子「原生」,其余全被分解。

收尾还有两个容易忽略但关键的:pm_implicit_barriers(tinygrad/codegen/__init__.py:281-284)按 local 内存的写读/写后读自动插 BARRIER;pm_add_control_flow(tinygrad/codegen/__init__.py:387)把 RANGE/END 结构化成可打印的循环嵌套。

5. 成程序段:linearize → render → compile

5.1 linearize:把 DAG 摆成一行

linearize(tinygrad/codegen/late/linearizer.py:8-40)是一次带优先级的拓扑排序:LOAD 尽量早、STORE 尽量晚、RANGE 开门早、END 关门晚(tinygrad/codegen/late/linearizer.py:22-31),并用 run_count 把循环体内的计算尽量挪进/挪出循环。产物是 list[UOp]——IR 从「图」正式变成「指令序列」。

5.2 pm_to_program:状态机式的四步

do_to_program(tinygrad/codegen/__init__.py:471-496)之后,PROGRAM 节点靠 pm_to_program(tinygrad/codegen/__init__.py:461-467)逐格生长:

PROGRAM(SINK) ──do_linearize──► PROGRAM(SINK, LINEAR)
   ──do_render──► PROGRAM(SINK, LINEAR, SOURCE)     # ISA 渲染器则 do_assemble 直接出字节
   ──do_compile──► PROGRAM(SINK, LINEAR, SOURCE, BINARY)

渲染与编译都是 renderer 的方法:ctx.render(list[UOp]) -> str 与 ctx.compiler.compile_cached(src) -> bytes(tinygrad/codegen/__init__.py:451-459)。编译产物按 renderer 类型+target+配置为键缓存(to_program_key,tinygrad/codegen/__init__.py:502-503)。

6. 渲染器家族:同一 IR,多种打印

6.1 基类的「能力声明」

Renderer(tinygrad/renderer/__init__.py:64-89)的核心是一组能力字段,lowering 全程读它们做决策:

字段	作用
`has_local` / `has_shared` / `has_threads`	能不能用 local 内存 / 共享内存 / CPU 线程
`global_max` / `local_max` / `shared_max`	维度与共享内存上限
`tensor_cores`	支持的 TC 列表
`code_for_op`	「原生」算子表,不在表里的都被 decomp 展开
`extra_matcher`	渲染前的最后一批私有重写规则

6.2 CStyleLanguage:一张模式表打印 C 方言

Metal/CUDA/OpenCL/CPU(clang)/WGSL 共用 CStyleLanguage(tinygrad/renderer/cstyle.py:120),渲染就是一张 base_rewrite 模式表(tinygrad/renderer/cstyle.py:11-75),逐 UOp 打印:

(UPat(Ops.RANGE, name="x"), lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"),
(UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
(UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"{ctx.render_access(bidx)} = {ctx[var]};"),

(tinygrad/renderer/cstyle.py:17-21、:61)

各后端只覆盖差异:MetalRenderer(tinygrad/renderer/cstyle.py:349)写 kernel 前缀与 threadgroup 地址空间,CUDARenderer(tinygrad/renderer/cstyle.py:402)写 __global__/__shared__,WGSL 单有 WGSLRenderer(tinygrad/renderer/wgsl.py:55)。

6.3 PTX/LLVM/NIR:不走 C 的三支

PTXRenderer(tinygrad/renderer/ptx.py:137):每个算子直接映射到 PTX 汇编码(asm_for_op,tinygrad/renderer/ptx.py:18-37),配 ptx_matcher 做布尔转谓词寄存器、半精度升 float 等私有改写(tinygrad/renderer/ptx.py:40-55);
LLVMRenderer(tinygrad/renderer/llvmir.py:148):打印 LLVM IR 文本,交给 LLVM 编 x86/arm;
NIRRenderer(tinygrad/renderer/nir.py:117):打印 Mesa NIR,给某些 GPU 驱动栈用;
ISARenderer(tinygrad/renderer/isa/__init__.py:39):更新的一支——不生成文本,直接做指令选择 + 线性扫描寄存器分配(tinygrad/codegen/__init__.py:431-438),do_assemble 出机器码(x86 已有实现)。

6.4 Estimates:性能模型的数据从哪来

Estimates.from_uops(tinygrad/renderer/__init__.py:26-62)在线性化后的指令表上数 FLOPS、LOAD/STORE 字节、并按 RANGE 乘出总量——BEAM 的剪枝、DEBUG=2 的 GFLOPS/GB/s 显示,全靠它。

7. 关键细节与坑

IF 只剩一种合法形态。 线性化清理规定:只有 gated STORE 能变 IF/ENDIF,其它 IF 直接 panic(tinygrad/codegen/__init__.py:408-414)。控制流是刻意受限的。
弱 dtype 是「推迟决定」的 dtype。 Python 常量建图时是 weakint/weakfloat,直到 pm_lower_weak/pm_commit_weak 才定宽(tinygrad/codegen/__init__.py:348-353 一带)——这就是为什么 Tensor(1) + Tensor(1.0) 的类型提升「自动」工作。
decomp 顺序敏感。 注释明说某些化简必须在 weak 提交前、某些必须在「index 还是 weakint」时做(tinygrad/codegen/__init__.py:345-347),动 pipeline 顺序是 tinygrad 贡献者最常见的踩坑方式。
BEAM 与小 kernel 不划算。 每个候选都要真机编译+跑三遍,BEAM 适合训练场景的热点 kernel;BEAM 是 ContextVar 可按作用域开关,IGNORE_JIT_FIRST_BEAM 之类的环境变量专门处理首跑开销。
渲染器不优化。 所有性能决策在渲染前已完成,渲染层是「笨」的——这让加后端变成纯体力活,是后端数量的来源。
看不出:local_size 的最终自动调优(optimize_local_size,tinygrad/engine/realize.py:103-127)与 BEAM 的分工边界在代码里没有文档化说明,只能从调用点推断:BEAM 管 AST 层面,optimize_local_size 管已编译 PROGRAM 的发射维度微调 (inferred)。

8. 代码地图

主题	文件路径	符号名
编译入口	`tinygrad/codegen/__init__.py`	`to_program`、`do_to_program`、`full_rewrite_to_sink`
轴优化	`tinygrad/codegen/opt/postrange.py`	`Scheduler`、`shift_to`、`apply_opt`、`apply_opts`
启发式	`tinygrad/codegen/opt/heuristic.py`	`hand_coded_optimizations`
BEAM	`tinygrad/codegen/opt/search.py`	`beam_search`、`actions`、`get_kernel_actions`、`_time_program`
TensorCore	`tinygrad/codegen/opt/tc.py`	`TensorCore`
GPU 维度	`tinygrad/codegen/gpudims.py`	`add_gpudims`、`get_grouped_dims`
归约/lowering	`tinygrad/codegen/__init__.py`	`reduce_ranges_to_acc`、`fix_group_for_reduce`、`do_devectorize`、`pm_implicit_barriers`
decomp	`tinygrad/codegen/decomp/`	`get_transcendental_patterns`、`pm_dtype_decomps`
线性化	`tinygrad/codegen/late/linearizer.py`	`linearize`、`pm_add_control_flow`
访存合并	`tinygrad/codegen/late/coalesce.py`	`memory_coalescing`
渲染基类	`tinygrad/renderer/__init__.py`	`Renderer`、`Estimates`
C 系渲染	`tinygrad/renderer/cstyle.py`	`CStyleLanguage`、`base_rewrite`、`MetalRenderer`、`CUDARenderer`
非 C 渲染	`tinygrad/renderer/ptx.py`、`tinygrad/renderer/llvmir.py`、`tinygrad/renderer/nir.py`、`tinygrad/renderer/isa/__init__.py`	`PTXRenderer`、`LLVMRenderer`、`NIRRenderer`、`ISARenderer`

1. 它要解决的小问题​

2. 管线总览​

3. 优化段:优化 = 给循环轴换类型​

3.1 AxisType:轴的「职业」​

3.2 shift_to:一切优化的原子操作​

3.3 谁来选 Opt:启发式、BEAM、或外部指定​

3.4 BEAM:把「编译器调参」变成「真机计时搜索」​

3.5 TensorCore:WMMA 节点​

4. lowering 段:从「数学」到「机器」​

5. 成程序段:linearize → render → compile​

5.1 linearize:把 DAG 摆成一行​

5.2 pm_to_program:状态机式的四步​

6. 渲染器家族:同一 IR,多种打印​

6.1 基类的「能力声明」​

6.2 CStyleLanguage:一张模式表打印 C 方言​

6.3 PTX/LLVM/NIR:不走 C 的三支​

6.4 Estimates:性能模型的数据从哪来​

7. 关键细节与坑​

8. 代码地图​