TorchInductor Pipeline 与算子融合

torch.compile Walkthrough 用 ((a + b) * c.relu()).sum(dim=-1) 走了一遍 torch.compile 的各个阶段，add、relu、mul、sum 四个算子最后变成一个 kernel。那个例子太小，看不出 Inductor 是怎么决定融合的：哪些节点能放进同一个 kernel，有冲突时先合并哪一对。这篇换一个稍大的例子，专门看这个决定的过程。

例子有 10 个 ATen 算子，有分支，也有 reduction。x 是 f32[1024, 1024]，bias、gamma、beta 是 f32[1024]：

def model(x, bias, gamma, beta):
    t = x * 1.1 + bias                # broadcasting + pointwise
    a = torch.relu(t)
    p = a * a                         # two branches
    q = a * 0.1
    s = p + q                         # merge branches
    m = s.mean(dim=1, keepdim=True)   # reduction
    y = (s - m) * gamma + beta        # pointwise epilogue
    return y

Eager 模式下它是 10 个 kernel。torch.compile 之后只剩一个：

triton_per_fused_add_mean_mul_relu_sub_0

这个 kernel 不是在某一步生成的，四个阶段各做了一部分决定。 下文按 TorchInductor 的流水线顺序讲，篇幅主要给 Scheduler：can_fuse 判断两个节点能不能融合，score_fusion 决定有冲突时先融合哪一对，fuse_nodes 是外层的融合循环。最后把同一个例子交给 TVM 的融合策略，比较两者的结果。

01Prerequisites#

这些背景在别的文章里讲过：

1.1Tensor Layout#

一个 tensor 是一块一维内存加上 size 和 stride。元素 的位置是

连续的 [1024, 1024] 张量 stride 是 [1024, 1]，offset 就是 1024·i + j。Inductor 日志里的 index 表达式都是这个式子：loop body 里的 index0 = 1024*p0 + p1，MemoryDep('arg0_1', 1024*d0 + d1, {d0: 1024, d1: 1024})，后面的字典是每个循环变量的范围。

广播不复制数据。x + bias 里，bias 被当成 stride 为 [0, 1] 的 [1024, 1024] 张量：每一行读同样的 1024 个元素，index 里没有行号，日志写成 MemoryDep('arg1_1', d1, ...)。

两个推论在后文会用到：

• view 只改 size 和 stride，不复制，所以要求新形状能用 stride 表达出来。transpose 之后的张量不连续，再 view 可能报错；reshape 在不能 view 时会复制。
• 长度为 1 的维度，stride 取什么值都指向同一个元素。buf0 的 size=[1024, 1], stride=[1, 1024] 和 stride=[1, 1] 是同一个布局。

1.2Iteration Space#

Inductor 生成的每个 kernel 是一个循环嵌套，Scheduler 用 group.iteration = (numel, rnumel) 描述它：

# Pointwise，例如 op1：group.iteration = (1048576, 1)
for x in range(1048576):        # numel：输出元素个数
    buf1[x] = ...

# Reduction，例如 op0：group.iteration = (1024, 1024)
for x in range(1024):           # numel：输出元素个数
    acc = 0
    for r in range(1024):       # rnumel：每个输出要归约的元素个数
        acc += ...
    buf0[x] = acc

两个节点要合进一个 kernel，就要能放进同一个循环嵌套。numel、rnumel 都相同的节点可以直接合并；pointwise (N, 1) 和 reduction (M, R) 要求 N = M × R，并且 pointwise 的循环能拆成 (M, R) 两层，这是后文后端检查的内容。生成 Triton 代码时，x 维分给各个 program（tl.program_id(0) 和 XBLOCK），r 维在 program 内部用 R0_BLOCK 遍历；一行能一次放进 block 时就是 persistent reduction。

1.3Functionalization#

原地修改和 view 会让依赖关系变得隐蔽：y = x.view(-1); y.add_(1) 改的是 x，但图里看不出 x 被写过。AOTAutograd 的 functionalization 把这类程序改写成没有 mutation 的形式（示意）：

# 原程序                        # functionalize 之后
def f(x):                       def f(x):
    x.add_(1)                       x1 = aten.add.Tensor(x, 1)
    return x * 2                    out = aten.mul.Tensor(x1, 2)
                                    aten.copy_.default(x, x1)   # 对输入的修改最后写回
                                    return out

