后端方言

概述

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

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

自定义运算符和后端特定运算符之间的区别：虽然自定义运算符出现在急切模式、ATen 方言和边缘方言中，但后端特定运算符仅由边缘方言之后发生的流程引入。

何时使用

此方言允许引入不符合规范 ATen 运算符集中定义的模式，并且不出现在上述任何方言（ATen 方言和边缘方言）中的运算符。如果您的用例满足以下一个或多个条件，请考虑使用后端运算符

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

• 在模块已降低到后端之后需要重新跟踪图模块。当我们重新跟踪时，后端运算符可以转换回原始子图（在 ATen 方言中），而普通的自定义运算符不会处理这种情况。

• 您的后端特定运算符没有通用的 CPU 内核，但只有特定后端的内核。通过使用原始子图作为默认内核并保持图模块可运行，使用后端运算符可以解决此问题。

• 或者，如果您认为这可能有点过头，并且只需要更轻量级的东西，并且在编译阶段只需要 Python 代码，则可以使用委托。

API

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

1. 注册一个与子图具有相同输入和输出的运算符。此运算符将不具有目标特定的实现（同样，在编译阶段也不需要），但它需要给出与子图相同的结果。

2. 创建一个允许编译器查找子图并将其替换为替换的模式。

3. 编写一个流程来替换子图和新的运算符。

为了促进此过程，我们提供了一个 API 来帮助减少 ExecuTorch 用户执行这些步骤的工作量。

流程基础设施入口点

为了将边缘运算符降低到后端运算符，流程将执行模式匹配以识别图中感兴趣的边缘运算符，然后将它们替换为等效的后端运算符。有两个 API 可以注册此类流程

• transform()。ExportProgram 上的 API，允许用户提供自定义流程。请注意，此 API 没有受到任何验证器的保护，因此程序的健全性无法保证。

• ExecutorchBackendConfig.passes。如果在此处添加，则流程将成为从后端方言到 ExecutorchProgram 的降低过程的一部分。

示例：其中一个流程是 QuantFusion。此流程采用“规范量化模式”，即“dequant - some_op - quant”，并将此模式融合成一个后端特定的单个运算符，即 quantized_decomposed::some_op。另一个更简单的示例是 此处，我们用 ExecuTorch 理解的运算符替换 sym_size 运算符

模式绑定装饰器

我们提供了一个装饰器 bind_pattern_to_op 来帮助用户轻松地将他们的后端运算符注册到 EXIR 中。此装饰器采用

• 一个 torch.Library 对象，它指示此后端运算符属于哪个库或命名空间。

• 一个名称或模式。如果我们已在 torch.Library 对象中定义了后端运算符的模式，则只需要名称。否则，如果传递了模式字符串，我们可以注册模式。

此装饰器应添加到我们尝试匹配的模式（然后降低到此后端运算符）上边缘方言。这样，我们就将此模式注册为该后端运算符的 CompositeImplicitAutograd 内核。

然后，可以在 Pass 中访问/使用该算子。 CompositeImplicitAutograd 内核确保

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

2. 确保 ExportProgram 的可回溯性。回溯后，后端算子将被分解为模式中使用的 ATen 算子。

示例

假设一个简单的程序，其中包含加法和 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,)

现在我想编写一个 Pass 将 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 的内核，并且它已准备好用于 Pass 中。一个简单的 Pass 看起来像这样

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

  • 模式：builtin.add

  • 后端：executor

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

  • 模式：builtin.mul

  • 后端：executor

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

  • 模式：builtin.sub

  • 后端：executor

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

  • 模式：builtin.floordiv

  • 后端：executor

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

  • 模式：builtin.div

  • 后端：executor

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

  • 模式：builtin.float

  • 后端：executor

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

  • 模式：builtin.gt

  • 后端：executor

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

  • 模式：builtin.lt

  • 后端：executor

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

  • 模式：builtin.ge

  • 后端：executor

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

  • 模式：builtin.le

  • 后端：executor

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

  • 模式：builtin.eq

  • 后端：executor

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

  • 模式：builtin.divmod

  • 后端：executor

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

  • 模式：operator.ne

  • 后端：executor

• 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

  • 模式：来源

  • 后端：量化