导出 IR 规范¶

导出 IR 是 torch.export 结果的中间表示 (IR)。要详细了解导出 IR，请阅读此文档。

导出的 IR 是一种规范，包含以下部分

计算图模型定义。
图中允许的运算符集。

方言是使用下面定义的运算组成的导出 IR 图，但具有旨在用于特定目的的附加属性（例如对运算符集或元数据的限制）。

当前存在的 EXIR 方言是

ATen 方言
Edge 方言
后端方言

这些方言表示捕获的程序从程序捕获到转换为可执行格式所经历的阶段。例如，ExecuTorch 编译过程从将 Python 程序捕获到 ATen 方言开始，然后将 ATen 方言转换为 Edge 方言，Edge 方言转换为后端方言，最后转换为用于执行的二进制格式。

ATen 方言¶

ATen 方言将用作 ExecuTorch 编译管道的入口点。这是 eager 模式 PyTorch 程序首次变为导出 IR 图。在此阶段，执行函数化，删除任何张量别名和变异，从而可以进行更灵活的图转换。此外，所有张量都转换为连续格式。

此方言的目标是尽可能忠实地捕获用户的程序（同时保持有效的导出 IR）。用户在 eager 模式下调用的已注册自定义运算符将在 ATen 方言中保持原样。但是，我们应避免通过 Pass 在图中添加自定义运算符。

目前，ATen 方言的功能是进一步降低到 Edge 方言。但是，将来我们可以将其视为其他导出用例的通用集成点。

ATen 方言属性¶

ATen 方言图是具有以下附加属性的有效导出 IR 图

call_function 节点中的所有运算符要么是 ATen 运算符（在 torch.ops.aten 命名空间中）、高阶运算符（如控制流运算符）或已注册的自定义运算符。已注册的自定义运算符是注册到当前 PyTorch eager 模式运行时中的运算符，通常通过 TORCH_LIBRARY 调用（暗示模式）。有关如何注册自定义运算符的详细信息，请访问此处。
每个运算符还必须具有元内核。元内核是一个函数，给定输入张量的形状，可以返回输出张量的形状。有关如何编写元内核的详细信息，请访问此处。
输入值类型必须是“Pytree-able”。因此，输出类型也是 Pytree-able，因为所有运算符输出都是 pytree-able。
ATen 方言的运算可以选择使用动态数据类型、隐式类型提升和张量的隐式广播。
所有张量内存格式均为 torch.contiguous_format。

ATen 运算符定义¶

运算符集定义可在此处找到。

Edge 方言¶

此方言旨在引入对边缘设备有用但不一定对通用（服务器）导出有用的专业化。但是，我们仍然不进一步专门针对每种不同的硬件。换句话说，我们不想引入任何新的硬件相关概念或数据；除了用户原始 Python 程序中已存在的那些。

Edge 方言属性¶

Edge 方言图是具有以下附加属性的有效导出 IR 图

OpCall 节点中的所有运算符要么来自预定义的运算符集（称为 “Edge 运算符”），要么是已注册的自定义运算符。Edge 运算符是具有数据类型专业化的 ATen 运算符。这允许用户注册仅适用于某些数据类型的内核，以减小二进制大小。
图的输入和输出，以及每个节点，都不能是标量。即，所有标量类型（例如 float、int）都转换为张量。

使用 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)
print(edge_dialect_program.exported_program)

此时，用户定义的图转换可以通过 edge_dialect_program.transform(pass) 运行。顺序很重要。注意：如果自定义 Pass 正在接触 node.target，请注意，此阶段的所有 node.target 都是“Edge 运算”（更多详细信息见下文），而不是 ATen 方言中的 torch 运算。有关 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 运算模式中任何参数的受支持数据类型。例如，我们可以执行以下操作

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 内部，我们提供了一个验证器来检查在自定义转换后，图是否仍然符合这些数据类型约束。

运算符集 (WIP)¶

查看 edge.yaml，获取指定了数据类型约束的运算符的完整列表。我们正在逐步扩展此运算符集，并致力于为所有核心 ATen 运算提供数据类型约束。

后端方言¶

请参阅此文档