PyTorch 编译栈指南：在 ATen IR 上编写图变换（FX Transformer / Subgraph Rewriter / PassManager / Partitioner）

导读

本文基于 PyTorch 官方用户指南（docs/source/user_guide/torch_compiler/torch.compiler_transformations.md），系统讲解如何在 ATen IR（即 torch.export 之后得到的 FX Graph / GraphModule）上编写、组织与调度图变换。由于 ATen IR 处于 FX Graph 层，任何为 FX 图编写的变换都能直接套用——读者将掌握四种核心能力：用原生节点遍历改写图、用 torch.fx.Transformer 以"解释执行 + 重建图"的方式做一对一/一对多/一对零变换、用 subgraph_rewriter 做多对一模式替换、以及用 PassManager 与 CapabilityBasedPartitioner 对整图做流水线化调度与按算子能力切分融合。文中所有示例均可在本仓库源码（torch/fx/ 与 torch/fx/passes/）中找到对应的实现佐证。

---

1. ATen IR 与图变换的基本面

在 PyTorch 2.x 编译栈中，torch.export 会把模型导出为 ATen IR——一组由 torch.ops.aten.* 算子节点构成的 FX Graph。该图以 torch.fx.GraphModule 的形式存在，因此：

任何为 FX Graph 编写的变换，都可以直接应用到 ATen IR 上。如果你熟悉编写 FX 图变换，那么在 ATen IR 上做变换是完全一样的。

这意味着 FX 生态中成熟的节点操作 API、Tracer、Interpreter、Transformer、PassManager 等基础设施，全部对 ATen IR 生效。这也是 torch.compile 的分解（decomposition）、算子融合（fusion）、内存规划等优化得以统一构建在 FX 之上的原因。

1.1 变换的最直接方式：遍历并改写节点

最直接的变换方式是遍历给定图，直接操作图中的节点。下面这个例子将图中所有的 torch.ops.aten.add.Tensor 调用替换为 torch.ops.aten.mul.Tensor：

import torch

