匯出 IR 規格¶

Export IR 是 torch.export 結果的中間表示形式 (Intermediate Representation, IR)。若要深入了解 Export IR 的詳細資訊，請閱讀此文件。

Exported IR 是一個規格，包含以下部分：

計算圖模型的定義。
圖中允許的運算符集合。

一個方言 (dialect) 是一個 Exported IR 圖，由以下定義的運算組成，但具有額外的屬性（例如對運算符集合或元數據的限制），旨在用於特定目的。

目前存在的 EXIR 方言有：

ATen 方言
Edge 方言
後端方言

這些方言代表一個捕獲的程式從程式捕獲到轉換為可執行格式所經歷的階段。例如，ExecuTorch 編譯過程從將 Python 程式捕獲到 ATen 方言開始，然後將 ATen 方言轉換為 Edge 方言，Edge 方言轉換為 Backend，最後轉換為用於執行的二進位格式。

ATen 方言¶

ATen 方言將用作 ExecuTorch 編譯管道的入口點。這是 eager mode PyTorch 程式首次變成 Exported IR 圖。在此階段，會執行功能化 (functionalization)，移除任何張量別名和突變，並允許進行更靈活的圖轉換。此外，所有張量都會轉換為連續格式。

此方言的目標是盡可能忠實地捕獲使用者的程式（同時保持有效的 Exported IR）。使用者在 eager mode 中呼叫的已註冊自定義運算符將按原樣保留在 ATen 方言中。但是，我們應避免通過 pass 在圖中添加自定義運算符。

目前，ATen 方言的功能是進一步降級到 Edge 方言。但是，在未來，我們可以將其視為其他 export 用例的通用整合點。

ATen 方言屬性¶

ATen 方言圖是一個有效的 Export IR 圖，具有以下額外屬性：

call_function 節點中的所有運算符都是 ATen 運算符（在 torch.ops.aten 命名空間中）、高階運算符（例如控制流運算符）或已註冊的自定義運算符。已註冊的自定義運算符是註冊到當前 PyTorch eager mode 運行時中的運算符，通常使用 TORCH_LIBRARY 呼叫（暗示 schema）。有關如何註冊自定義運算符的詳細資訊，請參閱此處。
每個運算符也必須具有 meta kernel。meta kernel 是一個函數，給定輸入張量的形狀，可以返回輸出張量的形狀。有關如何編寫 meta kernel 的詳細資訊，請參閱此處。
輸入值類型必須是 “Pytree-able”。因此，輸出類型也是 Pytree-able，因為所有運算符的輸出都是 pytree-able。
ATen 方言的 Ops 可以選擇使用 Dynamic dtypes、隱式類型提升和張量的隱式廣播。
所有張量內存格式都在 torch.contiguous_format 中。

ATen 運算符定義¶

運算符集合定義可在此處找到。

Edge 方言¶

此方言旨在引入對 Edge 設備有用的專業化，但不一定對常規 (伺服器) export 有用。但是，我們仍然避免進一步專業化到每個不同的硬體。換句話說，我們不希望引入任何新的硬體相關概念或數據；除了使用者原始 Python 程式中已存在的概念或數據。

Edge 方言屬性¶

Edge 方言圖是一個有效的 Export IR 圖，具有以下額外屬性：

OpCall 節點中的所有運算符都來自預定義的運算符集合，稱為“Edge 運算符”，或已註冊的自定義運算符。Edge 運算符是具有 dtype 專業化的 ATen 運算符。這允許使用者註冊僅適用於某些 dtypes 的 kernel，以減少二進位大小。
圖的輸入和輸出，以及每個節點的輸入和輸出，不能是 Scalar。即，所有標量類型（例如 float、int）都會轉換為 Tensor。

使用 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 ops”（更多詳細資訊如下），而不是像 ATen 方言中的 torch ops。有關 pass 編寫的教程，請參閱此處。執行完所有這些 pass 後，to_edge() 將確保圖仍然有效。

Edge 運算符¶

如前所述，edge 運算符是具有類型專業化的 ATen 核心運算符。這意味著 edge 運算符的實例包含一組 dtype 約束，這些約束描述了 ExecuTorch 運行時及其 ATen kernel 支持的所有張量 dtypes。這些 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。

在 edge.yaml 中收集並記錄這些 dtype 限制後，EXIR 會讀取該檔案，並將這些限制載入到 EXIR Edge 運算子中。這使得開發人員可以方便地學習 Edge op schema 中任何參數支援的 dtypes。例如，我們可以執行：

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 限制。

Op Set (WIP)¶

請查看 edge.yaml 以取得指定 dtype 限制的運算子的完整清單。我們正在逐步擴展此運算子集，並致力於為所有核心 ATen ops 提供 dtype 限制。

Backend Dialect¶

請參閱此文件