本文將說明 XLA 高階最佳化工具 (HLO) 模組從初始狀態到最終可執行檔的歷程。有時我們會省略「模組」，只稱其為「HLO」。
最佳化前 HLO
首先，我們從預先最佳化的 HLO 模組開始。最佳化前 HLO 不包含視為 XLA 內部運算的運算 (ops)，例如 fusion 或 bitcast。作業在這個階段沒有版面配置，或是有版面配置但會遭到忽略。最佳化前 HLO 通常是由 TensorFlow 和 JAX 等高階架構產生。使用 XLA 旗標 -xla_dump_to 時，最佳化前的 HLO 會傾印至檔案，檔案名稱後置字串為「before_optimizations.txt」。
最佳化 HLO 模組
XLA：GPU 管道會執行一連串的傳遞作業，將最佳化前的 HLO 轉換為最佳化後的 HLO。這些傳遞內容可以依語意分組，並依下列順序執行：
分片相關票證
包括 Shardy Partitioner 等傳遞項目，或是 SPMD 分片項目。
最佳化過程
包括合法化和簡化程序。
集體最佳化階段
與最佳化階段類似，但著重於集體作業。
版面配置指派傳遞
每個 HLO 作業都會指派一個版面配置，這是指令形狀的一部分。版面配置會控管張量在記憶體中的實際配置方式。
含有版面配置的形狀範例：
f32[10,20,30]{2,0,1}
元素類型之後是形狀的邏輯維度，接著是次要到主要順序的版面配置排列。在本例中，最小的次要維度是 30，第二小的次要維度是 10，主要維度則是 20。
版面配置指派的目標是使用貪婪策略，盡可能減少所需的實體轉置次數。這項程序會從特定版面配置限制開始 (例如，cuDNN/cuBLAS 程式庫預期會有連續維度)，然後在 HLO 圖表中「向下」和「向上」傳播版面配置。在版面配置傳播結束時，部分指令可能會有衝突的版面配置，一個是從運算元傳播，另一個是從使用者傳播。為解決這項衝突，系統會插入 copy HLO 指令，將版面配置從運算元版面配置變更為指令版面配置。
版面配置正規化傳遞
由於實體形狀較難判斷，版面配置正規化會嘗試重新編寫形狀，使其使用預設版面配置 {rank-1, rank-2, …, 0}。在上述範例中，正規化形狀為 f32[20,10,30]{2,1,0}。變更版面的複製作業會重新編寫為 transpose 和 bitcast 的組合。由於我們目前無法將所有作業正規化，因此仍有部分作業可能採用非預設的版面配置，最明顯的例子就是 gather 和 dot。在正規化作業和非正規化作業的界線之間，會有 bitcast 作業代表轉置，也就是指派版面配置的轉置，使其在實體上成為無作業。
版面配置正規化也會明確表示一些隱含的轉置，這很重要，因為程式碼產生器可以使用專用發射器處理明確的轉置。舉例來說，從技術上來說，重塑作業允許運算元和結果之間有不同的實體配置 (例如，由於等級不同)。在版面配置正規化傳遞中執行的 ReshapeDecomposer 傳遞，會將重塑轉換為 transpose、重塑 bitcast 和 transpose 的序列。
版面配置指派最佳化傳遞後
這裡最重要的傳遞內容是 Triton 融合 (GEMM 融合 + Softmax/Layernorm 融合) 或重新編寫至程式庫呼叫。自動調校也會在這個步驟中執行，XLA 會在不同發射器之間選擇，為捲積或點積挑選最佳演算法，為 Triton 發射器處理的融合作業找出最佳分塊等。
Fusion 票證
主要有 PriorityFusion 和 Multi-Output 融合這兩種。
在 PriorityFusion 中，我們會根據成本模型形成融合。融合時，如果作業可融合到所有使用者，我們會允許作業與多位使用者重複。如果可以，我們也會允許擴充現有的 Triton Softmax 融合。
Multi-Output 融合是獨立的傳遞，可融合共用運算元的作業/融合。此外，它也可以透過新增額外輸出，將運算元/運算元融合併入使用者，而不會重複，因此要融合的運算元其他使用者可以重新導向這些輸出。這個傳遞程序必須小心，以免在 HLO 圖表中導入週期。
多重輸出融合後，系統會執行常見的子運算式消除 (HloCSE 傳遞)，如果先前重複的作業最終出現在同一個融合中，可能會合併這些作業。
多張融合後票證
與集合相關的多次傳遞 (例如將集合轉為非同步，或強制執行集合的特定相對順序)。
最後，我們會執行 CopyInsertion，加入副本，確保就地作業不會覆寫其他地方仍需要的資料。
最佳化完成後，如果使用 -xla_dump_to 標記，最佳化 HLO 會傾印至檔案，檔案名稱後置字串為「after_optimizations.txt」。如要在實際變更 HloModule 的中繼傳遞後傾印 HLO，可以使用 -xla_dump_hlo_pass_re=.* 旗標 (或特定規則運算式，將其限制為特定傳遞)。
排程
沒有時間表的 HLO 模組在處理作業的順序方面仍具有某種程度的自由度。任何尊重運算元/結果關係和控制項依附元件的拓撲排序都是有效的。排程會決定要使用哪個特定順序。這個階段的主要考量是記憶體用量上限，這取決於張量的生命週期。在初始步驟中，我們會嘗試不同的排程器演算法，並選擇可將尖峰記憶體用量降到最低的排程。請注意，目前我們還不會使用實體緩衝區 (這會在「緩衝區指派」中發生)，並模擬記憶體用量。
然後，LatencyHidingScheduler 傳遞會執行，並盡量提高運算與通訊的重疊程度。但這樣做可能會再次增加記憶體用量。
最後，如果記憶體尖峰用量高於可用記憶體量，我們會執行 HloRematerialization。這個階段會嘗試減少記憶體用量，但會犧牲效能，例如可能會分割部分融合，並複製部分作業，以縮短緩衝區生命週期。如果發生重新具體化，建議您調查如何減少模型端的記憶體需求 (例如使用較小的批次大小)。
緩衝區指派
在降低至 LLVM IR 之前，我們會立即執行緩衝區指派傳遞，將緩衝區切片指派給 HLO 圖表中的每個指令。緩衝區指派作業會分幾個步驟執行：
HloDataflowAnalysis會將HloValues(基本上是邏輯緩衝區) 指派給指令。對於就地作業，運算元的HloValue可以重複使用。一個作業可能會定義多個HloValue(例如使用元組結果形狀)。HloAliasAnalysis會嘗試合併緩衝區以進行別名作業，並計算從HloValue到HloBuffer的對應。BufferAssignment會計算HloBuffers對應至大緩衝區內緩衝區切片的對應項，確保不會將相同緩衝區切片用於生命週期重疊的不同HloBuffers。對於可能會別名的作業，稍微重疊是可以的 (一個HloBuffer的結束時間可能與另一個HloBuffer的開始時間一致)。使用-xla_dump_to標記時，緩衝區指派的相關資訊會傾印至檔案，檔案名稱後方會加上「after_optimizations-buffer-assignment.txt」後置字串。
Thunks
HLO 圖表經過最佳化和排程後，會降級為特定後端 (CPU 或 GPU) 的線性 thunk 序列。
在 XLA 中，Thunk 是執行階段執行的獨立工作單元抽象化。可能是已編譯的核心啟動、特定作業、程式庫呼叫、控制流程建構、集體通訊等。「Thunk 序列」代表特定後端的整個可執行檔。
Thunk 排放
將排定的 HLO 計算轉換為 thunk 序列的過程稱為「thunk 發出」。每個後端都有專屬的發射器類別，負責處理這項作業。
如果是 GPU 後端，則由 IrEmitterUnnested 處理。EmitHloComputation 會在運算中逐一查看排定的 HLO 指令清單，並分派至專門的 Emit... 方法 (例如EmitFusion、EmitConvolutionThunk、EmitWhile)。這些方法都會建構適當的 Thunk 物件，並將其附加至 thunk 序列。
對於 CPU 後端，ThunkEmitter 會執行這項角色，並以類似方式整理。最終 ThunkSequence 會嵌入 CpuExecutable 中。
請注意，HLO 模組的項目計算中，每項指令可能對應到最終 thunk 序列中的零個 (kTuple、kConstant 等)、一個或多個 (例如排序指令) thunk。
命令緩衝區：最佳化 GPU 執行作業
現代 GPU 硬體可讓您記錄一次 GPU 作業序列 (核心啟動、記憶體副本等)，然後以最少的 CPU 負荷多次重播該序列。這項效能最佳化功能至關重要，特別是對於有許多小型快速啟動核心的工作負載。XLA 會使用「命令緩衝區」做為 CUDA 圖形或 HIP 圖形的抽象化。核心介面是在 GpuCommandBuffer 中定義。
指令緩衝區在 thunk 序列中以 CommandBufferThunk 表示。
發射器不會直接從 HLO 指令產生這個 thunk。而是由在 ThunkSequence 本身執行的 CommandBufferConversionPass 完成。
這個階段會找出相容的連續子序列 Thunk (例如一系列的 KernelThunk 和 GemmThunk)。然後以單一 CommandBufferThunk 取代找到的子序列。新的 thunk 會將原始 thunk 的邏輯封裝為輕量型 CommandBufferCmd 物件的清單。當 CommandBufferThunk 首次在特定 GPU 串流上執行時，會將指令序列「記錄」到硬體指令緩衝區。在後續的所有執行作業中，它只會向 GPU 發出單一指令，以「重播」記錄的序列。這樣可避免啟動每個個別核心時，CPU 產生額外負荷。
執行檔
XLA 編譯管線的最終產品是獨立的平台專屬可執行檔。這個物件會封裝在目標裝置 (例如 CPU 或 GPU) 上執行已編譯程式所需的所有資訊。這是編譯器和執行階段之間的橋樑。PJRT 等現代執行階段會使用稍高層級的抽象概念 (請參閱 PjRtExecutable)，但這些概念最終會包裝後端專屬的可執行檔。
Executable 包含編譯期間產生的幾項重要資訊。雖然確切內容會因後端而異，但通常包括：
已編譯的程式碼：這是將在裝置上執行的低階機器碼。 如果是 CPU，通常是一或多個物件檔案。如果是 GPU，這是 PTX 或 HSACO 格式的已編譯裝置程式碼，會在執行階段載入 GPU。
執行計畫 (ThunkSequence)：執行階段邏輯的核心。這是 Thunk 物件的線性序列。每個 thunk 都代表一個工作單元，例如啟動核心、呼叫程式庫函式 (例如 cuBLAS) 或處理控制流程。執行階段會透過這個序列進行疊代，藉此執行程式。
記憶體配置 (BufferAssignment)：這項重要的中繼資料是由 BufferAssigner 產生，說明運算的完整記憶體配置。這個物件會指定每個緩衝區的大小，以及如何為參數、輸出和暫時值分配及重複使用記憶體。執行階段會使用這項資訊分配裝置記憶體，並將正確的指標傳遞至每個 thunk。
(選用) HLO 模組：為了進行偵錯和剖析，可執行檔通常會保留對最終最佳化 HloModule 的參照，該模組是從中編譯而來。
編譯器會為每個特定後端協調最終可執行檔的建立作業。編譯器實作的 RunBackend 方法是編譯程序的最後一個步驟，會將所有編譯的構件封裝到 Executable 物件中。GpuCompiler 和 CpuCompiler 分別以 GPU 和 CPU 為目標。
當使用者在可執行檔上呼叫 Execute... 時，執行階段會使用 BufferAssignment 分配記憶體，然後叫用 ThunkSequence，在裝置上使用已編譯的程式碼啟動作業。