背景概述

随着深度学习模型规模的不断扩大，模型训练和推理的性能优化变得越来越重要。PyTorch 2.x 引入了 torch.compile 机制，通过即时编译技术将 Python 级别的算子图转换为高效的 NPU 内核代码，从而显著提升模型执行效率。理解 torch.compile 的整体工作流程，对于开发者优化模型性能、排查性能瓶颈具有重要意义。

本文以一个简单的张量运算示例为切入点，详细剖析 torch.compile 从 FX 图生成到最终 Inductor IR 产出的完整流程，重点介绍 post_grad_pass 优化阶段和 lowering 降级阶段的核心处理逻辑。

示例程序

@torch.compile
def simple_optimization_demo(x: torch.Tensor):
    a = x + 0.0         
    h = a - 2.0
    b = a * 1.0          
    c = b - 0.0          
    d = c + 2.0
    e = torch.relu(d)    
    f = torch.sigmoid(e)
    g = f.view(-1)
    return g

# 测试
x = torch.randn([4, 32], device="npu")
result = simple_optimization_demo(x)

在得到 FX 图之后，torch.compile 会依次执行以下核心处理流程。

---

1. 准备工作

1.1 记录输入 Shape

在编译开始时，系统会记录输入张量的形状信息，为后续的算子融合和内存布局优化提供基础数据。

1.2 对图进行预处理

为了支持内存布局优化（如将 NCHW 格式转换为 NHWC 格式），需要将计算图中的 view 操作替换为 reshape 操作。这一转换通过 view_to_reshape 函数完成，详细原理可参考相关技术文档。

1.3 记录原始输出 Stride

保存原始输出的 stride 信息，确保后续优化不会破坏张量的内存访问模式。

---

2. post_grad_pass 优化阶段

post_grad_pass 以递归方式处理计算图中的每个子图，执行一系列优化操作，为后续的代码生成做准备。

2.1 消除死代码（eliminate_dead_code）

遍历整个图的节点，如果某个节点不是输出节点，且没有任何其他节点依赖于该节点的输出，则该节点可以被安全移除。这一步骤可以有效减少冗余计算。

2.2 局部重排序（reorder_for_locality）

该优化仅在推理模式下启用，通过重新排列节点执行顺序来优化访存局部性，提升缓存命中率。

2.3 用户自定义预处理优化

系统会从配置文件中检测用户是否定义了自定义的预处理优化规则，并据此调整优化策略。

2.4 批量操作识别与融合

初始时 post_grad_fusion_options 为空，运行时系统会注册 batch_aten 的加减乘除 pattern。随后通过 generate_fusion_from_config 检查配置信息，遍历所有融合规则进行模式匹配，匹配成功后更新图中的节点。

进入 apply_group_batch_fusion 后，系统从输出到输入逆向遍历所有节点。从当前节点出发，根据特定的融合规则寻找所有可融合的候选节点，同时排除已在 fused_set 中标记为已融合的节点。当候选节点数量大于最小融合集大小时，执行融合操作；使用贪心策略从候选中找到满足 rule.graph_search_options 的独立子集进行融合。

2.5 消除无意义指令（remove_noop_ops）

从节点表中读取 placeholder 节点，遍历所有节点，将满足特定表达式的节点在保证不破坏输入输出内存共享关系和现有别名的条件下进行替换。具体而言，若能有效提取节点参数，且发现当前节点与源节点的元数据相同、替换前后不改变程序语义，则将所有对当前节点的引用替换为对源节点的引用，并删除当前节点。

2.6 预定义模板匹配

应用预定义的模板匹配规则，识别可优化的计算模式。

2.7 并行运算与通信融合

根据配置选择性能优化策略，逐个处理用户指定的优化选项，找到对应的优化 pass，比较优化前后的结果并记录有效优化。同时实现并行运算、DDP 通信模块融合，处理用户自定义的 pass，并通过稳定拓扑排序保证节点的相对顺序。涉及的关键函数包括：

• _fuse_ddp_communication_passes
• _fuse_ddp_bucket_size
• post_grad_custom_post_pass
• stable_topological_sort

2.8 分解高阶算子（decompose_auto_functionalized）

注册一个 graph_pattern，匹配 torch.ops.higher_order.auto_functionalized 和 triton_kernel_wrapper_functional 的算子并执行分解操作。

2.9 重新编译更新 FX 图

完成上述优化后，重新编译并更新 FX 图，确保后续阶段基于优化后的图结构进行代码生成。

---

3. Lowering 降级阶段

Lowering 阶段的主要功能是将 post_grad_pass 处理后的节点翻译为 pointwise/computeBuffer 的 Loop level IR。

3.1 AOT 模式下的常量折叠

当处于 AOT 模式且启用常量折叠时：

1. 分离常量图：将常量计算部分分离出来，对常量图进行编译和代码生成，实现编译期完成所有常量计算。
2. 简化主图：使主图仅包含真正的动态计算部分，减少运行时开销。

3.2 节点到 Inductor IR 的转换

在 graph.run(*example_inputs) 中，对图中的每个节点进行转换，生成对应的 Inductor IR。具体实现方式为：根据每个节点的 op 类型，调用 lowering 中对应的函数完成初步 Inductor IR 的生成。如果算子不在 lowering 中，则检查其是否存在于白名单和 implicit_fallbacks 配置中。最终根据 node.op 找到对应的执行函数，由解释器调用 class GraphLowering(torch.fx.Interpreter) 中的相应函数。

3.3 数据类型检查与内核代码生成

后续检查数据类型是否受支持，然后通过 graph.compile_to_fn() 编译生成内核代码，最终产出 Inductor IR。

以下为示例程序经过 lowering 后生成的 Inductor IR 片段（设备名为示意）：

TensorBox(StorageBox(
  Pointwise(
    'npu',
    torch.float32,
    def inner_fn(index):
        i0, i1 = index
        tmp0 = ops.load(arg0_1, i1 + 32 * i0)
        tmp1 = ops.constant(0.0, torch.float32)
        tmp2 = tmp0 + tmp1
        return tmp2
    ,
    ranges=[4, 32],
    origin_node=None,
    origins=OrderedSet([add])
  )
))

从上述 IR 可以看出，连续的逐点算子（如加法、乘法、减法等）被逐步融合到同一个 Pointwise 计算中，最终仅需一次内核启动即可完成全部计算，大幅减少了内核启动开销和内存访问次数。

---

总结

torch.compile 的整体流程可以概括为：FX 图构建 → post_grad_pass 优化 → Lowering 降级 → 内核代码生成。其中 post_grad_pass 阶段通过死代码消除、算子融合、无意义指令移除等一系列优化手段，显著精简了计算图结构；lowering 阶段则将优化后的图转换为高效的 Inductor IR，最终生成高性能的 NPU 内核代码。理解这一流程有助于开发者更好地利用 torch.compile 进行模型性能优化，也为深入探索 PyTorch 编译栈提供了清晰的脉络。