def replace_add_with_mul(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
    for node in gm.graph.nodes:
        if node.op == "call_function" and node.target == torch.ops.aten.add.Tensor:
            node.target = torch.ops.aten.mul.Tensor

要点解析：

• gm.graph.nodes 返回图内全部节点，可按拓扑顺序遍历；
• node.op == "call_function" 过滤出函数调用节点；
• node.target 是 torch.ops.aten.add.Tensor 这一算子入口对象（OpOverload），直接整体替换 node.target 即完成算子改名，原有 args / kwargs 保持不变。

1.2 删除与追加节点：利用 FX 图工具函数

除了就地改写目标，我们还可以通过 FX 的 Graph 工具函数删除节点、追加新节点。例如在每次 add 之后插入一个 torch.ops.aten.relu.default：

import torch

def insert_relu_after_add(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
    for node in gm.graph.nodes:
        if node.op == "call_function" and node.target == torch.ops.aten.add.Tensor:

            # 指定插入点：在该作用域内添加到图中的节点
            # 会被插入到 `node` 之后
            with gm.graph.inserting_after(node):
                # 插入一个 op 为 `torch.ops.aten.relu.default` 的 call_function 节点
                new_relu_node = gm.graph.call_function(torch.ops.aten.relu.default, args=(node,))
                # 将所有使用 `node` 的地方替换为使用 `new_relu_node`
                node.replace_all_uses_with(new_relu_node)

这里两个关键 API 的语义：

• gm.graph.inserting_after(node) 上下文管理器设定后续节点创建的插入位置；
• node.replace_all_uses_with(new_node) 将图中所有以 node 为输入（user）的边重定向到 new_relu_node，实现"插入 + 重定向消费者"的一步到位。

更丰富的图工具函数（如 graph.erase_node、graph.call_method、graph.get_attr、graph.lint 等）可查阅 torch.fx.Graph 的完整文档。

1.3 变换的分类学：三个轴

虽然可以直接操纵图实现一切变换，但变换大体上可以按三个轴归类：

轴	维度	典型用例
轴 A	一对 X 映射（如分解 decomposition） vs 多对一映射（如融合 fusion）	A.1 分解；A.2 融合
轴 B	前向迭代（如形状传播 shape propagation） vs 反向迭代（如死代码消除 DCE）	B.1 前向；B.2 反向
轴 C	依赖局部节点信息（如 out-variant 转换） vs 依赖全局图信息（如内存规划 memory planning）	C.1 局部；C.2 全局

官方对使用频率的预判排序为：A.1（一对多）+ B.1（前向）+ C.1（局部） 最常见，其次是 A.2（多对一），再次是 B.2（反向）+ C.2（全局）。

虽然直接操纵图可以完成所有变换，但对于频率最高的 level 1 与 level 2 用例，PyTorch 额外提供了一些辅助工具，下文逐一展开。

---

2. Transformer：一对 X 映射与局部信息利用

对于 level 1 用例（一对 X 映射、前向迭代、仅依赖局部节点信息），可以借助 torch.fx.Transformer 类：它逐一执行每个节点并重建一张新图，在重建过程中应用你所指定的变换。

2.1 Transformer 的底层机制

从源码看，Transformer 是 Interpreter 的一个特殊子类：

• 它是一个符号化解释器，无需真实输入即可运行（"does not require arguments to run, as Interpreter does. Transformer works entirely symbolically"）；
• 它在构造时创建一张新图 new_graph，并内嵌一个 TransformerTracer；
• 覆写的 call_function 会调用 self.tracer.create_proxy("call_function", target, args, kwargs)，把算子调用"落"到新图上（torch/fx/interpreter.py#L641-L645）；
• 对外暴露 transform() 方法，返回变换后的新 GraphModule。

因此，Transformer 等价于"按原图拓扑解释执行一遍，同时把每个节点转录进一张新图，转录时应用你覆写的逻辑"。

2.2 One-to-One Pass（一对一映射）

如果要把算子 A 换成算子 B，只需在运行 GraphModule 时，每次遇到 A 就返回 B：

class ReplaceAddWithMul(torch.fx.Transformer):
    def call_function(self, target, args, kwargs):
        if target != torch.ops.aten.add.Tensor:
            return super().call_function(target, args, kwargs)
        return super().call_function(torch.ops.aten.mul.Tensor, args, kwargs)

transformed_graph_module = ReplaceAddWithMul(graph_module).transform()

super().call_function(target, args, kwargs, meta) 会创建一个 call_function FX 节点，并返回用给定参数运行该算子的结果——这个返回值会作为后续节点转录时的 Proxy 输入，从而把数据流自动接续到新图上。

2.3 One-to-X Pass（一对多映射）

如果要做一对多映射（例如把一个算子 A 拆成 B、C 两个算子），则调用两次 super().call_function 创建两个 FX 节点，并返回运行算子 C 的结果：

class ReplaceAddWithMulSub(torch.fx.Transformer):
    """
    Original:
        def f(x, y):
            return x + y

    After pass:
        def f(x, y):
            z = x * y
            return z - y
    """
    def call_function(self, target, args, kwargs):
        if target != torch.ops.aten.add.Tensor:
            return super().call_function(target, args, kwargs)

        x, y = args

        mul_res = super().call_function(torch.ops.aten.mul.Tensor, args, {})
        return super().call_function(torch.ops.aten.sub.Tensor, (mul_res, y), {})

transformed_graph_module = ReplaceAddWithMulSub(graph_module).transform()

这正是**分解（decomposition）**类变换的范式：一个高层算子被拆成若干底层原语，中间结果通过多次 super().call_function 连接。

2.4 One-to-None Pass（一对零映射，即删除算子）

如果想删除某个算子，只需直接返回传入该函数的输入值（相当于"短路"该节点，将其从新图中省略）：

class RemoveDetachPass(torch.fx.Transformer):
    def call_function(self, target, args, kwargs):
        if target not in (
            torch.ops.aten.detach.default,
            torch.ops.aten.detach_copy.default,
        ):
            return super().call_function(target, args, kwargs, meta)

        assert len(args) == 1
        return args[0]

transformed_graph_module = RemoveDetachPass(graph_module).transform()

原文档中的 RemoveDetachPass 用 super().call_function(target, args, kwargs, meta) 透传非目标节点（meta 为可选的额外参数），对 detach / detach_copy 则直接返回其唯一输入 args[0]，实现无副作用的去 detach 优化。

2.5 利用局部节点信息：借助算子 Schema 改写参数

Transformer 的第三个典型用法是读取当前节点的局部信息并据此改写参数。例如，把图中所有标量（scalar）参数转换为张量：遍历每个 call_function 节点，对其每个参数，如果它是 float / int / bool 且对应算子 Schema 中该参数的类型是 torch.TensorType，则用 torch.tensor(value) 进行转换。

def args_map(target, fn, args, kwargs):
    assert isinstance(args, tuple)
    assert isinstance(kwargs, dict)
    args = list(args)
    kwargs = kwargs.copy()

    # 根据传入的函数更新参数
    def update(key, args, schema):
        args[key] = fn(args[key], schema)

    # 逐个更新 Schema 中的参数
    for i, schema in enumerate(target._schema.arguments):
        if schema.name in kwargs:
            update(schema.name, kwargs, schema)
        elif not schema.kwarg_only and i < len(args):
            update(i, args, schema)
    return tuple(args), kwargs

class ScalarToTensorPass(torch.fx.Transformer):
    def call_function(self, target, args, kwargs):
        def try_coerce(value, arg):
            return (
                torch.tensor(value)
                if isinstance(value, (float, int, bool))
                and type(arg.type) == torch.TensorType
                else value
            )

        args, kwargs = args_map(target, try_coerce, args, kwargs)
        return super().call_function(target, args, kwargs)

transformed_graph_module = ScalarToTensorPass(graph_module).transform()

技术细节：

• target._schema.arguments 是 torch.ops.aten.* 算子注册时的完整参数 Schema，包含参数名、kwarg_only 标记与类型；
• args_map 辅助函数按 Schema 把"位置参数 / 关键字参数"统一映射到回调 fn 上；
• 类型判断依赖 arg.type == torch.TensorType（来自 torch 的公共类型系统），从而只对算子期望张量、实际却传入 Python 标量的参数做 torch.tensor 提升。

注意：原文档中的 ScalarToTensorPass 内含一行调试用的 breakpoint()，实际落地为 pass 时请删除，以免触发交互式调试器。

---

3. Subgraph Rewriter：多对一模式替换

对于 level 2 用例（多对一映射），可以利用 FX 的 subgraph rewriter（源码位于 torch/fx/subgraph_rewriter.py）。给定一个 pattern，它会找到与模式匹配的算子子图，并把每个匹配的子图替换为 replacement。

重要注意：这是一个就地（inplace）操作，会直接修改传入的图模块。

pattern 与 replacement 必须是可调用函数或 GraphModule，且内部使用与目标图相同的算子（ATen 算子），这样 subgraph rewriter 才能在图中找到对应模式。pattern/replacement 可调用对象的输入会被当作通配符参与匹配。

3.1 基本用法示例

from torch.fx import subgraph_rewriter

def replace_patterns(graph_module):
    def pattern(x, y):
        x = torch.ops.aten.add.Tensor(x, y)
        x = torch.ops.aten.mul.Tensor(x, y)
        return x

    def replacement(x, y):
        return torch.ops.aten.sub.Tensor(x, y)

    replaced_patterns = subgraph_rewriter.replace_pattern_with_filters(
        traced_module, pattern, replacement
    )

这段代码将图中所有 "add 之后再乘同一个 y" 的连续结构，整体替换为一个 sub 算子——即典型的算子融合雏形：多个算子坍缩为一个，同时保持语义等价。

3.2 返回值：ReplacedPatterns

subgraph rewriter 返回一个 ReplacedPatterns 列表，每个元素描述一次成功匹配与替换：

@dataclass
class ReplacedPatterns:
    # 发现匹配的起始节点
    anchor: Node
    # 模式子图中的节点到大图中的节点的映射
    nodes_map: Dict[Node, Node]
    # 被添加到图中的节点列表
    replacements: List[Node]

该数据结构与源码 torch/fx/subgraph_rewriter.py 中的 ReplacedPatterns 定义完全一致。

注意： subgraph rewriter 新建的节点不会携带被匹配节点上已有的元数据（metadata），但你可以借助：

• ReplacedPatterns.nodes_map 找到原图中被匹配的节点；
• ReplacedPatterns.replacements 找到变换后图中新替换的节点。

从而在需要传播 shape/dtype 等元数据时，自行建立新旧节点之间的对应关系。

---

4. Pass Manager：多 Pass 的调度与校验

PassManager（源码位于 torch/fx/passes/infra/pass_manager.py）用于在给定图模块上依次运行多个 pass。初始化 PassManager 时传入要执行的 pass 列表并设置若干标志；要运行整组 pass，直接把图模块作为参数调用 PassManager 实例即可。

4.1 构造与运行

from torch.fx.passes.infra.pass_manager import PassManager

pm = PassManager(
    passes=[replace_add_with_div, replace_div_with_mul],
    run_checks_after_each_pass=True,
    suppress_check_failures=False,
)
graph_module_out = pm(graph_module)

结合源码（torch/fx/passes/infra/pass_manager.py#L180-L195），PassManager.__init__ 支持的参数为：

参数	默认值	说明
passes	[]	pass 列表；每个 pass 是可调用对象，修改模块并返回 PassResult
constraints	[]	pass 间约束列表；约束是可调用对象，接收两个 pass (A, B)，返回 True 表示 A 依赖 B
steps	1	pass 组最多运行轮数；若某轮图不再变化则提前停止
run_checks_after_each_pass	False	是否在每个 pass 之后运行校验与 lint
suppress_check_failures	False	校验失败时是否抛异常

4.2 运行语义：PassResult 与图不再变化即停止

PassManager.__call__ 的执行流程（torch/fx/passes/infra/pass_manager.py#L254-L317）：

1. 若约束尚未验证，先调用 solve_constraints() 依据约束对 pass 做拓扑排序（存在环时抛 RuntimeError，见 _topological_sort_passes 与 this_before_that_pass_constraint）；
2. 运行当前 check 校验模块；
3. 在 steps 轮内逐 pass 执行：每个 pass 应返回 PassResult（否则抛 TypeError，提示用 pass_result_wrapper 包装），GraphModule 每轮后调用 recompile()，并在 run_checks_after_each_pass 开启时运行 check；
4. 若某轮没有任何 pass 报告修改（modified 均为 False），则提前终止循环。

4.3 添加公共校验：set_checks / add_checks

要为每个 pass 之后运行的检查集添加公共校验，调用 set_checks(check: Callable)，传入一个可调用函数。当 run_checks_after_each_pass 标志开启时，该 check 会在图模块上运行每个 pass 后被调用。

pm = PassManager(passes=[replace_add_with_div, replace_div_with_mul])

def check_div_target(graph_module):
    for node in graph_module.graph.nodes:
        if node.op == "call_function" and node.target != torch.div:
            raise ValueError("Target should be div!")

pm.add_checks(check_div_target)

pm(graph_module)    # 在 replace_div_with_mul pass 之后抛出 ValueError

源码中该 API 名为 add_checks（torch/fx/passes/infra/pass_manager.py#L236-L252），它还会校验 check 函数只接收一个参数（模块），否则抛 TypeError。在这个示例中，由于最后一个 pass 把所有节点替换为 torch.div（而非 torch.div 之外的算子，这里以 check_div_target 的判断为准），校验函数检测到不满足"target 应为 div"的节点即抛出 ValueError——这是用**图不变量（graph invariant）**守护 pass 流水线正确性的标准写法。

---

5. Partitioner：按能力切分与匹配子图

除了变换本身，FX 还提供若干常用的**图切分（partitioning）**工具。

5.1 SubgraphMatcher：模式子图查找

要在大图中查找与特定模式匹配的子图，可以使用 FX 的 SubgraphMatcher（源码位于 torch/fx/passes/utils/matcher_utils.py）。

类属性（构造参数）：

参数	默认值	说明
pattern (Graph)	—	目标匹配模式；图中的 placeholder 节点在匹配时被当作通配符
match_output (bool)	False	为 True 时模式图中的 output 节点作为模式的一部分；为 False 时匹配时忽略 output 节点
match_placeholder (bool)	False	为 True 时模式图中的 placeholder 节点作为模式的一部分；为 False 时 placeholder 用作通配符
remove_overlapping_matches (bool)	True	为 True 时，发生重叠匹配时只返回第一个匹配
ignore_literals (bool)	False	为 True 时不检查字面量是否相等，将其视为通配符

这些默认值在源码 torch/fx/passes/utils/matcher_utils.py#L62-L89 中有明确体现；构造时若模式图为空图或含死代码，会抛出 ValueError / AssertionError。

示例： 在大模型中查找与 PatternModel 相同结构的 addmm 子图：

from torch.fx.passes.utils.matcher_utils import SubgraphMatcher

class LargeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._weight = torch.nn.Parameter(torch.ones(3, 3))
        self._bias = torch.nn.Parameter(torch.ones(3, 3))

    def forward(self, x):
        return torch.ops.aten.addmm.default(self._bias, x, self._weight)

large_model_graph = torch.export(LargeModel(), inputs).graph

class PatternModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._weight_1 = torch.nn.Parameter(torch.ones(5, 5))
        self._bias_1 = torch.nn.Parameter(torch.ones(5, 5))

    def forward(self, x):
        return torch.ops.aten.addmm.default(self._bias_1, x, self._weight_1)

pattern_graph = torch.export(PatternModel(), inputs).graph

subgraph_matcher = SubgraphMatcher(pattern_graph)
match_result = subgraph_matcher.match(large_model_graph)

注意示例中的 inputs 需按实际模型输入自行构造（例如 (torch.randn(3, 3),)）。match 函数返回一个 InternalMatch 列表：

@dataclass
class InternalMatch():
    # 发现匹配的起始节点
    anchors: List[Node]
    # 模式子图中的节点到大图中的节点的映射
    nodes_map: Dict[Node, Node] = field(default_factory=dict)
    # 目标图中与模式 placeholder 匹配的节点
    placeholder_nodes: List[Node] = field(default_factory=list)
    # 匹配子图中被 output 返回的节点
    returning_nodes: List[Node] = field(default_factory=list)

该定义与源码 torch/fx/passes/utils/matcher_utils.py#L35-L59 一致（额外字段 name_node_map 仅在启用 SubgraphMatcherWithNameNodesMap 时可用）。从源码的匹配逻辑（_nodes_are_equal、_is_contained、_remove_overlapping_matches）可以推断：matcher 采用同构子图匹配——节点 op 与 target 逐项相等（placeholder 默认通配），且匹配子图不允许"泄漏"到模式之外（除 returning nodes 外的节点不得有图外使用者），同时支持重叠匹配去重。

5.2 CapabilityBasedPartitioner：最大可支持子图切分

要找到支持特定不变量（invariant）的最大节点子图，可以使用 FX 的 CapabilityBasedPartitioner（源码位于 torch/fx/passes/infra/partitioner.py）。

类属性（构造参数）：

参数	说明
graph_module (torch.fx.GraphModule)	待切分的图模块
operator_support (OperatorSupportBase)	用于判定图中某节点是否"支持"被纳入 partition 的对象
allows_single_node_partition (bool)	为 True 时允许形成单节点 partition
non_compute_ops (Optional[Sequence[str]])	被视为"非计算"算子的集合（如 torch.ops.aten.view、_operator.getitem），避免 partitioner 产出只含这些非计算算子的图
allowed_single_node_partition_ops (Optional[Sequence[str]])	允许处于单节点 partition 中的算子集合

源码中还额外提供了 skip_horizontal_fusion 参数（默认 False），用于禁止水平（同层）融合，可按需开启。

OperatorSupportBase 与组合：

OperatorSupportBase 由 partitioner 用来判定某个节点是否属于 partition，做法是覆写 is_node_supported 函数。此外可以组合多个 OperatorSupportBase：

• chain：任一个返回 False 则整体返回 False（逻辑与）；
• any_chain：任一个返回 True 则整体返回 True（逻辑或）。

二者实现分别基于 all(...) 与 any(...) 组合各子 OperatorSupportBase 的判定结果，便于把"dtype 约束""算子白名单""名称黑名单"等原子规则搭建成复合支持策略（同文件还提供了 OpSupports.decline_if_input_dtype、OpSupports.decline_if_node_in_names 等现成原子规则）。

示例： 找出所有 add/mul 节点组成的分区并融合为子模块：

from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
from torch.fx.passes.operator_support import any_chain, OperatorSupportBase

class AddMulOperatorSupport(OperatorSupportBase):
    def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
        return node.op == "call_function" and node.target in [
            torch.ops.aten.add.Tensor, torch.ops.aten.mul.Tensor,
        ]

capability_partitioner = CapabilityBasedPartitioner(
    graph_module,
    op_support,   # 传入 AddMulOperatorSupport() 实例
)

# 返回分区列表（每个分区包含属于该分区的节点）
partition_list = capability_partitioner.propose_partitions()
# 将分区融合为图模块，并在图中插入 call_module 节点
fused_graph_module = capability_partitioner.fuse_partitions(partition_list)

工作流程分两步：

1. propose_partitions() 依据 operator_support 计算出满足"所有节点均受支持"且尽量大的连续分区集合（从源码看，Partition 以 nodes 字典维护分区成员，_DependencyViewer 用于推导节点间下游依赖以约束分区连通性）；
2. fuse_partitions(partition_list) 借助 torch.fx.passes.utils.fuser_utils.fuse_by_partitions 把每个分区融合成一个子模块（call_module 节点），并将原子算子"下沉"进子模块的图中——这正是编译器中算子融合 + 子图下沉的典型落地路径。

---

6. 组合成一条完整编译流水线

综合全文，在 ATen IR 上构建一条编译优化流水线的标准范式可以归纳为：

1. 模式查找与匹配：用 SubgraphMatcher 在 torch.export 得到的 ATen 图上定位感兴趣的子图（如 addmm、卷积 + BN 等复合结构）；
2. 子图改写与融合：用 subgraph_rewriter 将匹配到的多算子结构替换为等价但更优的单个算子（多对一融合）；用 Transformer 实现一对一替换、一对多分解（detach 消除、标量提升等局部变换）；
3. 调度与校验：用 PassManager 把上述变换按依赖约束拓扑排序后顺序执行，并注册 check 函数在每个 pass 之后验证图不变量（如"所有 target 必须是 div"），run_checks_after_each_pass 保证错误第一时间暴露；
4. 能力切分与落子图：用 CapabilityBasedPartitioner + 自定义 OperatorSupportBase 找出满足后端执行能力的最大子图，fuse_partitions 生成可交给特定后端（如 Inductor、自定义 runtime）执行的子模块。

这套"查找 → 改写 → 调度校验 → 切分"的流程，与 PyTorch 编译栈中 torch.compile / torch.export 生态的实际架构保持一致：所有组件都构建在 torch.fx 的 Graph / GraphModule 抽象之上，因此对 ATen IR 与普通 FX 图一视同仁。读者可以参照 torch/fx/subgraph_rewriter.py、torch/fx/passes/infra/pass_manager.py、torch/fx/passes/utils/matcher_utils.py 与 torch/fx/passes/infra/partitioner.py 等源码继续深入，也可以在本仓库的 torch/_inductor/、torch/fx/passes/ 目录中寻找这些工具在生产级 pass 中的真实调用场景。