之后每个节点都是纯函数，节点之间的依赖就是图里的边。Inductor 的 post-grad pass 在这样的图上改写、重排、删除节点，不用担心漏掉隐藏的写；最后再由 reinplace_inplaceable_ops 把能原地做的操作改回原地。

同一阶段还会用到 FakeTensor：只有 shape、dtype、stride、device 等元信息、没有数据的 tensor。在它上面执行算子只推导输出的元信息，不做计算，Inductor 靠它在编译时知道每个中间结果的形状。

1.4Dependency Graph#

Scheduler 把节点和它们之间的读写依赖看成一张有向无环图（DAG）。几个概念：

• 拓扑序：每个节点都排在它依赖的节点之后。执行顺序必须是一个拓扑序。
• 祖先：沿依赖边往回能走到的所有节点。node1 是 node2 的祖先，说明 node2 直接或间接用了 node1 的结果。
• DCE（dead code elimination）：删掉结果没人用的节点。
• CSE（common subexpression elimination）：相同的表达式只算一次。t0 = load(x[i]) 之后再出现 load(x[i])，直接复用 t0。Inductor 生成 kernel 代码时做 CSE。

融合两个节点，等于要求它们在同一个 kernel 里一起执行。如果两者之间隔着第三个节点，就可能形成环：

所以 can_fuse 和融合循环都要检查中间节点和环。

1.5Extern Kernel and Template#

Inductor 自己生成的主要是 pointwise 和 reduction kernel。矩阵乘、卷积默认调用外部库（cuBLAS、cuDNN），在 IR 里是 ExternKernel。外部 kernel 是黑盒，Inductor 不能往里面加代码，所以融合在它的两侧断开：它的输入必须先写进内存，它的输出也是一个真实的 buffer。

max_autotune 打开后，Inductor 会为 mm 等算子测量多个候选实现，其中包括用 Triton 模板生成的 kernel（IR 里是 TemplateBuffer）。模板 kernel 是 Inductor 生成的代码，可以继续往里加东西：

• Epilogue fusion：在模板写回结果之前，接上后续的 pointwise，例如 mm 之后的 + bias 和 relu。
• Prologue fusion：在模板读输入时，先做前置的 pointwise，例如对输入做 dtype 转换。

02Operator Fusion#

算子融合用一个更大、更复杂的 kernel，换更少的 kernel launch 和更少的中间数据读写。三个算子 A → B → C 顺序执行时，每个中间结果都要写回 HBM 再读出来；融合后，中间结果留在寄存器里：

收益有三项：

• Global memory 读写减少，通常是最大的一项。
• Kernel launch 减少，每次 launch 都有固定的 CPU 和调度开销。mode="reduce-overhead" 用 CUDA Graph 解决的也是这个问题，只是从另一个方向入手。
• 数据局部性更好：上一步的结果在寄存器里直接给下一步用。

在本例上估算一下第一项。一个 [1024, 1024] 的 f32 张量是 4 MiB。Eager 模式下，10 个 kernel 里有 9 个读一个完整张量，p + q 读两个，除 mean 外每个都写一个完整张量：

访存量差了 10 倍，这是 pointwise 和 reduction 这类 memory-bound 算子值得融合的原因。

代价也有三项：

• 寄存器压力：融合后的 kernel 要同时保存多个算子的中间值。需要的寄存器超过可用的数量时会发生 register spill，溢出的值进 local memory，省下的访存又加回来一部分。
• 代码变大：instruction cache 压力更大，编译时间也更长。
• 并行配置只能选一套：分开执行时，每个 kernel 可以各自选 block size、num_warps、tile size；融合后一套参数要同时服务多个算子，不一定对每个都最优。

所以融合不是“能融就融”。编译器要回答三个问题：能不能融（合法性），值不值得融（收益估计），候选冲突时先融哪个（顺序）。Inductor 分别用 can_fuse、score_fusion 的分数和 fuse_nodes 的贪心循环回答它们。

03torch.compile#

