后端方言¶
概述¶
Backend dialect 是 edge dialect 的一种特殊变体,因为它在后端特定的图形转换之后包含后端特定的节点和元数据。Backend dialect 是一个可选阶段,仅当我们想将 backend-awareness 引入图形时才需要。更具体地说,后端方言中的图形可能包含仅对目标后端有意义的运算符或委托的降低模块(参见 delegate doc)。一个用例是,如果我们想将运算符融合为单个运算符,例如,将连续的 addmm + relu 融合到单个运算符 addmm_relu,我们可以在这里执行此操作。
本文将介绍如何引入 backend 特定的 Operator。
自定义操作和后端特定操作之间的区别:自定义操作以 Eager 模式、ATen 方言和边缘方言显示,而后端特定操作仅通过边缘方言之后发生的通道引入。
适用情形¶
此方言允许引入不符合规范 ATen 运算符集中定义的架构的运算符,并且这些运算符不会出现在上述任何方言(ATen 方言和边缘方言)中。如果您的使用案例满足以下一个或多个条件,请考虑使用后端运算符:
您的后端提供了一个库,该库优化了相当于子图的特定运算符。例如,(相当于 linear + relu)可以在某个后端更快地执行。
linear_relu
在 graph 模块已经降低到 backend 后,需要对其进行 retrace。当我们回溯时,后端运算符可以转换回原始子图(在 ATen 方言中),而普通的自定义运算不会处理这个问题。
您的后端特定 Operator 没有通用的 CPU 内核,而只有某个后端的内核。使用后端运算符可以通过使用原始子图作为默认内核并保持图形模块可运行来解决此问题。
或者,如果您担心它可能有点矫枉过正,并且只想要更轻量级的东西,并且只需要编译器阶段的 Python 代码,则可以使用 delegate。
蜜蜂属¶
对于运算符/子图替换,常见流程为:
注册一个与子图具有相同输入和输出的运算符。这个 operator 不会有特定于 target 的实现(而且,在编译阶段不需要),但它需要给出与 subgraph 相同的结果。
创建一个模式,允许编译器查找 subgraph 并将其替换为替换。
编写一个传递以将 subgraph 替换为 new 运算符。
为了简化该过程,我们提供了一个 API 来帮助减少 ExecuTorch 用户执行这些步骤的工作量。
通过 Infra 入口点¶
为了将边缘运算降低到后端运算,传递将执行模式匹配以识别图中感兴趣的边缘运算,然后将它们替换为等效的后端运算。有两个 API 可用于注册此类通行证:
transform()
.ExportProgram 上的一个 API,允许用户提供自定义通行证。请注意,这不受任何验证者保护,因此无法保证程序的健全性。如果在此处添加,则该通道将成为从后端方言到 ExecutorchProgram 的降低过程的一部分。
示例:QuantFusion 就是这样一个通道。此通道采用 “规范量化模式”,即。“dequant - some_op - quant”,并将此模式融合为特定于后端的单个运算符,即 .另一个更简单的例子是这里,我们将运算符替换为 ExecuTorch 可以理解的运算符quantized_decomposed::some_op
sym_size
Pattern Binding Decorator¶
我们提供了一个 decorator 来帮助用户轻松地将他们的后端操作者注册到 EXIR 中。此装饰器采用:bind_pattern_to_op
一个 Object 中,它指示此 backend Operator 属于哪个库或命名空间。
torch.Library
名称或架构。如果我们已经在对象中定义了 backend 运算符的架构,则只需要一个名称。否则,如果传入了 schema 字符串,我们可以注册 schema。
torch.Library
这个装饰器应该被添加到我们试图在 edge dialect 上匹配的模式中(然后降低到这个后端操作)。这样,我们将此模式注册为此后端 Operator 的内核。CompositeImplicitAutograd
然后可以从通道访问/使用运算符。内核确保:CompositeImplicitAutograd
用户无需编写 (CPU) 可运行内核。
确保 的可回溯性。回溯后,backend 运算符将被分解为 Code Pattern 中使用的 ATen 运算。
ExportProgram
例¶
假设一个简单的程序同时包含 add 和 relu 运算符:
def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
z = x + y
return torch.ops.aten.relu.default(z)
降低到 edge dialect 后,它变为:
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 来 merge 和 into ,第一步是写一个 pattern: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
然后我们需要从 fused operator 命名空间创建一个 operator 库,然后在我们的模式上使用装饰器:
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
通过这种方式,我们将模式注册为内核,并且它已准备好在 pass 中使用。一个简单的传递如下所示: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,)
操作集¶
目前有使用 API 的后端 Operator。bind_pattern_to_op
executorch_prims::add.int(SymInt a, SymInt b) -> SymInt
模式:builtin.add
backend:执行程序
executorch_prims::mul.int(SymInt a, SymInt b) -> SymInt
模式:builtin.mul
backend:执行程序
executorch_prims::sub.int(SymInt a, SymInt b) -> SymInt
模式:builtin.sub
backend:执行程序
executorch_prims::floordiv.int(SymInt a, SymInt b) -> SymInt
模式:builtin.floordiv
backend:执行程序
executorch_prims::truediv.int(Scalar a, Scalar b) -> Scalar
模式:builtin.div
backend:执行程序
executorch_prims::sym_float.Scalar(Scalar a) -> Scalar
模式:builtin.float
backend:执行程序
executorch_prims::gt.int(SymInt a, SymInt b) -> bool
图案:builtin.gt
backend:执行程序
executorch_prims::lt.int(SymInt a, SymInt b) -> bool
图案:builtin.lt
backend:执行程序
executorch_prims::ge.int(SymInt a, SymInt b) -> bool
图案:builtin.ge
backend:执行程序
executorch_prims::le.int(SymInt a, SymInt b) -> bool
模式:builtin.le
backend:执行程序
executorch_prims::eq.int(SymInt a, SymInt b) -> bool
模式:builtin.eq
backend:执行程序
executorch_prims::mod.Scalar(SymInt a, SymInt b) -> SymInt
模式:builtin.divmod
backend:执行程序
executorch_prims::neg.Scalar(Scalar a) -> Scalar
图案:operator.ne
backend:执行程序
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
模式:源
后端:量化