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