自定义编译器传递和分区器

传递

Passes 大致可以按几个维度进行分类：

轴 A：

1. 创建一对多映射（例如分解）

2. 创建多对一映射（例如融合）

轴 B：

1. 执行前向迭代（例如，形状传播）

2. 执行反向迭代（例如，死代码消除）

轴 C：

1. 依赖于本地节点信息（例如，输出变体转换）

2. 依赖于全局图信息（例如内存规划）

我们对这些用例的使用频率预测如下：

1. A.1, B.1, C.1

2. A.2

3. B.2, C.2

级别 1

对于级别 1 的用例（创建一对多映射、执行前向迭代以及查看本地节点信息），我们可以利用一个名为 ExportPass的辅助类。 这是一种基于 解释器 的方法，我们在此方法中执行每个节点并重新生成图，但应用了指定的转换。这使我们能够通过确保在传递过程中创建的所有节点都符合 IR 规范来保留 IR 规范，包括确保元数据（如堆栈跟踪、FakeTensor 值和 torch.nn.Module 层次结构）根据所做的转换得到保留和更新。

要实现此传递，我们可以创建ExportPass的子类并实现暴露的函数。当使用图模块调用时，它将运行该图模块并创建一个包含由该传递指定的更改的新图。这意味着传入的图模块必须在 CPU 上可运行，并且在此传递运行后，这一不变性将得以保持。

一对一传递

一对一映射的一个示例，如果我们想要用另一个操作 B 替换操作 A， 我们可以运行给定的 fx.GraphModule，每次遇到操作 A 时，返回操作 B。

考虑以下示例：

class ReplaceInPlaceReluWithOutOfPlaceReluPass(ExportPass):
    """
    relu_ is the in-place version. Replace it with relu, which is the
    out-of-place version
    """

    def call_operator(self, op, args, kwargs, meta):
        if op != torch.ops.aten.relu_.default:
            return super().call_operator(op, args, kwargs, meta)
        return super().call_operator(Op(torch.ops.aten.relu.default), args, kwargs, meta)

# To create a pass
replace_pass = ReplaceInPlaceReluWithOutOfPlaceReluPass()
# To run a pass
new_graph_module = replace_pass(graph_module).graph_module

super().call_operator(op, args, kwargs, meta) 调用会创建一个 call_function FX 节点，并返回使用给定参数运行该操作符的结果。

一对多传递

如果我们想要执行一对多映射，例如将操作 A 替换为另外两个操作 B 和 C，那么我们需要调用两次 super().call_operator 来创建两个 FX 节点，一个包含操作 B，另一个包含操作 C，并返回运行操作 C 的结果。

例如：

class ReplaceAddWithMulSub(ExportPass):
    """
    Original:
        def f(x, y):
            return x + y

    After pass:
        def f(x, y):
            z = x * y
            return z - y
    """
    def call_operator(self, op, args, kwargs, meta):
        if op != torch.ops.aten.add.default:
            return super().call_operator(op, args, kwargs, meta)

        x, y = args

        mul_res = super().call_operator(
            torch.ops.aten.mul.default,
            args,
            {},
            meta
        )

        return super().call_operator(
            torch.ops.aten.sub.default,
            (mul_res, y),
            {},
            meta
        )

一对零传递

如果我们想要移除一个算子，只需返回传入函数的值即可：

class RemoveDetachPass(ExportPass):
    def call_operator(self, op, args, kwargs, meta):
        if op not in (
            torch.ops.aten.detach.default,
            torch.ops.aten.detach_copy.default,
        ):
            return super().call_operator(op, args, kwargs, meta)

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

利用本地信息

利用本地节点信息的一个示例是，如果我们希望将图中的所有标量转换为张量，我们可以运行给定的fx.GraphModule，并将每个包含标量的参数转换为张量。它可能看起来像这样：

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

    # Update the argument based on the function passed
    def update(key, args, schema):
        args[key] = fn(args[key], schema)

    # Update each argument in the schema
    for i, schema in enumerate(self.op._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)

class ScalarToTensorPass(ExportPass):
    def call_operator(self, op, 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(op, try_coerce, args, kwargs)
        return super().call_operator(op, args, kwargs)

级别 2

对于创建多对一映射，我们可以利用FX的子图重写器。 给定一个pattern，它会创建与模式匹配的操作符子图， 然后将每个匹配的子图替换为replacement。

注意

This is an inplace operation.

pattern 和 replacement 输入必须是使用与 EXIR 图匹配时所使用的相同操作（ATen 操作）编写的可调用函数，以便子图重写器能够在图中找到正确的模式。模式/替换可调用函数的输入将被视为通配符。

考虑以下示例：

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
)

子图重写器返回一个 ReplacedPatterns 列表：

@dataclass
class ReplacedPatterns:
    # Node from which the match was found
    anchor: Node
    # Maps nodes in the pattern subgraph to nodes in the larger graph
    nodes_map: Dict[Node, Node]
    # List of nodes that were added into the graph
    replacements: List[Node]

注意

The nodes created by the subgraph rewriter will not have the metadata that
is normally in EXIR nodes (`stack_trace`, `val`, `nn_module_stack`).