Inductor 是 torch.compile 的默认 backend。torch.compiler.list_backends() 列出可用的 backend：inductor、cudagraphs、openxla、tvm。它们都接收 Dynamo 捕获的 FX Graph；cudagraphs 是一种 GPU 执行机制，用 CUDA Graph 降低重复执行的调度开销，另外三个是面向不同硬件的深度学习编译器。调试用的 eager（只做 Dynamo 捕获）和 aot_eager（再经过 AOTAutograd）不出现在列表里，它们不生成优化后的 kernel，用来定位问题出在哪一层。

Inductor 的入口是 _compile_fx_inner：

@time_and_log(attr="compilation time (in seconds)")
def _compile_fx_inner(
    gm: GraphModule,                          # FX 图和执行它需要的模块属性
    example_inputs: Sequence[InputType],      # 示例输入：推导 shape、dtype、device
    compile_region_name: str | None = None,
    **graph_kwargs: Unpack[_CompileFxKwargs], # 是否 CUDA Graph、是否 backward、是否 C++ wrapper 等
) -> OutputCode:

它收到的 gm 已经经过 Dynamo 捕获和 AOTAutograd 处理，节点都是 aten.* 算子。上层的 compile_fx() 和 fx_codegen_and_compile() 根据配置选编译策略，默认用 _InProcessFxCompile，实际工作在它的 codegen_and_compile() 里。

04TorchInductor Pipeline#

Inductor 的主流程分为四个阶段，每个阶段回答一个问题：

• Graph Preprocessing：What to compute？ 规范化 FX Graph、补全元信息、做图级优化：哪些计算保留、删除、替换或组合。输出是预处理后的 FX Graph。
• Lowering：How to represent？ 决定算子怎么表示、是否分解、是否 fallback、何时物化结果。输出是 Inductor IR，包含 IR 操作、buffer、布局和依赖信息。
• Scheduling：How to organize？ 决定执行顺序、算子融合、部分重排。输出是 SchedulerNode / FusedSchedulerNode。
• Code Generation：How to emit？ 决定循环与索引表达式、tile 大小、warp 配置、模板实现和调用方式。输出是 kernel 代码和调用它们的 wrapper。

最后编译、封装成 CompiledFxGraph（AOTInductor 是 CompiledAOTI）。

“融合”在四个阶段里是四件不同的事：Preprocessing 里是 pattern matching 把一个子图换成另一个子图；Lowering 里是把 Pointwise 内联进使用者的表达式；Scheduling 里才是把多个调度节点合成一个 kernel；Codegen 里是 kernel 内的公共子表达式消除和删掉只在 kernel 内部使用的 buffer。本例的 10 个算子到 1 个 kernel，后三件都用上了。

05Graph Preprocessing#

在创建 GraphLowering 之前，codegen_and_compile() 依次做了这些处理：

• view_to_reshape：把 aten.view.default 节点的 target 改成 aten.reshape.default。view 要求输入的 stride 能直接表达新形状，否则报错；Inductor 之后可能改变中间结果的布局（padding、channels-last 等），原来合法的 view 可能不再合法。reshape 在不能 view 时会退回到 copy，所以更安全。
• fake_tensor_prop：如果能从输入里检测到已有的 FakeTensorMode 就复用，否则新建一个；然后用 FakeTensorProp(gm, mode).propagate(*example_inputs) 沿图执行每个算子，把结果的元信息挂到节点上。这一步不做真实计算。
• _recursive_record_original_output_strides：先递归处理 invoke_subgraph 高阶算子里的子图，再记录当前图输出的 stride。Inductor 内部可以给中间结果选别的布局，但图的输出必须和 eager 的 stride 一致。

5.1post_grad_passes#

_recursive_post_grad_passes 先递归处理子图，再对当前图调用 post_grad_passes(gm, is_inference)。它对 forward 图和 backward 图各调用一次，此时的图已经规范化、functionalize 过。和融合有关的有两类：模式匹配（pass_patterns 等）在图里找到预定义的算子组合，换成等价的实现；_remove_profiler_ops 删掉 record_function 等 profiler 节点，避免它们挡住后面的融合。其余是死代码消除、CPU 专用优化、分布式通信优化和收尾的分解。

reinplace 放在最后有原因：前面所有 pass 都在 functional 图上工作，没有 mutation 和 alias，改写更容易保证正确；改完之后再把能原地的操作恢复成原地。

