导出 IR 规范¶
导出 IR 是 torch.export
结果的中间表示 (IR)。要了解导出 IR 的更多详细信息,请阅读此 文档。
导出的 IR 是一个规范,由以下部分组成:
计算图模型的定义。
图中允许的操作符集合。
方言 (dialect) 是由下面定义的这些操作组成的一个导出 IR 图,但具有额外的属性(例如,对操作符集或元数据的限制),这些属性是为了特定目的而设计的。
目前存在的 EXIR 方言是:
这些方言代表了捕获程序从程序捕获到转换为可执行格式所经历的阶段。例如,ExecuTorch 编译过程从 Python 程序捕获到 ATen 方言开始,然后 ATen 方言转换为 Edge 方言,Edge 转换为 Backend,最后转换为用于执行的二进制格式。
ATen 方言¶
ATen 方言将作为 ExecuTorch 编译管道的入口点。这是命令式 PyTorch 程序首次成为导出 IR 图。在此阶段,会执行函数化,移除任何张量别名和变异,并允许进行更灵活的图转换。此外,所有张量都会转换为连续格式。
此方言的目标是尽可能忠实地捕获用户的程序(同时保持有效的导出 IR)。用户在命令式模式下调用的已注册自定义操作符将在 ATen 方言中保持不变。但是,我们应该避免通过 pass 在图中添加自定义操作符。
目前,ATen 方言的功能是进一步下推到 Edge 方言。然而,未来我们可以将它视为其他导出用例的通用集成点。
ATen 方言的属性¶
ATen 方言图是一个有效的导出 IR 图,并具有以下附加属性:
在
call_function
节点中的所有操作符要么是 ATen 操作符(在torch.ops.aten
命名空间下)、高阶操作符(如控制流操作符),要么是已注册的自定义操作符。已注册的自定义操作符是注册到当前 PyTorch 命令式运行时中的操作符,通常使用TORCH_LIBRARY
调用(意味着有 schema)。有关如何注册自定义操作符的详细信息,请在此 处 查找。每个操作符还必须有一个元内核 (meta kernel)。元内核是一个函数,给定输入张量的形状,可以返回输出张量的形状。有关如何编写元内核的详细信息,请在此 处 查找。
输入值类型必须是“Pytree-able”(可 Pytree 化)。因此,输出类型也必须是 Pytree-able,因为所有操作符的输出都是 Pytree-able 的。
ATen 方言的操作符可以选择使用动态数据类型、隐式类型提升和张量的隐式广播。
所有张量的内存格式都为
torch.contiguous_format
。
Edge 方言¶
此方言旨在引入对 Edge 设备有用的专用化,但不一定适用于通用(服务器)导出。但是,我们仍然会避免进一步专用于每种不同的硬件。换句话说,我们不想引入任何新的与硬件相关的概念或数据;除了用户原始 Python 程序中已有的。
Edge 方言的属性¶
Edge 方言图是一个有效的导出 IR 图,并具有以下附加属性:
OpCall 节点中的所有操作符要么来自预定义的运算符集,称为“Edge Operators”,要么是已注册的自定义运算符。Edge 操作符是具有数据类型专用的 ATen 操作符。这允许用户注册仅适用于某些数据类型的内核,以减小二进制文件大小。
图的输入和输出,以及到每个节点的输入和输出,不能是标量 (Scalar)。即所有标量类型(如 float、int)都转换为 Tensor。
使用 Edge 方言¶
Edge 方言在内存中用 exir.EdgeProgramManager
Python 类表示。它包含一个或多个 torch.export.ExportedProgram
对象,其中包含方法的图表示。
import torch
from executorch import exir
class MyModule(torch.nn.Module):
...
a = MyModule()
tracing_inputs = (torch.rand(2, 2),)
aten_dialect_program = torch.export.export(a, tracing_inputs)
edge_dialect_program: exir.EdgeProgramManager = exir.to_edge(aten_dialect_program)
print(edge_dialect_program.exported_program)
此时,可以通过 edge_dialect_program.transform(pass)
运行用户定义的图转换。顺序很重要。注意:如果自定义 pass 正在修改 node.target
,请注意此时所有的 node.target
都是“Edge ops”(更多细节见下文),而不是像 ATen 方言中的 torch ops。有关 pass 编写的教程可以在 这里 找到。执行完所有这些 pass 后,to_edge()
将确保图仍然有效。
Edge 操作符¶
如前所述,Edge 操作符是具有数据类型专用的 ATen 核心操作符。这意味着 Edge 操作符的实例包含一组数据类型约束,这些约束描述了 ExecuTorch 运行时及其 ATen 内核支持的所有张量数据类型。这些数据类型约束以 edge.yaml 中定义的 DSL 表示。以下是数据类型约束的一个示例:
- func: sigmoid
namespace: edge
inherits: aten::sigmoid
type_alias:
T0: [Bool, Byte, Char, Int, Long, Short]
T1: [Double, Float]
T2: [Float]
type_constraint:
- self: T0
__ret_0: T2
- self: T1
__ret_0: T1
这表示如果 self
张量是 Bool, Byte, Char, Int, Long, Short
类型之一,那么返回的张量将是 Float
。如果 self
是 Double, Float
类型之一,则返回的张量将是相同的数据类型。
在收集了这些数据类型约束并将它们记录在 edge.yaml 中之后,EXIR 会消耗该文件,并将约束加载到 EXIR Edge 操作符中。这使得开发人员可以方便地了解 Edge op schema 中任何参数的受支持数据类型。例如,我们可以这样做:
from executorch.exir.dialects._ops import ops as exir_ops # import dialects ops
sigmoid = exir_ops.edge.aten.sigmoid.default
print(sigmoid._schema)
# aten::sigmoid(Tensor self) -> Tensor
self_arg = sigmoid._schema.arguments[0]
_return = sigmoid._schema.returns[0]
print(self_arg.allowed_types)
# {torch.float32, torch.int8, torch.float64, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool}
print(_return.allowed_types)
# {torch.float32, torch.float64}
这些约束对于编写该操作符的自定义内核非常有用。此外,在 EXIR 内部,我们提供了一个验证器来检查图在自定义转换后是否仍然符合这些数据类型约束。