PyTorch Compile 3：后端

图有了，为什么还不能直接编译#

Dynamo 交出来的 FX graph，从后端的角度看还很「粗糙」，主要有三点不合用：

• 算子太多、太高层。 PyTorch 暴露给用户的算子数量庞大（论文里统计约有三千个，含各种重载）。让每个后端都为这几千个算子分别写一遍高效实现，既不现实也没必要。
• 带着副作用。 图里可能还有原地改写张量的算子（如 x.add_(1)、copy_），以及视图和别名带来的隐式数据共享。这些副作用让编译器很难安全地重排、融合算子。
• 只有前向。 训练要靠反向传播更新参数，但 Dynamo 只追踪了前向计算，图里没有反向。

后端要的是一张更规整的图：算子种类少、纯函数（无副作用）、训练时还得带上反向。把粗糙的前向图整理成这样一张图，正是 AOTAutograd 和算子降解这两步在做的事——它们夹在 Dynamo 和 Inductor 之间，对应第一篇编译栈里 Dynamo 之后、Inductor 之前的那两层。

AOTAutograd：补出反向图，顺便把图捋顺#

AOTAutograd 是 PyTorch 里一个可复用的组件，被许多编译后端共用，主要干两件事：补出反向图，以及把图整理成后端好处理的形式。

先说反向图。eager 模式下，PyTorch 的反向图是运行时动态生成的：前向每算一步就在一条「磁带」（tape）上记一笔，反向时再回放这条磁带求梯度。但编译器需要提前（ahead-of-time，AOT 的由来）就拿到完整的反向图，才能对它做优化。AOTAutograd 的做法是：用假张量（fake tensor，只有 shape、dtype 等元信息、没有真实数据的张量）作输入，跑一遍 eager 的自动微分引擎，把前向和反向一起记录成一张联合图（joint graph）。这里用假张量，是因为此时只想拿到计算的结构、不需要真算数值——这也和上一篇符号执行只追踪结构、不碰真实值一脉相承。（顺带一提：有些依赖数据具体值的算子在假张量上跑不了，Dynamo 会在这些地方 graph break，把它们放到图外用真实张量执行。）

联合图记录了前向和反向的全部计算，但真正跑的时候两者是分开的（前向先跑，反向按需再跑）。AOTAutograd 于是用一个最小割算法（min-cut）把联合图切成独立的前向图和反向图，切分目标是优化显存占用。切的过程中还会顺带做一类优化叫重算（rematerialization）：反向要用到的一些中间激活值，如果重新算一遍比一直存着更划算，就不在前向里保存它、而是到反向里按需重算，用一点额外计算换显存。

第二件事是函数化（functionalization）：把图里所有原地改写、带副作用的算子，换成语义等价的纯函数版本（比如 add_ 换成 add），让后端最终面对的是一张没有副作用的纯函数图，重排和融合才有安全保证。

算子降解：把几千个算子压到一小撮#

上面提到 PyTorch 的算子多达几千个。算子降解（decomposition）就是用来收拾这个问题的：把复杂的高层算子，用一组更基础的算子重新表达出来，后端于是只需要支持这一小组基础算子。

降解规则写起来就是一个普通的 Python 函数，用别的算子把某个算子实现一遍。比如 log2 可以降解成 log：

1
log2_scale = 1 / math.log(2)
2

3
@register_decomposition(torch.ops.aten.log2)
4
def log2(x):
5
    return torch.log(x) * log2_scale

这样后端就不必单独实现 log2，只要会 log 和 mul 即可。降解会被递归地展开、规范化，直到全部落在基础算子集合上。论文里提到写作时共有 191 条降解规则（算上重载是 387 条），大部分并非 Inductor 专用，任何后端都能通过 torch._decomp 复用。

这套「把大算子集压缩成小算子集」的努力，就是编译栈里 PrimTorch 那一层：它定义一组规范的、数量小得多的基础算子，外加一个降解库把庞大的算子表面改写到这组基础算子上。降解这一步实际是搭在 AOTAutograd 上做的（AOTAutograd 调用时会带上一份想要的降解规则字典），所以第一篇编译栈图里 AOTAutograd 和 PrimTorch 是紧挨着的两层。经过降解，Inductor 要操心的算子种类大大减少；万一碰到一个它确实不认识、也没有降解规则的算子，就退回去调用原始的 PyTorch 实现（fallback），保证不会卡死。

