导出 IR 规范¶
导出 IR 是 的结果的中间表示 (IR)。要了解有关 Export IR 的详细信息,请阅读此文档。torch.export
导出的 IR 是一个规范,由以下部分组成:
计算图模型的定义。
图表中允许的运算符集。
方言是由定义的操作组成的导出的 IR 图 ,但具有其他属性(例如对运算符集或 metadata) 进行
当前存在的 EXIR 方言是:
这些方言表示捕获的程序所经历的阶段 程序捕获转换为可执行格式。例如, ExecuTorch 编译过程从捕获 Python 程序开始到 ATen Dialect,则 ATen Dialect 转换为 Edge Dialect,Edge 转换为 Backend,并且 finally 转换为二进制格式执行。
ATen 方言¶
ATen dialect 将作为 ExecuTorch 编译的入口点 管道。这是 Eager 模式 PyTorch 程序第一次变成 Exported IR 图。在这个阶段,执行函数化,删除任何张量 别名和突变,并允许更灵活的图形转换 被制造。此外,所有张量都转换为 continuous 格式。
这种方言的目标是尽可能忠实地捕获用户的程序 (同时保持有效的 Exported IR)。用户已调用的已注册自定义运算符 在 EAGER 模式下,将在 ATen 方言中保持原样。然而,我们应该避免 通过 pass 在 Graph 中添加自定义操作。
目前,ATen dialect 的功能是进一步降低到 Edge dialect。 但是,在未来,我们可以将此视为 其他导出用例。
ATen Dialect 属性¶
ATen 方言图是有效的导出 IR 图,具有以下附加 性能:
节点中的所有运算符要么是 ATen 运算符(在命名空间中),要么是高阶运算符(如 control flow 运算符)或已注册的自定义运算符。已注册的自定义运算符是 注册到当前 PyTorch Eager 模式运行时的运算符,通常 with call (暗示 schema)。有关如何注册 自定义运算符
call_function
torch.ops.aten
TORCH_LIBRARY
每个 operator 还必须有一个 meta kernel。元内核是一个 函数,该函数在给定输入张量的形状时,可以返回 output 张量。有关如何编写元内核的详细信息,请点击此处。
输入值类型必须为 “Pytree-able”。因此,输出 types 也是 Pytree 支持的,因为所有运算符输出都是 Pytree 支持的。
ATen dialect 的 Ops 可以选择工作动态 dtypes、implicit type 张量的提升和隐式广播。
所有张量内存格式均采用 。
torch.contiguous_format
边缘方言¶
这种方言旨在引入对 Edge 有用的特化 设备,但不一定用于常规 (服务器) 导出。然而,我们仍然 保留进一步专注于每个不同的硬件。换句话说,我们 不想引入任何新的硬件依赖概念或数据;此外 那些已经存在于用户的原始 Python 程序中的。
Edge Dialect 属性¶
边缘方言图是有效的导出 IR 图,具有以下附加 性能:
OpCall 节点中的所有算子都来自预定义的算子集, 称为“Edge Operators”,或已注册的自定义运算符。Edge 运算符是 具有 dtype 特化的 ATen 运算符。这允许用户注册 仅适用于某些 dtypes 的内核,以减少二进制大小。
图形的输入和输出以及每个节点的输入和输出都不能是标量。即 所有标量类型(如 float、int)都转换为 Tensor。
使用 Edge Dialect¶
Edge 方言用 中的 Python 类表示
记忆。这包含一个或多个 s,这些
包含方法的图形表示形式。exir.EdgeProgramManager
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 ops” (更多详细信息见下文),而不是 ATen 方言中的 torch ops。
可以在此处找到有关传递编写的教程。毕竟,这些通行证是
executed,将确保图形仍然有效。edge_dialect_program.transform(pass)
node.target
node.target
to_edge()
边缘运算符¶
如前所述,边缘算子是类型为 专业化。这意味着 edge 运算符的实例包含一组 dtype 约束,用于描述 ExecuTorch 运行时及其 ATen 内核。这些 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
也就是说,如果 tensor 是 类型之一,则返回的 tensor 将是 。如果 是 之一,则返回张量将具有相同的 dtype。self
Bool, Byte, Char, Int, Long, Short
Float
self
Double, Float
在 edge.yaml 中收集并记录这些 dtype 约束后,EXIR 使用该文件,并将约束加载到 EXIR Edge 运算符中。这 便于开发人员了解任何参数支持的 dtypes 在 Edge op 架构中。例如,我们可以这样做:
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}
这些 constraints 对于想要为此 Operator 编写自定义内核的人来说很有帮助。同样在 EXIR 中,我们提供了一个验证器,用于检查图形在自定义转换后是否仍然符合这些 dtype 约束。