第 3 级

对于创建传递的第三种方式，我们可以利用最基础的 PassBase。 要创建一个传递，我们可以继承此类并实现函数 call 以包含传递内容。此外，我们还可以实现函数 requires 和 ensures，它们将在函数 call 之前和之后被调用。请注意， 这些函数也可以在 ExportPass 中被重写。要在图模块上运行一个传递，我们可以直接将图模块传递给该类的实例。

考虑以下示例：

class ReplaceAddPass(PassBase):

    def __init__(self, replace_op):
        self.replace_op = replace_op

    def call(self, graph_module):
        for node in gm.graph.nodes:
            if node.op == "call_function" and node.target == torch.add:
                node.target = self.replace_op

    # Optional to implement, will be called before call()
    def requires(self, graph_module) -> None:
        for node in graph_module.graph.nodes:
            if node.op == "call_function" and node.target == torch.add:
                return
        raise ValueError("No torch.add ops!")

    # Optional to implement, will be called after call()
    def ensures(self, graph_module: torch.fx.GraphModule) -> None:
        pass

# To create a pass
replace_add_with_div = ReplaceAddPass(torch.div)
# To run a pass
replace_add_with_div(graph_module)

传递管理器

PassManager 是一个用于在给定图模块上运行多次传递的类。在初始化 PassManager 实例时，我们传入想要运行的传递列表并设置几个标志。要在图模块上运行这些传递集合，我们可以直接将图模块传递给 PassManager 实例。

一个示例：

from executorch.exir.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)

要添加一组在每个传递后运行的通用检查，我们可以调用函数 set_checks(check: Callable)，该函数接受一个可调用函数作为输入。如果设置了 run_checks_after_each_pass 标志，则 check 将在图模块的每次传递运行后被调用。

一个示例：

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)    # raises ValueError after replace_div_with_mul pass

分区器

我们可以使用几种基于 FX 图的分区器来对图进行分区。然而，这些分区器不一定能生成符合 IR 规范的图，因此在使用时请谨慎。

子图匹配器

要在图中查找与特定模式匹配的子图，我们可以利用 FX 的 SubgraphMatcher。

类属性：

• pattern (Graph): 目标匹配模式。图中的占位符节点在匹配时将被视为通配符。

• match_output (bool): 如果为 True，模式图中的输出节点将被视为目标模式的一部分。如果为 False，则在匹配过程中忽略输出节点。

• match_placeholder (bool): 如果为 True，模式图中的占位符节点将被视为目标模式的一部分。如果为 False，占位符节点将用作通配符。

• remove_overlapping_matches (bool): 如果为 True，在匹配重叠的情况下，将仅返回第一个匹配项。

• ignore_literals (bool): 如果为 True，将不检查字面量是否相等，而是将其视为通配符。

考虑以下示例：

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 = to_edge(export(LargeModel(), large_inputs)).exported_program().graph_module.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 = to_edge(export(PatternModel(), pattern_inputs)).exported_program().graph_module.graph

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

match 函数返回一个 InternalMatch 列表：

@dataclass
class InternalMatch():
    # Nodes from which the match was found
    anchors: List[Node]
    # Maps nodes in the pattern subgraph to nodes in the larger graph
    nodes_map: Dict[Node, Node] = field(default_factory=dict)
    # Nodes in target graph that are matched placeholder in pattern
    placeholder_nodes: List[Node] = field(default_factory=list)
    # Nodes in matched subgraph returned by output
    returning_nodes: List[Node] = field(default_factory=list)

基于能力的分区器

要查找支持特定不变性的最大节点子图，我们可以利用 FX 的 CapabilityBasedPartitioner。

类属性

• graph_module (torch.fx.GraphModule): 我们要进行划分的图模块。

• operator_support (OperatorSupportBase): 用于确定图中的节点是否受分区支持的对象。

• allows_single_node_partition (bool): 如果为 True，允许形成单节点分区。

• non_compute_ops (Optional[Sequence[str]]): 一组被视为“非计算”的操作（例如 torch.ops.aten.view 和 _operator.getitem），以便分区器不会创建仅包含这些非计算操作的图

• allowed_single_node_partition_ops (Optional[Sequence[str]]): 允许包含在单个节点分区中的一组操作。

这个 OperatorSupportBase 类被分区器用来确定图中的特定节点是否属于该分区。这是通过覆盖is_node_supported函数来实现的。你可以通过使用 chain（如果任何一个OperatorSupportBase返回False，则返回False）和 any_chain （如果任何一个OperatorSupportBase返回True，则返回True）来链式调用多个OperatorSuppportBase。

考虑以下示例：

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,
)

# Returns a list of partitions (list of nodes that belong in each partition)
partition_list = capability_partitioner.propose_partitions()

如果您查看基于能力的分区器，您可能会发现一个 fuse_partition 函数，该函数将返回一个修改后的图，其中分区作为子模块，并通过 call_module 节点在顶层图中调用这些子模块。然而，这不符合 IR 规范，因为我们要不允许 call_module 节点。

组合

我们还提供了一个组合辅助函数： generate_pattern_op_partitions

Args:

