匯出 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 約束。

後端方言¶

請參閱此文件