XLA 架構

XLA (加速線性代數) 是一種機器學習 (ML) 編譯器,可將線性代數最佳化,可改善執行速度和記憶體用量。本頁提供 XLA 編譯器目標和架構的簡要總覽。

目標

如今,XLA 支援多個機器學習架構前端 (包括 PyTorch、TensorFlow 和 JAX),且屬於 OpenXLA 專案的一部分。OpenXLA 是適用於機器學習的開放原始碼編譯器技術生態系統,由頂尖的機器學習硬體和軟體機構協同開發。建立 OpenXLA 專案前,XA 是在 TensorFlow 專案中開發,但基本目標仍保持不變:

  • 加快執行速度。編譯子圖形以縮短短期作業的執行時間,並消除執行階段的負擔;融合管道的作業以降低記憶體負擔,以及專門處理已知的張量形狀,以便進行更積極的持續傳播。

  • 改善記憶體用量。分析及排定記憶體用量,消除多個中繼儲存空間緩衝區。

  • 減少對自訂作業的依賴。改善自動融合的低層級運算效能,滿足原本手動融合自訂運算的效能,以移除對許多自訂作業的需求。

  • 提升可攜性。輕鬆為新款硬體編寫新後端,讓大量的機器學習模型可以在該硬體上執行未經修改的。這與將個別單體式運算用於新硬體的做法相比,這種做法需要重寫模型,才能使用這些運算。

運作方式

XLA 編譯器會從 StableHLO 中定義的機器學習架構取得模型圖形,並編譯成不同架構的機器操作說明。StableHLO 會定義版本化作業集 (HLO = 高階作業),提供機器學習架構和編譯器之間的可攜性層。

一般來說,將模型圖轉換為目標最佳化執行檔的編譯程序包含下列步驟:

  1. XLA 會在目標獨立的 StableHLO 圖形上執行多項內建最佳化和分析傳遞,例如 CSE、目標獨立作業融合,以及用於分配執行階段記憶體以進行運算的緩衝區分析。在這個最佳化階段,XLA 還會將 StableHLO 方言轉換為內部 HLO 方言。

  2. XLA 會將 HLO 運算傳送至後端,以便進一步進行 HLO 層級最佳化,這次只是考量目標特定的資訊和需求。舉例來說,GPU 後端可能會執行特別適合 GPU 程式設計模型的作業融合,並判斷如何將運算分割為串流。在這個階段,後端可能也會以模式比對特定作業或其中一組作業,與最佳化的程式庫呼叫。

  3. 後端接著會執行特定目標的程式碼產生作業。XLA 隨附的 CPU 和 GPU 後端會使用 LLVM 進行低階 IR、最佳化和程式碼產生。這些後端會發出高效率方式代表 HLO 運算所需的 LLVM IR,然後叫用 LLVM,從這個 LLVM IR 發出原生程式碼。

在這個程序中,XLA 編譯器的模組化,因此在替代後端可輕鬆鎖定某些新的硬體架構。GPU 後端目前透過 LLVM NVPTX 後端支援 NVIDIA GPU。CPU 後端支援多個 CPU ISAs。