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。