Inductor：把图翻译成 define-by-run 的循环 IR#

到 Inductor 这一步，拿到的已经是一张规整的图：纯函数、算子种类少、训练时带反向。Inductor 是 PyTorch 官方的参考后端，也是 torch.compile 的默认后端，它先把这张 FX 图降低（lower）成自己的一套中间表示（IR），再在这套 IR 上做优化和代码生成。

Inductor 的 IR 有个特点，叫 define-by-run 的循环级 IR（loop-level IR）：循环体不用一个固定的数据结构来描述，而是直接用一段可执行的 Python 代码来定义「怎么算出一个元素」。举例来说，log2 那个算子降解后，在 Inductor 里大致长这样：

1
def inner_fn(index):
2
    i0, i1 = index
3
    tmp0 = ops.load("arg0", i0 * s1 + i1)   # 从输入张量取一个元素
4
    tmp1 = ops.log(tmp0)
5
    tmp2 = ops.constant(1.4426950408889634) # 1/log(2)
6
    tmp3 = ops.mul(tmp1, tmp2)
7
    return tmp3                              # 算出输出的这个元素

这个 inner_fn 描述的是「给定坐标 index，输出张量的这个位置该怎么算」：从输入 load 一个元素，取对数、乘以常数，得到结果。张量的尺寸和步长用 SymInt 那样的符号（s0、s1）表示——这正是上一篇的动态 shape 能一路贯通到后端的原因。整套循环级 IR 在写作时只有 54 个基础操作：ops.load/ops.store 读写内存、ops.reduction 做归约、ops.masked 处理条件、随机数相关的几个，以及其余的逐元素数学运算。Inductor 为大约 433 个 PyTorch 算子写了降低规则（算上重载是 1605 个）；遇到没有降低规则的算子，同样退回原始 PyTorch 实现。

用可执行 Python 定义循环体（而不是一套死的数据结构）的好处是：IR 直接拥有 Python 的全部表达力，能自由组合、嵌入逻辑，省掉大量模板代码。更巧的是 ops.* 是一个可替换的命名空间——把它指向「分析」用的实现，跑一遍 inner_fn 就能记录下这个算子读写了哪些内存；把它指向「代码生成」用的实现，跑一遍就能吐出对应的 Triton 或 C++ 代码。同一段 IR，换一套 ops.* 的含义，既能拿来分析、又能拿来生成代码。

算子融合与内存规划#

降低成 IR 之后，Inductor 进入调度（scheduling）阶段，决定三件事：哪些算子融合成一个 kernel、kernel 之间按什么顺序跑、以及怎么规划显存（哪些缓冲区可以复用或提前释放）。

调度器先把 IR 里的每个缓冲区变成一个调度节点，再根据每个节点读写了哪些内存，在节点之间连出依赖边。这些依赖边标注了符号化的内存地址，融合是否合法就靠它判断：比如一个节点顺序地写某个缓冲区、而下游节点却要逆序去读它，两者就不能融合。在合法的前提下，调度器用一个打分函数给所有可行的融合排序，反复贪心地融合，直到没有新的融合机会为止。打分优先考虑：融合属于哪一类（逐元素 / 归约 / 模板）、能省下多少内存往返、两个节点在原图里离得近不近。

融合为什么这么重要？因为深度学习里大量算子是访存受限（memory-bound）而非计算受限的——真正的瓶颈是把数据在显存和计算单元之间搬来搬去，而不是算术本身。如果每个逐元素算子都单独做成一个 kernel，那么每个 kernel 都要把输入从显存读进来、把中间结果写回显存，中间张量白白在显存里走一个来回。把一串逐元素算子融进同一个 kernel，中间结果就留在寄存器里，一个元素从头算到尾，只在最前和最后各访问一次显存。对访存受限的负载，省下的这些显存往返就是实打实的加速——这也是为什么打分函数把「省下的内存往返字节数」作为关键指标。

codegen：GPU 出 Triton，CPU 出 C++#

融合和调度定下来之后，每个 kernel 的循环体要落成真正的代码。Inductor 有两条代码生成路径：

