StableHLO 發展藍圖

在撰寫本文時,StableHLO 已準備好取代 MHLO/HLO,成為編譯器介面。這類檔案可由 TensorFlow、JAX 和 PyTorch 產生,也能供 XLA 和多個第三方 PJRT 外掛程式使用,並且具備 MHLO/HLO 提供的所有公開功能,以及其他功能。

本文件說明 StableHLO 專案的後續步驟,將問題追蹤工具中顯示的持續性工作分類,並將這項工作安排至已規劃的成果。

目前的里程碑

我們目前的里程碑遵循兩大趨勢:

  1. 為整個 OpenXLA 社群盡量發揮 StableHLO 的效益。
  2. 為所有 OpenXLA 成員提供一致的開發人員體驗。
  • MHLO 淘汰:在 24 年第 4 季,我們開始探索 MHLO 的內部淘汰作業,將常用的傳遞作業 (包括標準化和資料夾模式) 遷移至 StableHLO。一旦內部確認遷移程序無須處理,我們就會分享 RFC,並提供遷移至 StableHLO 的外部時間表。這項作業很可能會在 2025 年第 1 季完成,我們預計會給予團隊充裕的時間和支援,以便在 2025 年第 1 季內遷移至 StableHLO。
  • 將硬體獨立最佳化功能遷移至 StableHLO:根據上述趨勢,我們希望 StableHLO 能成為整合硬體獨立圖表簡化的最佳位置,這樣所有 PJRT 外掛程式 (包括從 StableHLO 轉換至非 XLA 編譯器 IR 的程式) 都能獲得最大效益。這項目標的一部分,是將 Google AI EdgeJAX-Enzyme 專案和其他專案 (全部位於 StableHLO 存放區) 中使用的模式整合起來。部分整合作業已經開始,但大部分的工作流程將在 2025 年第 1 季完成。
  • OpenXLA 元件化:我們已開始在 openxla/xla 中為 HLO 建立專屬元件,類似於 StableHLO 存放區設定 (ref),並開始將所有 OpenXLA 後端移至 PJRT 外掛程式後方。我們也致力於修正這些 PJRT 外掛程式中發現的重大使用者體驗問題,包括在 StableHLO 外掛程式中提供精確的 StableHLO 版本通訊功能,以便新外掛程式能立即使用新功能 (參考資料)。
  • 讓組合運作 e2e:在 24 年第 3 季,我們在 HLO 中加入了組合,讓抽象層級能夠獲得完整的編譯器堆疊支援。在 Q4'24 中,我們教導 XLA inliner 有關複合函式的知識,並在 HLO/StableHLO 中新增了 pass,以便將未知的複合函式與其分解作業內嵌。我們目前正在研究如何新增專屬 JAX API,以便從架構中產生組合 (PyTorch API 已存在),以及新增 Colab 文件,說明如何正確使用組合,預計在 24 年第 4 季完成。

過去的里程碑

在 2024 年第一季,我們發布了 StableHLO 1.0,實作了高優先順序的改善項目,包括完成公開的 opset 規格、參考解譯器、擴充向前和向後相容性,以便透過相容性單元測試支援裝置端部署、透過複合運算擴充功能、指定的動態支援、發布完整的測試資料套件,並評估黃金結果等等。