「Inductor-1」从torch.compile开始认识整个流程

本文基于 PyTorch 2.14.0a0（commit 8b5ec6266b0）撰写，文中涉及的代码锚点与行为描述均以该 commit 为准。

整体流程

flowchart TD
    A[用户编写的 PyTorch 代码] --> B[TorchDynamo
图捕获前端]
    
    B --> C{是否满足
Guard 条件?}
    C -- 是 --> D[生成 FX Graph]
    C -- 否 --> E["重新编译
（形状/设备等变化）"]
    E --> B
    
    B --> F[遇到 Graph Break?]
    F -- 是 --> G[切割图
回退到 Eager 模式]
    G --> D
    F -- 否 --> D
    
    D --> H[AOTAutograd
提前自动微分与算子分解]
    H --> I[生成前向+反向联合图
并分解复杂算子]
    
    I --> J[Inductor
后端编译器]
    J --> K["转换为 Loop-based IR
（展开形状、循环）"]
    K --> L["调度器（Scheduler）
算子融合、内存规划"]
    L --> M["代码生成器（Codegen）
生成内核代码"]
    
    M --> N{目标硬件?}
    N -- GPU --> O[生成 Triton 内核
并通过 Triton 编译]
    N -- CPU --> P[生成 C++ 内核
并通过编译器编译]
    
    O --> Q[加载并执行内核]
    P --> Q
    Q --> R[输出结果]
    
    style A fill:#f9d0c4,stroke:#333
    style B fill:#b5e7a0,stroke:#333
    style H fill:#b5e7a0,stroke:#333
    style J fill:#b5e7a0,stroke:#333
    style Q fill:#f9d0c4,stroke:#333

进入torch.compile()之后的主要功能分为三部分：Dynamo,AOTAutoGrad,Inductor。三者分别对应不同的功能

1. TorchDynamo

目标是在不改变用户体验的前提下，将用户的PyTorch代码捕获生成一个可优化的FX图。相比于torch.jit.trace和torch.jit.script，Dynamo可以更加透明的支持标准Python代码，包括各种条件判断、循环、函数调用等。

1.1 核心机制

TorchDynamo不会直接解析源代码，而是在CPython的解释器层面工作，替换当前函数的frame求值函数，获得对每条Python字节码的控制，流程如下：

1.1.1 拦截字节码

当用户调用torch.compile 或者使用torch.compile 装饰器修饰函数的时候，Dynamo会将函数的__Code__对象替换为一个包装版本，时期运行时进入Dynamo的frame求值器，之后逐条分析函数的字节码指令，例如 LOAD_FAST, BINARY_OP, CALL_FUNCTION等。

1.1.2 建立符号环境

Dynamo维护一个映射：变量名->Proxy对象。Proxy是对FX图节点的封装，代表改变量在计算图中的符号值。当遇到例如x = a + b的时候不会执行实际的加法计算，而是通过以下步骤来构建出计算图：

从环境中查找a和b 对应的proxy，（或者常量）
向FX图中添加一个call_function(target=operator.add, args=(proxy_a, proxy_b)) 节点。
创建一个新的Proxy作为结果，并绑定到x变量上。
这样，这句x = a + b就加入了这个计算图，而不是一句实际的推理代码了。

1.1.3 处理Tensor Op

当遇到PyTorch函数的时候，例如torch.add,torch.nn.functional.relu或者tensor.sum()时，会执行如下逻辑：

如果当前函数的输入包含一个Proxy，那么这行调用就会被记录为一个FX节点(其实就是图上的一个算子或者函数(本质还是算子))
如果输入全是const，那么就会执行constant fold，直接将计算后的结果返回

1.1.4 处理控制流

并非所有的Python代码都可以被转换为计算图，例如各种控制流、循环等操作，这些操作通常依赖运行期间的真实值，(当然部分循环不依赖运行期的值，是可以直接展开循环的，例如Decode Layer展开)。这时Dynamo的处理逻辑为：

条件分支是静态可判断的，那么会直接计算出真实运行需要走哪个分支，然后只将改分支的代码加入到计算图中
条件分支依赖运行时tensor，Dynamo无法在编译期确定真实的推理路径，就会触发Graph Break，停止当前的图构建，将已经捕获的部分编译为一个子图，然后这个条件分支部分的逻辑使用eager mode来在运行期进行真实的分支判断，条件分支之后的代码会继续捕获，然后生成多个子图进行编译。
循环次数是静态的，那么就会循环展开，将所有的代码加入到图中，例如decode layer。
循环次数无法确定时，也就是依赖tensor的shape，那么会尝试使用Proxyshape来展开循环，或者触发Graph Break

1.1.5 支持Python的数据结构

Python中通常都存在较多的List,Tuple,Dict,等数据结构，Dynamo会在符号环境中维护这些抽象表示，相关的索引、迭代、解包等操作，也会在符号层面模拟，然后产生图节点：

x = [a, b]会常见一个包含两个Proxy的ListProxy
y = x[0]会直接返回a 的proxy，而不会产生图节点。
如果是动态的索引操作，触发GraphBreak。

1.1.6 Guard机制

Dynamo 生成的图是针对特定执行条件优化的，需要记录各种的shape等条件，如果后续调用时这些信息出现了变化，那么就无法利用此前已经编译好的产物，此时需要重新编译。会触发Guard的信息包括：

tensor的shape、stride、dtype、device.
tensor的dtype和value被用来控制流程
全局变量的值被访问
可以通过工具来查看生成的Guards，torch._dynamo.explain

1.1.7 Graph Break

以下几种情况会触发GraphBreak：

调用未被注册为可追踪的Python函数
依赖tensor真实值的控制流
执行不支持的Python操作，例如打印、异常处理
图断裂后会导致性能降低，Dynamo会尽量通过内联函数、自动处理常见的Python模式来减少断裂，某些情况则无法避免了。torch._dynamo.explain可以查看图断裂的位置和原因。

开发提示：

通过添加 @torch._dynamo.allow_in_graph 可以声明函数为可追踪
注册自定义函数到Dynamo追踪表

import torch._dynamo as dynamo
@dynamo.register("my_custom_fn")
def my_custom_fn(x): 	
	return x.sum()

使用 torch.compile 的 fullgraph=True
使用 torch._dynamo.disable 排除特定函数

1.2 输出FX GraphModule

经上述步骤，Dynamo最终产出一个torch.fx.GraphModule，图中包含：

ATen算子调用
Python函数调用，自定义函数会尝试内联或者保留
可能存在多子图

1.3 关键优化和特性

函数内联
全局变量和闭包处理：Guard可以捕获闭包中引用的变量，并进行Guard来检查他们是否变化
torch._dynamo.optimize装饰器：允许用户手动控制编译过程。
动态shape支持：通过torch._dynamo.mark_dynamic可以标记某些动态维度为动态，最终使用ProxyShape来生成更通用的代码。

2. AOTAutoGrad

Ahead-of-Time Autograd 是 torch.compile() 流水线的第二个核心阶段，位于 Dynamo 之后。它的职责一句话概括：在编译期把反向计算图也生成出来，与前向图合并，并完成算子分解，为 Inductor 提供一张”纯函数、无副作用、可任意融合”的完整计算图。

为什么这件事必须发生在编译期，而不是直接复用 PyTorch 现有的 eager 模式 autograd？这就是理解 AOTAutograd 的起点。

2.1 为什么需要 AOTAutograd

Eager 模式的 autograd 是一个”运行时”机制：

前向执行时，每个算子都会把”如何求导”记录到一张运行时动态构建的图上（autograd graph）
反向时，autograd 引擎从 loss 出发，沿着这张图逐算子反向传播
为了反向，前向的每个中间结果都必须保存在内存里，直到反向用完

这个机制对编译器来说有两个致命问题：

反向图在编译期不可见：Inductor 只能看到前向图，反向是运行时”黑盒”，无法提前优化
图里有”副作用”：x.add_(1) 这样的原地修改、view/reshape 这样的视图操作，都会限制编译器的融合和重排自由

AOTAutograd 的解决方案是：把求导过程从运行时搬到编译期。Dynamo 给出前向 FX 图后，AOTAutograd 直接对这张图做”符号求导”，展开出反向图，然后把前向+反向合并成一张 Joint Graph，交给 Inductor。

flowchart LR
    A[Dynamo 输出
前向 FX 图] --> B[函数化
Functionalization]
    B --> C[符号求导
生成反向图]
    C --> D[Joint Graph
前向 + 反向合并]
    D --> E[算子分解
Decomposition]
    E --> F[交给 Inductor]

2.2 核心步骤

2.2.1 函数化（Functionalization）

图编译器（Inductor）只理解纯函数：输入进来，输出出去，中间没有任何”修改现有数据”的操作。但 PyTorch 代码里充满了 in-place 操作，例如：

1
2
3

x.add_(1)          # 原地加 1
x.relu_()          # 原地 relu
y = x.view(2, 3)   # view 返回共享内存的新 tensor

函数化的做法是用 FunctionalTensor（torch/_subclasses/functional_tensor.py）拦截这些操作，改写成等价的纯函数形式：

# x.add_(1) 变成：
tmp = x + 1
x.copy_(tmp)       # 最终的写回保留，但语义上"读旧值 → 算新值 → 写回"是显式的

# view 则保留为视图操作，但会额外记录它与原始 tensor 的关联，
# 保证后续任何对 y 的修改都能被追踪回 x

对 Inductor 的意义：它拿到的图里没有任何隐式副作用，融合、重排、内存复用（memory planning）都可以放心做。这也是为什么 Inductor 的代码里会看到这样的注释（compile_fx.py 中）：「by the time we’re in inductor, we assume that AOTAutograd has already “taken care” of autograd」——Inductor 从来不需要理解 autograd 的存在。

2.2.2 符号求导，生成反向图

函数化之后，AOTAutograd 对前向图上的每个算子查它的反向规则（例如 add 的反向是把梯度原样传回两个输入），从输出端向输入端”反向追踪”，构建出一张完整的反向 FX 图。注意这不是运行时数值求导，而是编译期展开求导过程——反向图也是图，同样可以交给编译器优化。

2.2.3 Joint Graph：前向+反向合并成一张图

反向图生成后，AOTAutograd 把前向图和反向图拼成一张联合图。这样做的目的是让优化空间最大化：

前向和反向重复的计算可以只算一次（消除公共子表达式）
跨前反向边界的算子融合成为可能（例如前向的激活函数可以直接融合进反向的使用点）
在图层面统一做死代码消除等优化

注意：Joint Graph 的”切分”（把它重新拆回前向子图和反向子图）这个动作在不同版本的 PyTorch 中位置发生过变化。经典实现在 AOTAutograd 内部（torch/_functorch/partitioners.py 的 min_cut_rematerialization_partition），而当前版本中，Joint Graph 会直接交给 Inductor，由 Inductor 在 joint_graph_passes（torch/_inductor/fx_passes/joint_graph.py）中处理前反向的拆分与融合。学习时以本版本代码为准。

2.2.4 激活内存：保存还是重算？

前向计算出的中间结果（激活值），反向计算时需要再次用到。这里有一个经典权衡：

保存：前向结果存在内存里，反向直接用 → 省计算，费内存
重算（recompute）：反向时重新算一遍前向的中间结果 → 省内存，费计算

这个决策由 torch._functorch.config.activation_memory_budget（默认 1.0，即全部保存）等配置控制。对自研硬件来说，当片上内存紧张时，这是一个重要的调优旋钮。

2.3 算子分解（Decomposition）—— 与 Inductor 关系最大的一节

为什么分解

文章开头提到”算子拆分之后会多了很多可融合的 loop”。展开说：PyTorch 的高层算子（如 addmm、layer_norm、softmax）在 ATen 里是一个个”黑盒”实现，Inductor 无法看到内部结构，也就无法把它们内部的逐元素计算和周围的算子融合。把它们分解成基础算子（加减乘除、sum、max 等）之后，图上的小算子变多了，Inductor 的融合器可以把它们和相邻算子重新组合成更高效的内核。

两张分解表

分解不是 AOTAutograd 一个组件的事，实际存在两层：

分解表	位置	使用时机
`core_aten_decompositions`	`torch/_decomp/__init__.py`	AOTAutograd 阶段，对 Joint Graph 分解
Inductor 自己的分解表	`torch/_inductor/decomposition.py`	Inductor 的 post_grad 阶段

分解是”有条件”的——这正是自研硬件的控制点

分解函数返回 NotImplemented 就表示”这个算子我不拆”。以 decomposition.py 中的 addmm 为例（torch/_inductor/decomposition.py）：

@register_decomposition([aten.addmm])
def addmm(self, mat1, mat2, ...):
    # 只有部分形状（matvec 等）才值得拆成 mul + sum
    if mat1.device.type not in ["cpu", "mps"]:
        # 非 CPU/MPS 设备默认返回 NotImplemented —— 不拆！
        ...
        return NotImplemented
    ...

注意这里的 device 判断：matmul 类算子在非 CPU/MPS 设备上默认不分解。这就是前面 NPU 话题里讲过的第一道关卡——如果你在做自研硬件适配，希望某些算子（比如有专用加速单元的 matmul）保持原样不被拆成加减乘除，控制点就是这里：在分解表中让对应算子在你的设备上返回 NotImplemented。第二道关卡在 lowering 层（算子是被内联成循环还是保留为外部调用），那属于 Inductor 的内容，后面章节展开。

2.4 Inductor 在 AOTAutograd 输出上的介入点

AOTAutograd 做完之后，Inductor 接手的是一张（或几张）纯函数 FX 图。Inductor 在其上做的事情按顺序是：

Joint Graph Passes（torch/_inductor/fx_passes/joint_graph.py）：在前向+反向联合图上做跨边界优化（模式匹配、常量折叠等）
前向/反向子图各自的 Post-Grad Passes：分解、融合模式匹配、冻结等
Lowering + Scheduler + Codegen：后面的章节展开

如果你需要在这条链路上插入自己的图变换，Inductor 提供了官方的自定义 pass 挂载点（torch/_inductor/config.py 中）：

挂载点	时机
`joint_custom_pre_pass` / `joint_custom_post_pass`	Joint Graph 优化前后
`post_grad_custom_pre_pass` / `post_grad_custom_post_pass`	前向/反向子图优化前后

实现方式是实现 torch._inductor.custom_graph_pass.CustomGraphPass 接口，然后赋给对应的 config 项。

三层图变换的归属（小结）

Dynamo 管”抓图”：把 Python 代码变成 FX 图
AOTAutograd 管”造反向图和拆算子”：函数化、符号求导、Joint Graph、算子分解
Inductor 管”融合和生成代码”：图优化 pass、Lowering、Scheduler、Codegen

2.5 调试 AOTAutograd

设置 TORCH_LOGS="+aot" 可以打印 AOTAutograd 阶段的日志
aot_config.aot_id（torch/_functorch/_aot_autograd/graph_capture.py）支持导出调试图
用我们自己写的调试脚本 agent_space/inductor_debug.py --breakpoint 在 _compile_fx_inner 处停住时，看到的 gm 就是 AOTAutograd 处理完交给 Inductor 的图——前向子图或 Joint Graph，这是我们观察 AOTAutograd 输出最直接的方式

3. Inductor

Inductor 是 torch.compile() 的默认后端，也是我们整个系列的主角。前面两节讲了 Dynamo 如何”抓图”、AOTAutograd 如何”造反向图和拆算子”，从这一节开始，图真正进入编译。

3.1 Inductor 是什么

一句话：Inductor 是一个输入 FX 图、输出可执行代码的后端编译器。它吃下 AOTAutograd 产出的纯函数图，吐出可以直接调用的 Python 模块（host 端代码 + 编译好的内核二进制）。

和传统编译器（GCC/LLVM）不同，Inductor 的工作对象不是标量指令，而是张量运算；它的优化目标也不是寄存器分配，而是算子融合、循环变换、内存复用。

3.2 核心设计理念：一切围绕”循环”

Inductor 最重要的设计选择是：它的中间表示（IR）不是”算子列表”，而是循环结构。

原因很朴素：GPU/CPU 上的张量计算，最终都是嵌套循环。add 是循环，relu 是循环，sum 是带累加的循环，卷积拆开也是循环。如果把所有计算都表示成循环，那么：

算子融合 = 循环合并：两个相邻算子的循环拼成一个，省掉中间结果的读写
循环变换（分块、重排、向量化）就是优化手段本身

所以 Inductor 的整个后端都在做一件事：把 FX 图上的算子 lower 成循环 IR，然后合并、重排这些循环，最后生成代码。

3.3 编译流水线概览

flowchart LR
    A[FX 图
纯函数、已分解] --> B[图优化 Pass
模式匹配、算子级融合]
    B --> C[Lowering
算子 → 循环 IR]
    C --> D[Scheduler
融合、排序、内存规划]
    D --> E[Codegen
IR → 内核源码]
    E --> F[编译与缓存
生成 .so]
    F --> G[加载执行]

四个阶段各司其职：

图优化（FX Passes）：在图层面做模式匹配与替换（比如把 conv + relu 合并成一个融合算子调用），这仍是”算子粒度”的优化
Lowering：把每个算子翻译成循环 IR。这是”算子 → 循环”的边界
Scheduler：核心中的核心。决定哪些循环合并成一个内核、按什么顺序执行、中间结果复用哪些内存
Codegen：把融合好的循环结构生成具体的内核源码——GPU 上生成 Triton 代码，CPU 上生成 C++/OpenMP 代码——再调用编译器编译，缓存起来复用