本例里这些 pass 都没有改动图：fx_graph_transformed.py 和 fx_graph_runnable.py 的算子序列相同，只多了 shape 标注和源码行注释。

回到 codegen_and_compile()，post-grad pass 之后它记录优化后的图、新旧节点的对应关系和编译指标，然后为 lowering 设置环境：V.set_fake_mode(fake_mode) 让后面的 shape 推导都走 FakeTensor；maybe_disable_comprehensive_padding 和 maybe_disable_graph_partition 按条件关掉 padding（给张量补元素来调整布局）和图分区（把一张图拆成多个子图）。

06Lowering#

torch.ops.aten.relu.default(x) 只说了“对张量做 ReLU”，没有说由哪个 kernel 执行、每个线程负责哪些元素、数据怎么分块和存取。Lowering 把这样的节点转换成 Inductor 能分析的表示：每个输出元素怎么算，读哪些 buffer，结果要不要写进内存。

6.1GraphLowering.run_node#

GraphLowering 是 torch.fx.Interpreter 的子类。GraphLowering.run() 只在父类的 run() 外面加了计时，遍历节点的逻辑来自父类；每个节点的 lowering 由子类重写的 run_node() 完成：

Interpreter.run_node() 用 fetch_args_kwargs_from_env(n) 把 FX 节点的输入引用解析成之前的 lowering 结果，再用 getattr(self, n.op)(n.target, args, kwargs) 分派。对 b = torch.relu(a)，FX 图里 b 的节点只引用节点 a；Inductor 在 env 里查到的是处理 a 时返回的对象，它不一定是已经物化的 Tensor，可以是一个延迟计算的 TensorBox。

GraphLowering 只支持 placeholder、call_function、get_attr、output 四种节点，call_method 和 call_module 直接 raise AssertionError。这不会造成问题：AOTAutograd 已经把方法调用和模块调用都展开成了 aten.* 的 call_function。

以 placeholder 为例，它从 self.args_iter 取得当前输入，对 Tensor 做三件事：

tensor = TensorBox.create(
    DonatedBuffer(                     # 1. 表示图输入的 buffer
        name=target,
        layout=FixedLayout(example.device, example.dtype, sizes, strides),  # 2. 设备、类型、形状、stride
    )
)                                      # 3. 包装成 Inductor 的 Tensor 表示

GraphLowering.run_node() 在进入父类分派之前还做了几件事：

• 记录 origins：当前 FX 节点，以及参数来自哪些节点。之后生成的 kernel 名（fused_add_mean_...）和源码注释都从这里来。
• _realize_inputs_at_stream_boundaries(n)：如果输入和当前节点在不同的 CUDA stream 上，提前把输入物化，防止跨 stream 的计算被内联到一起，让 Scheduler 能把它们拆成不同的 kernel。
• 检查几种特殊情况，见下面的折叠块。

6.2call_function#

大部分节点走 call_function。去掉日志和少数特例后，它的逻辑是：

def call_function(self, target, args, kwargs):
    if target not in lowerings:                       # 算子还没有 lowering
        if base_name in FALLBACK_ALLOW_LIST:
            make_fallback(target)                     # 注册一个 fallback lowering
        elif config.implicit_fallbacks:
            make_fallback(target, layout_constraint=...)   # 隐式 fallback
        elif get_decompositions([target]):
            raise MissingOperatorWithDecomp(...)      # 有分解规则，但没有 lowering
        else:
            raise MissingOperatorWithoutDecomp(...)

    layout_constraints = maybe_layout_constraints(target)
    if layout_constraints:                            # 有些算子要求连续布局或特定 stride
        old_args, old_kwargs = args, kwargs
        args, kwargs = layout_constraints(n, *args, **kwargs)   # 必要时复制输入

    out = None
    if target in user_lowerings and target not in V.active_user_lowering_ops:
        out = user_lowerings[target](*args, **kwargs)  # 用户注册的 lowering 优先
    if out is None:                                    # 没有，或者返回了 None
        out = lowerings[target](*args, **kwargs)       # 标准 lowering

    if layout_constraints:
        self.propagate_mutation(n, old_args, old_kwargs, args, kwargs)  # 把 mutation 写回原始输入
    return out