• GPU → Triton。 Triton 是 OpenAI 的一门 GPU kernel 语言兼编译器，用相对简洁的 Python 式代码就能写出接近手写水平的 GPU kernel。Inductor 把循环级 IR 映射成 Triton：每个 kernel 一次处理一个 XBLOCK 大小的数据块，块尾对齐不上的部分用掩码盖掉，中间还做公共子表达式消除（把重复计算缓存进 tmp 变量）。归约有两套写法——小归约整块载入、留在寄存器或共享内存里算（persistent reduction），大归约则用块做累加器循环着算。矩阵乘、卷积这类复杂算子则走模板：手写的 Triton 骨架加上生成的部分拼起来。
• CPU → C++/OpenMP。 CPU 路径生成带 OpenMP 的 C++。有向量化和非向量化两个变体：向量化变体做分块，用 at::vec::Vectorized 一次处理 16 个元素、映射到 SIMD 指令；非向量化变体则是较标准的 C++。两者都用 #pragma omp for 做多线程并行。

选择直接生成 Triton 和 C++/OpenMP、而不是自己从零造一套代码生成器，是 Inductor 的一个有意设计（论文称之为「复用业界最好的语言」）：既能借力这两个生态里已有的优化，产出的代码也相对可读、好调试。

reduce-overhead 背后：用 CUDA graph 消除启动开销#

kernel 生成好了，还需要有人按顺序调用它们、算好每个张量的尺寸、分配和释放显存——这是包装层代码生成（wrapper codegen）干的活。包装层有两个版本：Python 版更灵活、能覆盖更多边角情况，C++ 版开销更低。

对小模型或小 batch 来说，GPU kernel 本身跑得很快，反而是 CPU 这边发起每个 kernel 的开销（算子分发、包装层代码）成了瓶颈——GPU 干等着 CPU 一个个把 kernel 发下来。第一篇提到的 reduce-overhead 模式针对的就是这种情况：它用 CUDA graph 把整段 kernel 启动序列录制下来，之后在 CUDA 驱动层把这段序列作为一个整体重放，省掉逐个启动的 CPU 开销，比 C++ 包装层还要低。

代价是 CUDA graph 有使用前提，只在满足条件时才启用，某些情况下（动态 shape、非 CUDA 张量等）会自动关掉。所以 reduce-overhead 恰好在「启动开销占大头」的场景——小的、shape 固定的模型——收益最明显。

max-autotune 在调什么#

default 模式下，Inductor 用一套启发式规则来决定 kernel 的配置（比如块大小）。第一篇提到的 max-autotune 模式则改用实测：对一个算子，在实际的输入 shape 上生成多个候选实现——Triton kernel 的不同块大小、分块方式，矩阵乘和卷积则有多个模板可选（有时还包括直接调用 cuBLAS 这类库实现）——把它们逐个跑一遍、计时，选最快的那个。这就是编译期的自动调优（autotuning）。

代价很直白：编译要慢很多（编译过程里真的在跑基准测试），换来运行更快。这正好呼应第一篇对 mode 的定位——max-autotune 用更长的编译时间换取更快的运行速度，适合编译一次、反复运行的场景。

小结#

后端这条链路可以串成一句：Dynamo 交出的 FX graph，先由 AOTAutograd 用假张量跑一遍自动微分、录出前向加反向的联合图，再用最小割切成前向图和反向图，并做函数化去掉副作用；接着算子降解把几千个高层算子压到一小撮基础算子上；然后 Inductor 把图降低成 define-by-run 的循环级 IR，在调度阶段做算子融合与内存规划，最后 codegen 成 GPU 上的 Triton 或 CPU 上的 C++/OpenMP，由包装层负责调用（必要时叠上 CUDA graph）。第一篇讲的两个 mode 到这里也有了出处：reduce-overhead 是用 CUDA graph 消除启动开销，max-autotune 是编译期实测选最快的 kernel。

论文在 A100 上做过测评：仅换后端、不动前端抓图的情况下，Inductor 相对 eager 的几何平均加速在 TorchBench 上约为推理 2.7×、训练 1.4×，在大多数配置下都快过其他 TorchDynamo 后端。

到这里，torch.compile 作为一个 JIT 编译器的原理就讲完了：它在运行时抓图、编译、缓存复用，整个过程都留在 Python 进程里。但要把模型部署到生产环境——脱离 Python、提前编好、用 C++ runtime 加载——需要的是另一套工具：torch.export 和 AOTInductor。