• graph_module (fx.GraphModule): 我们要进行划分的模块

• patterns (List[torch.fx.Graph]): 一系列以 torch.fx.Graph 形式表示的模式。这些图可以通过 exir.capture（推荐）或符号追踪（可能无法生成准确的边缘方言图）获取的 GraphModule 中的 graph 字段获得，也可以通过手动构建图模块来创建。

• op_support (OperatorSupportBase)：一种 OperatorSupportBase，可通过以下方式创建：

  • 直接子类化并实现 is_node_supported()

  • 获取 create_op_support() 的结果

  • 获取 create_pattern_support() 的结果

  • 多个 OperatorSupportBase 类通过 chain() 或 any_chain() 级联在一起

返回

• 由给定的 OperatorSupportBase 对象与给定模式图的并集所支持的是包含节点的分区（即最大可能子图）列表。

源分区器

对于更复杂的用例，用户可能希望根据更高层次的模块（torch.nn.Linear 或 torch.nn.functional.Linear）进行分区，这些模块现在被分解为它们的操作符（aten.permute, aten.addmm），我们有以下辅助函数：

get_source_partitions(graph: torch.fx.Graph, wanted_sources: List[Any]) -> Dict[Any, SourcePartition]

Args:

• graph: 我们要进行划分的图

• wanted_sources: 从此源分解出的节点来源列表。这可以是一个函数（例如 torch.nn.functional.linear）或一个叶模块类型（例如 torch.nn.Linear）

Returns:

• 将源（例如 torch.nn.modules.linear.Linear）映射到对应于从该类型模块展平的节点列表的 SourcePartitions 列表。

@dataclass
class SourcePartition():
    # Nodes in a particular partition
    nodes: List[Node]
    # Module type
    module_type: Type
    # Nodes in the graph that are needed as inputs to the partition
    input_nodes: List[Node] = field(default_factory=list)
    # Nodes in the partition that are being used by nodes outside of the partition
    output_nodes: List[Node] = field(default_factory=list)
    # Parameters that are being used
    params: List[str] = field(default_factory=list)

一个示例：

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(3, 3)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(3, 5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

inputs = (torch.randn(3, 3),)
edge_graph = to_edge(export(M(), inputs)).exported_program().graph_module.graph
print(edge_graph)
"""
graph():
    %arg0 : [#users=1] = placeholder[target=arg0]
    %_param_constant0 : [#users=1] = get_attr[target=_param_constant0]
    %permute_default : [#users=1] = call_function[target=torch.ops.aten.permute_copy.default](args = (%_param_constant0,), kwargs = {})
    %_param_constant1 : [#users=1] = get_attr[target=_param_constant1]
    %addmm_default : [#users=1] = call_function[target=torch.ops.aten.addmm.default](args = (%_param_constant1, %arg0, %t_default), kwargs = {})
    %_param_constant0_1 : [#users=1] = get_attr[target=_param_constant0]
    %permute_default_1 : [#users=1] = call_function[target=torch.ops.aten.permute_copy.default](args = (%_param_constant0_1,), kwargs = {})
    %_param_constant1_1 : [#users=1] = get_attr[target=_param_constant1]
    %addmm_default_1 : [#users=1] = call_function[target=torch.ops.aten.addmm.default](args = (%_param_constant1_1, %addmm_default, %t_default_1), kwargs = {})
    %relu_default : [#users=1] = call_function[target=torch.ops.aten.relu.default](args = (%addmm_default_1,), kwargs = {})
    %_param_constant2 : [#users=1] = get_attr[target=_param_constant2]
    %permute_default_2 : [#users=1] = call_function[target=torch.ops.aten.permute_copy.default](args = (%_param_constant2,), kwargs = {})
    %_param_constant3 : [#users=1] = get_attr[target=_param_constant3]
    %addmm_default_2 : [#users=1] = call_function[target=torch.ops.aten.addmm.default](args = (%_param_constant3, %relu_default, %t_default_2), kwargs = {})
    return [addmm_default_2]
"""

module_partitions = get_source_partitions(edge_graph, [torch.nn.Linear, torch.nn.ReLU])
print(module_partitions)
"""
{<class 'torch.nn.modules.linear.Linear'>: [
    ModulePartition(nodes=[_param_constant0, t_default, _param_constant1, addmm_default], module_type=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[arg0], output_nodes=[addmm_default], params=["_param_constant0", "_param_constant1"]),
    ModulePartition(nodes=[_param_constant0_1, t_default_1, _param_constant1_1, addmm_default_1], module_type=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[addmm_default], output_nodes=[addmm_default_1], params=["_param_constant0_1", "_param_constant1_1"]),
    ModulePartition(nodes=[_param_constant2, t_default_2, _param_constant3, addmm_default_2], module_type=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[relu_default], output_nodes=[addmm_default_2], params=["_param_constant2", "_param_constant3"])],

 <class 'torch.nn.modules.activation.ReLU'>: [
    ModulePartition(nodes=[relu_default], module_type=<class 'torch.nn.modules.activation.ReLU'>, input_nodes=[addmm_default_1], output_nodes=[relu_default], params=[])]}
"""