lowerings 是 Inductor 的 lowering 注册表，每一项是一个算子的 lowering 函数，用 @register_lowering 注册。分享时的版本里有 1811 项：

>>> from torch._inductor.lowering import lowerings
>>> len(lowerings)
1811
>>> lowerings.get(torch.ops.aten.add.Tensor)
<function make_pointwise.<locals>.inner at 0x79070c01ae50>

aten.add 的 lowering 是 make_pointwise 生成的闭包，返回一个 Pointwise；view 类算子返回 View，比如 permute：

@register_lowering(aten.permute, type_promotion_kind=None)
def permute(x, dims):
    assert isinstance(x, TensorBox)
    assert isinstance(dims, (list, tuple))
    return TensorBox(PermuteView.create(x.data, tuple(dims)))

Fallback 对融合有直接影响：make_fallback 注册的 lowering 直接调用 ATen kernel，生成 ExternKernel。它只能读写真实内存，所以输入要 realize，输出是一个 buffer，它在 Scheduler 里成为 ExternKernelSchedulerNode，一般的融合在它两侧断开。一个模型里如果有算子掉进 fallback，融合的机会也随之减少。

Lowering 返回后，run_node() 还要处理结果：

• 约束输出的 layout 和 stride：根据输出是否对用户可见、是否涉及 as_strided 等条件，调用 ir.ExternKernel.require_stride_order(...) 或 require_exact_strides(...)。

• 决定是否 realize。realize 不表示立刻执行 GPU 计算，而是让编译器明确“这个结果要存下来”，具体代码在后面生成。例如：

  python

  # Realize if the IRNode already has accumulated lots of reads
  if isinstance(result, TensorBox) and result.has_exceeded_max_reads():
      # Prevent excessive accumulation in a computed buffer, when
      # there are multiple branches each with small number of memory
      # reads, but they converge to a user.
      result = maybe_apply_channels_last_stride_order(result, n)
      result.realize_hint()

• 记录复用信息和符号形状信息。

最终，GraphLowering 把 FX 节点都转换成 Inductor IR，生成的 IR 操作记录在 self.operations 里。

6.3Inductor IR#

Inductor IR 有五类核心对象：

• TensorBox：对应一个上层的 tensor，.data 指向一个 StorageBox（中间可能隔着 View）。
• StorageBox：代表一块存储，所有别名指向同一个 StorageBox。realize() 时把里面的延迟节点替换成 ComputedBuffer，分配内存。
• View：StorageBox 外的坐标变换层，处理上层的 view 操作。
• Buffer：一块已分配内存的计算结果，有名字、有 layout，是 Scheduler 调度的基本单元。
• Pointwise / Reduction：纯粹的延迟计算描述，只有一个 inner_fn，没有内存，不参与调度。

为什么要多一层 StorageBox？因为别名。b = a 之后 a、b 是两个 TensorBox，指向同一个 StorageBox；对 a 的原地修改替换的是 StorageBox.data，b 自然看到新值。a.permute(1, 0) 得到的 c 不复制数据，只在 StorageBox 外加一层重新索引：

这是模型上的解释。AOTAutograd 的 functionalization 之后，图里一般只剩把结果写回输入的 copy_，很少再有 add_ 这样的原地操作。

没有 realize 的 Pointwise 只是一个函数。使用者通过 loader 读它时，实际调用的是它的 inner_fn，所以一串 pointwise 会层层内联，直到遇到需要 realize 的节点：

e = torch.sum(d) 分两步：先调用 d.make_loader() 拿到 relu_fn，把它内联进 Reduction 的 inner_fn；然后 make_reduction 主动调用 realize()，因为 reduction 的输出 shape 和输入不同，必须写进内存才能被后续节点使用。realize 前后，同一个 TensorBox 的 .data 链是：

StorageBox.realize() 就地把 self.data 换成一个包着它的 ComputedBuffer，这个 buffer 就是 Scheduler 可调度的计算单元。除了 reduction，graph 输入、buffer mutation、CUDA stream 边界、特定的 stride / layout 约束、输出是 view、多个使用者导致重复计算的风险过大、inner_fn 过于复杂，也都会触发 realize。各个阈值见 torch.compile Walkthrough。

6.4示例#

