导出 IR 规范¶

Export IR 是 torch.export 结果的中间表示 (IR)。要了解更多关于 Export IR 的详细信息，请阅读此文档。

导出的 IR 规范包含以下部分

计算图模型的定义。
图中允许使用的算子集。

一个 方言 (dialect) 是一个由下方定义的操作组成的导出 IR 图，但具有额外的属性（例如对算子集或元数据的限制），旨在用于特定目的。

当前存在的 EXIR 方言有

ATen 方言
Edge 方言
后端方言

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

ATen 方言¶

ATen 方言将用作 ExecuTorch 编译流程的入口点。这是急切模式 (eager mode) PyTorch 程序首次成为导出的 IR 图。在此阶段，会进行功能化处理，移除任何张量别名和变异，并允许进行更灵活的图转换。此外，所有张量都会被转换为连续格式。

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

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

ATen 方言属性¶

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

所有 call_function 节点中的算子都是 ATen 算子（在 torch.ops.aten 命名空间中）、高阶算子（如控制流算子），或已注册的自定义算子。已注册的自定义算子是注册到当前 PyTorch 急切模式运行时的算子，通常通过调用 TORCH_LIBRARY（隐含 schema）。关于如何注册自定义算子的详细信息可以在此处找到。
每个算子还必须有一个 meta kernel。meta kernel 是一个函数，给定输入张量的形状，可以返回输出张量的形状。关于如何编写 meta kernel 的详细信息可以在此处找到。
输入值类型必须是“Pytree-able”。因此，输出类型也是 Pytree-able 的，因为所有算子的输出都是 pytree-able 的。
ATen 方言的算子可以选择支持动态数据类型 (Dynamic dtypes)、隐式类型提升 (implicit type promotions) 和张量的隐式广播 (implicit broadcasting)。
所有张量的内存格式都是 torch.contiguous_format。

ATen 算子定义¶

算子集定义可以在此处找到。

Edge 方言¶

此方言旨在引入对边缘设备有用的专用化，但不一定适用于通用（服务器）导出。然而，我们仍然避免针对每种不同的硬件进行进一步的专用化。换句话说，除了用户原始 Python 程序中已有的概念或数据外，我们不想引入任何新的依赖于硬件的概念或数据。

Edge 方言属性¶

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

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

使用 Edge 方言¶

Edge 方言在内存中由 Python 类 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_dialect_program.transform(pass) 运行用户定义的图转换。顺序很重要。注意：如果自定义 pass 触及 node.target，请注意在此阶段所有的 node.target 都是“Edge 算子”（更多细节如下），而不是像在 ATen 方言中那样的 torch 算子。关于编写 pass 的教程可以在此处找到。所有这些 pass 执行后，to_edge() 将确保图仍然有效。

Edge 算子¶

如前所述，Edge 算子是具有类型专用化的 ATen 核心算子。这意味着 Edge 算子的一个实例包含一组 dtype 约束，描述了 ExecuTorch 运行时及其 ATen kernel 都支持的所有张量 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 算子 schema 中任何参数支持的 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}

这些约束对于希望为此算子编写自定义 kernel 的人很有帮助。此外，在 EXIR 内部，我们提供一个验证器来检查在自定义转换后，图是否仍然符合这些 dtype 约束。

算子集 (进行中)¶

请查阅 edge.yaml 获取指定了 dtype 约束的完整算子列表。我们正在逐步扩展此算子集，并计划为所有核心 ATen 算子提供 dtype 约束。

后端方言¶

请参阅此文档