跳到主要内容

数据截至 (上游 commit 118e812b316f)

03 · 调度器:从张量图到 kernel 序列

这一章讲什么: tinygrad 的「大脑」。一张任意大的 UOp 图,怎么决定哪些中间结果要落地成 buffer(切 kernel 的边界),怎么把剩下的全部融进 kernel,又怎么排出执行顺序。全程还是图重写,没有一行命令式调度代码。

1. 它要解决的小问题​

(a @ b).sum() + c 这张图理论上可以融成一个 kernel,但 (a @ b) 的结果如果后面还要单独用一次,就必须先落地。再放大:一个 transformer 的一次前向是上千个节点的一张图。

调度器要回答三个问题:

  1. 哪里切:哪些中间值必须写成 buffer(= 一个新 kernel 的输出),哪些可以融进下游;
  2. 怎么排:切出来的 kernel 之间有读写依赖,执行顺序是什么;
  3. 怎么省内存:中间 buffer 的生命期已知,能不能复用显存。

2. 全流程一张图​

 张量图(若干 Tensor 的 SINK)
      │ ① callify:分配输出 buffer,输入换 PARAM
      ▼
 一个 CALL(PARAM0, PARAM1, …)  ← 一张「可缓存的函数图」
      │ ② rangeify:movement op 消成 RANGE 算术;按消费者关系做融合决策
      ▼
 rangeified 图(STAGE = 候选 buffer,REDUCE 带 RANGE)
      │ ③ 化简 + STAGE→STORE
      ▼
 若干 STORE/END 子树
      │ ④ split_store:每棵 → 一个 kernel CALL(SINK + KernelInfo)
      │ ⑤ create_schedule:AFTER 链上建 RAW/WAR 依赖,拓扑排序
      ▼
 LINEAR(排好序的 kernel CALL 列表)+ ⑥ 内存规划(arena 子分配)

对应入口 create_linear_with_vars(tinygrad/schedule/__init__.py:181-210),它把 ②~⑥ 串成三个 graph_rewrite 加一个内存规划。下面逐段拆。

3. 第①步:callify——把「图」变成「一次函数调用」​

transform_to_call(tinygrad/tensor.py:219-239)在 Tensor.realize 之前跑,做三件事:

  1. 打编号:给需要落地的节点打 tag,编出一张 uop_list(add_tags,tinygrad/tensor.py:44-51);
  2. 分配输出 buffer:每个被 tag 的节点换成「新 buffer + STORE + AFTER」(replace_contig_with_store_after,tinygrad/tensor.py:53-61);
  3. 输入参数化:把输入 BUFFER/bound Variable 换成带槽位的 PARAM(pm_replace_buf,tinygrad/tensor.py:208-217)。

产物是一个 CALL(body, PARAM0, PARAM1, ...) 加一张 buffer_map(原节点 → 新 buffer)。之后同一张图再来一次,内容哈希相同,直接命中调度缓存(schedule_cache,tinygrad/schedule/__init__.py:118)——参数化让缓存键与具体 buffer 无关。

4. 第②步:rangeify——融合决策的核心​

这是整套调度器的心脏,思想一句话:不给张量选「计算位置」,而是给每个维度生成一个循环变量(RANGE),让 movement op 全部消成对这些变量的算术。

4.1 movement op 的「消失」​

apply_movement_op(tinygrad/schedule/indexing.py:169-185)定义了六种 movement op 对 RANGE 元组做什么:

op对 RANGE 的变换
SHRINK加偏移:a → a+off
PERMUTE重排元组
FLIP反向:a → (s-1)-a
EXPAND丢弃前几个(广播维不需要变量)
PAD变成带 valid 条件的表达式
RESHAPE展平成一维再按新形状取模/整除拆回(_apply_reshape,tinygrad/schedule/indexing.py:152-165)

重点看最后一行:reshape 合并这个传统编译器里的大麻烦,在这里被变成「一维化 + 符号化简」。这就是 tinygrad 不再需要旧式 "ShapeTracker/View 栈" 的原因。

4.2 哪里必须落地:realize map​

先跑 pm_generate_realize_map(tinygrad/schedule/indexing.py:45-54)标记「非落地不可」的点:

  • CONTIGUOUS 和用户 STORE(assign):语义要求;
  • MSTACK/MSELECT(多设备)的输入;
  • STORE 的源里若读到目标自身(WAR 冒险),源要先落地(realize_store_after_src,tinygrad/schedule/indexing.py:34-36)。