本例 lowering 之后只有两个 buffer：mean 的 reduction 产生 buf0，最后的 epilogue 产生图的输出 buf1。其余 8 个算子都没有 realize：

ir_pre_fusion.txt 里 op0 的 loop body 能看到内联的代价。下面把每串 ops 调用合成了一行：

class op0_loop_body:
    var_ranges = {p0: 1024, p1: 1024}
    index0 = 1024*p0 + p1                  # x[p0, p1]
    index1 = p1                            # bias[p1]
    index2 = p0                            # buf0[p0]
    def body(self, ops):
        relu   = relu(load(arg0_1, index0) * 1.1 + load(arg1_1, index1))  # a，给 a*a 的左边
        relu_1 = relu(load(arg0_1, index0) * 1.1 + load(arg1_1, index1))  # a，给 a*a 的右边
        mul_2  = relu * relu_1                                            # p = a * a
        relu_2 = relu(load(arg0_1, index0) * 1.1 + load(arg1_1, index1))  # a，给 a * 0.1
        mul_4  = relu_2 * 0.1                                             # q = a * 0.1
        add_3  = mul_2 + mul_4                                            # s = p + q
        reduction = ops.reduction(torch.float32, torch.float32, 'sum', add_3)
        ops.store_reduction('buf0', index2, reduction)

• a 被用了三次，就内联了三次：x * 1.1 + bias 和 relu 在 loop body 里出现三遍，x 也 load 了三遍。op1 的 loop body 也是这样从 load(arg0_1, index0) 开始重算。这就是 realize 规则里“多个使用者导致重复计算”的风险。
• 在 IR 里这只是表达式被复制，不一定是真的多算：codegen 生成 Triton 代码时会做公共子表达式消除（CSE），同一个 kernel 里相同的 load 和相同的计算只生成一次。
• reduction 的类型是 sum：buf0 存的是每一行的和，除以 1024 的那一步留给了使用 buf0 的 op1。
• buf0 的 layout 是 size=[1024, 1], stride=[1, 1024]，op1 的输出 buf1 是 size=[1024, 1024], stride=[1024, 1]。

为什么 a 和 s 都有两个使用者，却没有 realize？run_node 在 lowering 之后对有多个使用者的结果调用 mark_reuse(len(n.users))，它是否 realize 由 StorageBox.should_realize_on_reuse 决定。GPU 上相关的部分是：

def should_realize_on_reuse(self, users):
    if users > 1 and isinstance(self.data, (Pointwise, Reduction)):
        if self.has_large_inner_fn():        # inner_fn 的 op 数超过阈值（GPU 默认 30）
            return True
        return self.num_reads() > config.realize_reads_threshold   # 读的不同 buffer 数 > 4
    return False

num_reads() 数的是 inner_fn 读了几个不同的 buffer，不是 load 了几次。a 和 s 都只读 x 和 bias 两个 buffer，2 ≤ 4；s 的 inner_fn 展开后大约 20 个 op，也低于 30。所以 Inductor 判断重算比存下来便宜，两个都内联。如果 s 读了 5 个以上的输入，或者表达式再长一些，它就会成为一个单独的 buffer，op0 和 op1 都改成读它。

到这里 s 有两份副本，各在一个节点里。如果 Scheduler 不把 op0 和 op1 合起来，x 就要从 HBM 读两遍。

07Scheduling#

Lowering 完成后调用 GraphLowering.codegen()：

def codegen(self):
    with dynamo_timed("GraphLowering.codegen", log_pt2_compile_event=True):
        self.init_wrapper_code()                          # 初始化 wrapper 代码生成环境
        self._update_scheduler()                          # 创建 Scheduler，处理 self.operations
        V.debug.draw_orig_fx_graph(self.orig_gm, self.scheduler.nodes)
        self.wrapper_code.push_codegened_graph(self)
        self.scheduler.codegen()                          # 调度 IR 节点，生成 kernel 代码
        result = self.wrapper_code.generate(self.is_inference)  # 生成最终的 wrapper 代码
        self.wrapper_code.pop_codegened_graph()
        return result

7.1Scheduler._init#

Scheduler 初始化时用 self.nodes = [self.create_scheduler_node(n) for n in nodes] 把 ir.Operation 列表转成自己的节点，然后建立依赖：