XLA 架構

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

目標

XLA 目前支援多個機器學習架構前端 (包括 PyTorch、TensorFlow 和 JAX),並且屬於 OpenXLA 專案。這項專案是開放原始碼編譯器技術的生態系統,由頂尖機器學習硬體和軟體機構共同開發。建立 OpenXLA 專案之前,XLA 是在 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 編譯器是模組化的一環,很容易將運算單元放置在其他後端,以便指定某些新的 HW 架構。GPU 後端目前透過 LLVM NVPTX 後端支援 NVIDIA GPU。CPU 後端支援多個 CPU ISA。