后端方言

概览

后端方言 是一种特殊的 边缘方言，因为它包含在后端特定的图转换之后的后端特定节点和元数据。后端方言是一个可选阶段，只有当我们希望将后端感知引入图中时才需要。更具体地说，后端方言中的图可能包含仅对目标后端有意义的操作符或委托降低模块（参见 委托文档）。一个用例是，如果我们想将操作符融合为单个操作符，例如将连续的 addmm + relu 融合为单个操作符 addmm_relu，我们可以在这一阶段完成。

本文档介绍如何引入后端特定的算子。

自定义算子与后端特定算子的区别：自定义算子会出现在 eager 模式、ATen 方言和 edge 方言中，而后端特定算子仅由 edge 方言之后执行的转换过程引入。

何时使用

该方言允许引入不符合标准 ATen 算子集中定义的 schema、且未出现在上述任何方言（ATen 方言和边缘方言）中的算子。如果您的使用场景满足以下一个或多个条件，请考虑使用后端算子：

• 您的后端提供了一个库，该库优化了等效于某个子图的特定操作符。例如，linear_relu（等效于线性 + relu），可以在某些后端上更快地执行。

• 在图模块已被降低到后端后，需要对其进行重新追踪。当我们进行重新追踪时，后端算子可以转换回原始子图（ATen 方言），而常规自定义算子并不处理这种情况。

• 您的后端特定算子没有通用的 CPU 内核，仅包含针对特定后端的核心实现。通过使用后端算子，可以将原始子图作为默认内核，从而解决此问题并保持图模块可运行。

• 或者，如果您担心这可能过于复杂，只需一个更轻量级的方案，且仅在编译器阶段需要 Python 代码，那么您可以使用 delegate。

APIs

对于算子/子图替换，常见流程如下：

1. 注册一个与子图具有相同输入和输出的算子。该算子无需提供特定于目标的实现（在编译阶段也不需要），但必须产生与子图相同的结果。

2. 创建一个模式，使编译器能够找到子图并用替换项进行替代。

3. 编写一个传递以用新算子替换子图。

为了简化该流程，我们提供了一套 API，以帮助 ExecuTorch 用户减少执行这些步骤所需的工作量。

通过基础设施入口

要将边缘算子降低为后端算子，一个转换过程将对图进行模式匹配以识别目标边缘算子，然后用等效的后端算子替换它们。有两种 API 可用于注册此类转换：

• transform(). ExportProgram 上的一个 API，允许用户提供自定义的转换。请注意，此功能没有任何验证器保护，因此程序的正确性无法得到保证。

• ExecutorchBackendConfig.passes。如果添加到这里，该传递将作为从后端方言到ExecutorchProgram的降低过程的一部分。

示例：其中一个这样的传递是QuantFusion。这个传递采用“规范量化模式”，即“反量化 - 某个操作 - 量化”，并将此模式融合为一个特定于后端的单个操作，例如quantized_decomposed::some_op。另一个更简单的示例是这里，我们将sym_size个操作替换为ExecuTorch能够理解的操作。

模式绑定装饰器

我们提供了一个装饰器 bind_pattern_to_op，以帮助用户轻松地将后端操作符注册到EXIR中。该装饰器接受以下内容：

• 一个 torch.Library 对象，它表示这个后端操作符属于哪个库或命名空间。

• 一个名称或模式。如果我们已经在 torch.Library 对象中定义了后端操作符的模式，只需提供一个名称即可。否则，如果传入了一个模式字符串，我们可以注册该模式。

这个装饰器应该被添加到我们试图匹配的模式上（然后将其降低为边缘方言上的此后端操作）。这样，我们就将此模式注册为此后端操作的一个 CompositeImplicitAutograd 内核。

然后，操作符可以从传递中访问/使用。CompositeImplicitAutograd 内核确保：

1. 用户无需编写可在 CPU 上运行的内核。

2. 确保ExportProgram的可追溯性。一旦被追溯，后端操作符将被分解为模式中使用的ATen操作。

示例

让我们假设一个包含 add 和 relu 算子的简单程序：

def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = x + y
    return torch.ops.aten.relu.default(z)

在降低为边缘方言后，它变为：

graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%arg0_1, %arg1_1), kwargs = {})
    %aten_relu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.relu.default](args = (%aten_add_tensor,), kwargs = {})
    return (aten_relu_default,)

