核心 ATen 運算元集定義¶
本頁介紹了核心 ATen 運算元集 (opset) 的描述和背景資訊。本頁推薦給正在為 ExecuTorch 開發新核心庫或委託的開發者閱讀。此外,建議先熟悉 torch.export 作為前置條件;特別是要了解 torch FX 圖、運算元分解和函式化的概念。
被確定為核心 ATen 運算元的列表可在 PyTorch 文件網站的 IRs 頁面找到。
什麼是運算元集?¶
torch.export 對給定的 PyTorch 程式執行完整的圖捕獲,生成一個描述程式執行計算的圖 IR。運算元(即對張量執行的操作)是圖中計算的基本單元,通常對應於圖 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 庫中不僅定義了大量運算元,而且可能會新增新運算元,或現有運算元的 schema 可能會發生變化。如果沒有運算元分解,構建在 torch.export 生成的 IR 之上的後端將不得不處理大量的運算元以及不斷變化的 opset。核心 ATen opset 透過定義一個更小、更易於管理的運算元集來解決這個問題,該運算元集在開發時考慮了穩定性。
核心 ATen 運算元集的開發¶
儘管 ExecuTorch 使用核心 ATen opset,但它並非 ExecuTorch 獨有。核心 ATen opset 的主要設計目標之一是儘可能通用;絕大多數用例不會希望分解其中包含的運算元。因此,核心 ATen opset 所隱含的分解應該對絕大多數用例有用。
另一個關鍵考慮因素是儘可能保持 opset 最小化,但不能以強加對效能或開發者體驗產生深遠負面影響的分解為代價。
核心 ATen opset 是透過調研公共 GitHub 倉庫中的模型以及知名開源模型,梳理出 ATen 運算元列表後進行審查而開發的。調研過程的目的是獲得一個精簡的 ATen 運算元列表,該列表可以大致反映哪些 ATen 運算元使用最廣泛。這樣就可以優先審查最常用的運算元。
關於每個運算元是應該作為核心運算元還是由核心 ATen 分解表進行分解的決定是根據以下因素確定的:
檢查運算元的潛在分解;分解應該是使用其他 ATen 運算元對該運算元進行的相對直接的重新表達。
分解不應該看起來像是對運算元的完全實現。
分解不應該隨輸入的執行時特性而變化。
我們還考慮分解該運算元是否會影響輸出的精度、數值有效性或記憶體佈局。
考慮開發者是否會出於效能或其他原因希望在圖中保留該運算元。
例如,某個運算元或許可以分解,但在大多數平臺上它可以對映到單個硬體指令,在這種情況下,最好將其提升為核心運算元。