导出 IR 规范¶
导出 IR 是 torch.export
结果的中间表示 (IR)。要详细了解导出 IR,请阅读此 文档.
导出的 IR 是一个规范,包含以下部分
计算图模型的定义。
图中允许的操作符集合。
**方言**是一种由以下定义的操作符组成的导出 IR 图,但具有额外的属性(例如操作符集或元数据的限制),这些属性旨在用于特定目的。
目前存在的 EXIR 方言有:
这些方言代表了捕获的程序从程序捕获到转换为可执行格式所经历的阶段。例如,ExecuTorch 编译过程从 Python 程序捕获到 ATen 方言开始,然后将 ATen 方言转换为 Edge 方言,Edge 转换为后端,最后转换为二进制格式以供执行。
ATen 方言¶
ATen 方言将用作 ExecuTorch 编译管道的入口点。这是急切模式 PyTorch 程序首次成为导出 IR 图。在此阶段,将执行函数化,删除所有张量别名和变异,并允许进行更灵活的图转换。此外,所有张量都将转换为连续格式。
该方言的目标是尽可能忠实地捕获用户的程序(同时保持有效的导出 IR)。用户在急切模式下调用的注册自定义操作符将按原样保留在 ATen 方言中。但是,我们应该避免通过传递在图中添加自定义操作符。
目前,ATen 方言的功能是进一步降低到 Edge 方言。但是,在未来,我们可以将其视为其他导出用例的通用集成点。
ATen 方言属性¶
ATen 方言图是一个有效的导出 IR 图,具有以下附加属性:
所有在
call_function
节点中的操作符要么是 ATen 操作符(在torch.ops.aten
命名空间中),要么是高阶操作符(如控制流操作符),要么是注册的自定义操作符。注册的自定义操作符是注册到当前 PyTorch 急切模式运行时的操作符,通常使用TORCH_LIBRARY
调用(意味着模式)。有关如何注册自定义操作符的详细信息,请参见 此处。每个操作符都必须有一个元内核。元内核是一个函数,它根据输入张量的形状,可以返回输出张量的形状。有关如何编写元内核的详细信息,请参见 此处。
输入值类型必须是“可 Pytree 化”。因此,输出类型也是可 Pytree 化的,因为所有操作符的输出都是可 Pytree 化的。
ATen 方言的操作可以选择使用动态数据类型、隐式类型提升和张量的隐式广播。
所有张量的内存格式都为
torch.contiguous_format
。
边缘方言¶
该方言旨在引入对边缘设备有用的专门化,但不一定适用于通用(服务器)导出。但是,我们仍然保留对每个不同硬件的进一步专门化。换句话说,我们不想引入任何新的硬件相关概念或数据;除了用户原始 Python 程序中已经存在的那些。
边缘方言属性¶
边缘方言图是一个有效的导出 IR 图,具有以下附加属性
OpCall 节点中的所有操作符要么来自预定义的操作符集,称为“边缘操作符”,要么来自注册的自定义操作符。边缘操作符是具有数据类型专门化的 ATen 操作符。这允许用户注册仅适用于某些数据类型的内核,以减小二进制文件大小。
图的输入和输出,以及每个节点的输入和输出,都不能是标量。即所有标量类型(如 float、int)都转换为张量。
使用边缘方言¶
边缘方言用内存中的 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)
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()
将确保图仍然有效。
边缘运算符¶
如前所述,边缘运算符是具有类型特化的 ATen 核心运算符。这意味着边缘运算符的实例包含一组 dtype 约束,这些约束描述了 ExecuTorch 运行时及其 ATen 内核支持的所有张量 dtype。这些 dtype 约束在 edge.yaml 中定义的 DSL 中表达。以下是一个 dtype 约束的示例
- 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
之一,则返回张量将具有相同的 dtype。
在收集这些 dtype 约束并在 edge.yaml 中记录后,EXIR 会使用该文件,并将约束加载到 EXIR Edge 运算符中。这使得开发人员可以方便地了解 Edge op 模式中任何参数支持的 dtype。例如,我们可以执行以下操作
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 内部,我们提供了一个验证器来检查图在自定义转换后是否仍然符合这些 dtype 约束。