现在我想编写一个传递函数，将 add 和 relu 合并为 add_relu，第一步是编写一个模式：

# In the pattern, we can use edge ops and ATen ops interchangably
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = torch.ops.aten.add.Tensor(x, y)
    out = torch.ops.aten.relu.default(z)
    return out

然后，我们需要从融合算子命名空间创建算子库，并在我们的模式上使用装饰器：

lib = Library("foo_namespace", "DEF")

@bind_pattern_to_op(lib, "add_relu(Tensor self, Tensor other) -> Tensor")
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        z = torch.ops.aten.add.Tensor(x, y)
        out = torch.ops.aten.relu.default(z)
        return out

这样，我们就将该模式注册为一个内核到 add_relu，并且它已经准备好在传递中使用。一个简单的传递看起来像这样：

class AddReluFusionPass(ExportPass):
    def call(self, graph_module: GraphModule) -> PassResult:
        # decorator registers this pattern as a CompositeExplicitAutograd kernel, since there's no kernel registered before.
        @bind_pattern_to_op(lib, "add_relu")
        def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
            z = torch.ops.aten.add.Tensor(x, y)
            out = torch.ops.aten.relu.default(z)
            return out

        def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
            return torch.ops.foo_namespace.add_relu.default(x, y)

        subgraph_rewriter.replace_pattern(
            graph_module,
            _trace_and_lower_to_edge_ops(pattern),
            _trace_and_lower_to_edge_ops(replacement),
        )
        return PassResult(graph_module, True)

结果图如下所示：

graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %foo_namespace_add_relu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.foo_namespace.add_relu.default](args = (%arg0_1, %arg1_1), kwargs = {})
    return (foo_namespace_add_relu_default,)

操作集

目前使用 bind_pattern_to_op API 的后端操作符有以下这些。

• executorch_prims::add.int(SymInt a, SymInt b) -> SymInt

  • pattern: builtin.add

  • 后端：执行器

• executorch_prims::mul.int(SymInt a, SymInt b) -> SymInt

  • pattern: builtin.mul

  • 后端：执行器

• executorch_prims::sub.int(SymInt a, SymInt b) -> SymInt

  • pattern: builtin.sub

  • 后端：执行器

• executorch_prims::floordiv.int(SymInt a, SymInt b) -> SymInt

  • pattern: builtin.floordiv

  • 后端：执行器

• executorch_prims::truediv.int(Scalar a, Scalar b) -> Scalar

  • pattern: builtin.div

  • 后端：执行器

• executorch_prims::sym_float.Scalar(Scalar a) -> Scalar

  • pattern: builtin.float

  • 后端：执行器

• executorch_prims::gt.int(SymInt a, SymInt b) -> bool

  • pattern: builtin.gt

  • 后端：执行器

• executorch_prims::lt.int(SymInt a, SymInt b) -> bool

  • pattern: builtin.lt

  • 后端：执行器

• executorch_prims::ge.int(SymInt a, SymInt b) -> bool

  • pattern: builtin.ge

  • 后端：执行器

• executorch_prims::le.int(SymInt a, SymInt b) -> bool

  • pattern: builtin.le

  • 后端：执行器

• executorch_prims::eq.int(SymInt a, SymInt b) -> bool

  • pattern: builtin.eq

  • 后端：执行器

• executorch_prims::mod.Scalar(SymInt a, SymInt b) -> SymInt

  • pattern: builtin.divmod

  • 后端：执行器

• executorch_prims::neg.Scalar(Scalar a) -> Scalar

  • pattern: operator.ne

  • 后端：执行器

• quantized_decomposed::embedding_byte(Tensor weight, Tensor weight_scales, Tensor weight_zero_points, int weight_quant_min, int weight_quant_max, Tensor indices) -> Tensor

  • 模式: 源

  • 后端：量化

• quantized_decomposed::add(Tensor a, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, Tensor b, float b_scale, int b_zero_point, int b_quant_min, int b_quant_max, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max) -> Tensor qc

  • 模式: 源

  • 后端：量化

• quantized_decomposed::add.scalar(Tensor qa, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, ScalarType a_dtype, Scalar b, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max, ScalarType out_dtype) -> Tensor

  • 模式: 源

  • 后端：量化

• quantized_decomposed::add_relu(Tensor a, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, Tensor b, float b_scale, int b_zero_point, int b_quant_min, int b_quant_max, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max) -> Tensor qc

  • 模式: 源

  • 后端：量化