4.3 融合决策:按消费者关系配 RANGE​

run_rangeify(tinygrad/schedule/indexing.py:188-280)先建 consumer map,再逆序遍历图,给每个节点配输出 RANGE。注释里把规则写得明明白白(tinygrad/schedule/indexing.py:223-227):

  1. 节点在 realize map 里 → 造新 RANGE(它是 kernel 输出,维度独立);
  2. 只有一个消费者 → 继承消费者的 RANGE(融进下游,免费);
  3. 多个消费者且大家的 RANGE 一致 → 合并 valid,继续融;
  4. 多个消费者且 RANGE 对不上 → 对不上的轴部分 realize(_realize_axis,tinygrad/schedule/indexing.py:252-266)。

另外 REDUCE 还有一条特殊规则:它的「被消掉的轴」直接变成 REDUCE 自己的 RANGE(convert_reduce_to_reduce_with_ranges,tinygrad/schedule/indexing.py:113-119),PAD/STACK 则分别变成 WHERE 条件(tinygrad/schedule/indexing.py:107-129)。

走完这一步,图里再也没有 movement op——只剩算术、LOAD/STORE 风格的 INDEX、REDUCE 和 RANGE。「融合」不是被显式执行的,而是「没被选去落地」的自然结果。

4.4 原理演示​

# 示意,非源码 —— (x.reshape(6)*y).sum() 经 rangeify 后的样子
# 输入:x:(2,3), y:(6,)
# reshape 不产生循环,只是换坐标;只有一个消费者,全程同一组 RANGE
for r0 in range(6):                       # 唯一的新 RANGE,长在 sum 的输出上
    acc += x_flat[r0] * y[r0]             # x.reshape(6) 被消成「按 r0 一维读 x」

