核心 ATen 運算符集定義¶
本頁提供核心 ATen 運算符集 (opset) 的描述和背景。建議正在為 ExecuTorch 開發新核心程式庫或委派的人員閱讀本頁。也建議熟悉 torch.export
作為先決條件;特別是 torch FX 圖形、運算符分解和函數化的概念。
已被識別為核心 ATen 運算子的運算子列表,可以在 PyTorch 文件網站的IRs 頁面上找到。
什麼是運算子集合?¶
torch.export
對給定的 PyTorch 程式執行完整的圖形捕獲,產生一個圖形 IR,用於描述程式執行的計算。運算子(即對 Tensor 執行的操作)是圖形中的基本計算單元,通常對應於圖形 IR 中的唯一節點。運算子的主要來源是 ATen 函式庫;除了 ATen 運算子之外,開發者也可以定義自己的運算子(即自定義運算子)。
“ATen 運算子集合”或“ATen opset”是可以被用於表示 PyTorch 程式,在程式被捕獲到圖形 IR 後,ATen 運算子的集合。
功能性 ATen 運算子集合¶
torch.export
的程式捕獲機制產生一個功能化的圖形,該圖形僅允許使用功能性運算子(即不變更或別名輸入的運算子)。因此,torch.export
產生一個將包含功能性 ATen opset 的圖形,該集合僅包含功能性 ATen 運算子。
核心 ATen 運算子集合¶
導出的圖形可以通過應用運算子分解進一步轉換。此過程將指定的 ATen 運算子替換為其他 ATen 運算子的等效序列。例如,aten.hardsigmoid
可以替換為 aten.clamp(aten.clamp(self + 3, min=0), max=6) / 6
。
如果使用預設分解設定分解 PyTorch 程式,則產生的圖形 IR 將包含“核心 ATen”opset。此 opset 將是功能性 ATen opset 的一個子集,因為某些運算子將被分解。作為核心 ATen opset 一部分的 ATen 運算子(即核心 ATen 運算子)將不會在預設分解設定下被分解。通常,核心 ATen 運算子不容易通過分解用其他 ATen 運算子重新表示。
核心 ATen opset 背後的主要動機是減少 PyTorch 後端和編譯器在模型匯出後需要處理的運算子數量。不僅 ATen 函式庫中定義了大量的運算子,而且可能會添加新的運算子,或者現有運算子的架構可能會更改。如果沒有運算子分解,建立在 torch.export
產生的 IR 之上的後端將不得不處理大量的運算子介面,以及一個不斷變化的 opset。核心 ATen opset 通過定義一個更小、更易於管理的運算子集合來解決這個問題,這個集合的開發考慮到了穩定性。
核心 ATen 運算子集合的開發¶
雖然 ExecuTorch 使用核心 ATen opset,但它並非 ExecuTorch 專屬。核心 ATen opset 的主要設計目標之一是盡可能通用;絕大多數用例不希望分解其中包含的運算子。因此,核心 ATen opset 暗示的分解應該對絕大多數用例都有用。
另一個關鍵考量是盡可能保持 opset 的最小化,但不能以犧牲對性能或開發者體驗產生重大負面影響的分解為代價。
核心 ATen opset 的開發是透過審查一個 ATen 運算子列表來完成的,該列表是透過調查公開 GitHub 倉庫中的模型以及眾所周知的開源模型來建立的。調查過程的目的是獲得一個簡化的 ATen 運算子列表,該列表代表最常使用的 ATen 運算子。這樣,就可以首先審查最常用的運算子。
每個運算子是否應該是核心運算子,或者是否應該由核心 ATen 分解表分解的決定,取決於:
檢查運算子的潛在分解;分解應該是使用其他 ATen 運算子對運算子進行相對簡單的重新表示。
分解不應該看起來像是運算子的直接實現。
分解不應該根據輸入的執行時特性而變化。
我們還會考慮分解運算子是否會影響輸出的精度、數值有效性或記憶體佈局。
考慮開發人員是否希望在圖形中保留運算子以實現性能或其他原因。
例如,也許一個運算子可以被分解,但它可以在大多數平台上映射到單個硬體指令,在這種情況下,最好將其提升為核心運算子。
未來工作¶
在每個 ATen 運算子都被審查並給予“核心”或“預設分解”的指定之前,核心 ATen opset 不能被認為是完全完整的。然而,這是一項巨大的任務,並且存在大量不常用的運算子。這就是為什麼採用了一種方法,透過調查模型來確定哪些運算子最常用,從而可以優先考慮「影響較大」的運算子。
儘管如此,仍然有許多運算子尚未評估。計畫是根據需要繼續評估其他運算子;PyTorch 社群可以通過開啟 GitHub issue 或通過在 PyTorch 論壇上的這個帖子上發表評論來提出額外的核心運算子或額外的核心分解。