真实图里 x.reshape(6) 会变成 x[(r0//3, r0%3)] 这样的索引算术,再被 symbolic 化简。

5. 第③④步:落地成 kernel​

5.1 STAGE→STORE,再切成 kernel​

rangeified 图里,被选去落地的节点是 STAGE(候选 buffer,带 BufferizeOpts,tinygrad/schedule/indexing.py:56-61)。get_kernel_graph(tinygrad/schedule/rangeify.py:378-403)接着:

  1. 跑 symbolic + pm_reduce_simplify + pm_const_buffer_folding + pm_remove_bufferize 化简——其中 remove_bufferize(tinygrad/schedule/rangeify.py:51-103)会把「只为 permute 之类而落地」的 STAGE 再融回去,代价模型很朴素:访问的 buffer 超过 3 个、或 REDUCE 里读了 buffer,就保留(tinygrad/schedule/rangeify.py:86-97);
  2. pm_limit_bufs(tinygrad/schedule/rangeify.py:199):单 kernel 输入 buffer 数超设备上限时,强制多切几个;
  3. pm_add_buffers(tinygrad/schedule/rangeify.py:262):STAGE 变成真正的「新 BUFFER + STORE」;
  4. split_kernels:每棵 STORE/END 子树经 split_store(tinygrad/schedule/rangeify.py:359-371)变成一个 kernel——SINK 做 AST、KernelInfo 挂元数据、输入输出 buffer 列在 CALL 的 src 里。

5.2 一个 kernel 长什么样​

CALL(
  SINK( STORE(INDEX(PARAM_out, r0), ADD(LOAD(INDEX(PARAM_a, r0)), LOAD(INDEX(PARAM_b, r0)))) )
  , arg=KernelInfo(name="E_1024")
, PARAM_out, PARAM_a, PARAM_b )
  • SINK 里就是这个 kernel 要算的全部内容,RANGE r0 是它的循环结构;
  • KernelInfo(tinygrad/uop/ops.py:1222-1231)带 kernel 名、已应用/待应用的优化、BEAM 值;
  • CALL 的其余 src 是运行时要绑的 buffer——「AST + buffer 列表」就是 kernel 的全部。

6. 第⑤步:依赖排序成 LINEAR​

kernel 之间的顺序由 AFTER 节点表达(assign、跨 kernel 读写同一 buffer 都会产生 AFTER)。create_schedule(tinygrad/schedule/__init__.py:28-79)显式建依赖图:

依赖规则位置
RAW(写后读)读到某个 AFTER 状态的 kernel,排在产生该状态的 kernel 之后tinygrad/schedule/__init__.py:47-52
WAR(读后写)读旧状态的 kernel,必须排在「覆盖该状态」的写之前tinygrad/schedule/__init__.py:53-61

然后是标准 Kahn 拓扑排序(tinygrad/schedule/__init__.py:63-78),有环直接报错(那意味着 assign 图自相矛盾)。产物就是 UOp(Ops.LINEAR, src=(CALL, CALL, ...))。

6.1 附加加工​

create_linear_with_vars 收尾还做三件事:

  • resolve_linear_call(tinygrad/schedule/__init__.py:101-111):递归解析嵌套 LINEAR,把 PARAM 换成真实 BUFFER;
  • pm_copy_from_store(tinygrad/schedule/__init__.py:165-179):认出「纯搬运」的 kernel,降级成 COPY(走 SDMA/copy engine,不占 SM);
  • memory_plan_rewrite(tinygrad/schedule/memory.py:20-62):见下。

6.2 内存规划:arena 子分配​

所有中间 buffer 的生命期在 LINEAR 里已知(首次/末次出现的 kernel 序号)。memory_plan_rewrite:

  1. 按生命期做开/闭事件排序,用 TLSF 分配器(tinygrad/runtime/support/memory.py 的 TLSFAllocator)在每个 (设备, 是否 copy) lane 里摆 offset(tinygrad/schedule/memory.py:45-55);
  2. 整个 lane 只真分一块 arena,每个中间 buffer 换成对 arena 的 SHRINK+BITCAST 视图(tinygrad/schedule/memory.py:57-59)。

DEBUG>=1 时会打印省了多少(memory reduced from X MB -> Y MB,tinygrad/schedule/memory.py:61-62)。被外部引用的 buffer(held)不参与。

7. 关键细节与坑​

  • 调度是幂等且可缓存的。 callify 参数化 + UOp.key 内容哈希 + schedule_cache(tinygrad/schedule/__init__.py:118),同一图第二次调度是字典查询。SCACHE=0 可关(tinygrad/helpers.py:282)。
  • 多设备走另一条路。 shard 张量的 UNSHARD/MSTACK/MSELECT 在 multi_pm(tinygrad/schedule/multi.py:284)里先展开成逐设备图,再进同一套 rangeify;allreduce 是预先写好的函数(tinygrad/schedule/allreduce.py)。
  • CONTIGUOUS 是用户给的「融合断点」。 想要某个中间值强制落地(比如复用重的中间结果),显式 .contiguous();它会进 realize map(tinygrad/schedule/indexing.py:49)。反过来,remove_bufferize 可能会把你以为落地了的东西融回去。
  • split reduceop。 超大归约会按阈值拆成两个 kernel(提升并行度),在 split_reduceop(tinygrad/schedule/prepare.py:66-81)——prepare 阶段,不在 rangeify。
  • assign 的正确性全在 AFTER。 Tensor.assign 造 AFTER(STORE)(tinygrad/tensor.py:451),调度器靠它排 WAR;绕过这个机制自己拼 STORE 很容易造出race。
  • 看不出:BufferizeOpts.removable 之外,本 commit 里没有更精细的「融合收益模型」(如按算术强度/访存比决定融合);融合基本是「能融就融 + 少数保守退出条件」。这是读代码时的诚实结论,不是缺陷报告。

8. 代码地图​

主题文件路径符号名
callifytinygrad/tensor.pytransform_to_call、add_tags、pm_replace_buf
调度入口tinygrad/schedule/__init__.pycreate_linear_with_vars、lower_sink_to_linear、schedule_cache
rangeify 主循环tinygrad/schedule/indexing.pyrun_rangeify、IndexingContext、pm_generate_realize_map
movement op 消解tinygrad/schedule/indexing.pyapply_movement_op、_apply_reshape
化简与反落地tinygrad/schedule/rangeify.pypm_remove_bufferize、remove_bufferize、pm_limit_bufs
切 kerneltinygrad/schedule/rangeify.pyget_kernel_graph、split_store、pm_add_buffers
依赖排序tinygrad/schedule/__init__.pycreate_schedule、_split_after、_states
copy 降级tinygrad/schedule/__init__.pypm_copy_from_store、copy_kernel_to_copy_uop
内存规划tinygrad/schedule/memory.pymemory_plan_rewrite
多设备tinygrad/schedule/multi.py、tinygrad/schedule/allreduce.pymulti_pm、create_allreduce_function