本頁面由 Cloud Translation API 翻譯而成。
Switch to English

XLA的MLIR CodeGen

XLA在HloInstruction運行,並對此表示形式進行了許多優化,在目標設備之間共享了許多優化。從某種意義上講,將計算線性計劃,並且將內存緩衝區靜態分配給每個值。特定於設備的代碼生成器通過遍歷此序列並調用“發射器”來生成適合於該設備的表示(例如,CPU上每個XLA計算使用單個LLVM函數,或者封裝GPU操作的一系列“ thunk”,可能在生成時生成PTX)定位GPU)。

作為暫存步驟,我們目前正在XLA完成緩衝區分配階段之後立即攔截該過程,而改為在lhlo方言中發出MLIR模塊。從那裡開始,我們根據設備使用MLIR組件(主要是Linalg,仿射和GPU方言)執行代碼生成。

下面是通過使用lhlo作為代碼源輸入來逐步遷移XLA / GPU的記錄計劃。

任務

主辦設備
輸入格式 HloInstruction *(任務1) HloInstruction *(任務1)
輸出格式 xla :: Thunk(任務2) LLVM IR(任務3)
  • 任務1將主機和設備輸入格式從HloInstruction *更改為LHLO。
  • 任務2將主機的輸出格式從thunk更改為“主機的某些著陸墊”(請參見下文)。
  • 任務3將設備輸出從LLVM IR遷移到某種形式的MLIR。該項目是可選的,有關詳細信息,請參見“遷移設備LLVM IR”部分。

該項目優先考慮具有啟用了LHLO發射器的端到端可運行模型的優先級。這意味著以下按優先級排列的目標清單:

  • 使XLA / GPU可與LHLO發射器一起運行,而未修改現有的Thunk和發射器。
  • 視情況消除在LHLO中對HloInstruction *的引用:
    • 將傳統發射器切換為基於MLIR的發射器(例如Linalg),或
    • 機械轉換現有發射器以採用MLIR表示(使用GPU Dialect遷移到Standard)。

遷移暴徒(任務2)

xla :: gpu :: Thunk是以下數據結構:

  • 可以從主機調用(xla :: gpu :: Thunk :: ExecuteOnStream())。
  • 在其子類中攜帶各種數據。
  • 與BufferAllocation :: Slice和StreamExecutor交互。
  • 啟動內核
  • 調用所有運行時庫。

費用包括:

  • 表示特定於操作的配置數據(例如卷積配置)。
  • 遷移op形狀和操作數形狀。
  • 代表一堆笨拙的樹(while,condition等)。

遷移工作與LHLO /發射器遷移無關。在資源有限的情況下,它優先於LHLO /發射器遷移。

關於如何從LHLO降低主機端部分,我們有幾種選擇:

  • TFRT
    • (專業版)很棒的CUDA和HIP包裝器。
    • (Pro)易於實現的庫調用(cuDNN,cuBLAS,cuFFT等),因為TFRT操作由C ++代碼解釋。
    • (Con)主機端正在開發中,未經測試。
  • 中央處理器代碼
    • (親)極大的降低能力。創建一些循環和條件,即可完成。
    • (缺點)GPUDialect尚未對鏈/流/異步/設備分配建模。
    • (缺點)CUDA / HIP運行時支持最少(工具包路徑,版本,動態加載等)。
  • 現有(解釋)XLA運行時

決策:採用TFRT,但也支持在TFRT中壓縮CPU代碼。

遷移設備LLVM IR(任務3)

元素發射器通過逐個元素填充來生成目標op。每個輸出元素都取決於操作數中的一組元素。通過組合緩衝區和動態索引來描述所有元素。幾乎描述所有“數學”操作就足夠了,但是出於性能原因,直接在(Cpu | Gpu)ElementalIrEmitter中直接實現了很大一部分“數學”操作。

ElementalIrEmitter的獨特之處在於:

  • XLA / GPU和CPU之間共享大部分代碼。
  • 它代表了模型中看到的大部分操作,包括所有基於元素的操作。
  • 大多數融合僅取決於ElementalIrEmitter。
  • 它在結構上很簡單,因為它描述了op元素和操作數元素之間的數據依賴DAG。
  • 它主要是便攜式的和高級的(例如,與GPU kReduce和GPU kCopy不同)。
  • 動態形狀支持至少對於元素級操作很容易。

現在,對於所有操作(無論是否元素發射),每個XLA操作的結束狀態都有幾種形式:

  1. 設備代碼保持為LLVM IR。
  2. 重構舊的發射器,使其類似於LHLO-> MLIR LLVM方言:
    • (成本)如果我們最終要遷移到標準版,那將是一項繁瑣的工作。
    • (優點)這很簡單而且機械。可以在短期內完成。
    • (優點)與(1)相比,它沒有更多的好處。
  3. 重構舊的發射器,使其像LHLO-> MLIR GPU + Standard + Loops:
    • (成本)將現有的排放者提升為標準排放帶來了一些挑戰。指針和GEP需要轉換為MemRef和SubView。確保amdgpu完整性是另一個問題。
    • (成本)XLA / GPU嚴重依賴LLVM元數據:
      • 塊/線程索引的range
      • aligndereferenceableinvariant.loadalias.scopenoalias用於加載/存儲。
      • llvm.loop.unroll.disablellvm.loop.unroll.fullllvm.loop.vectorize.enable用於順序循環。
    • (好處)可以是長期的。更便攜。
  4. 將舊的發射器重構為LHLO-> Linalg,並編寫新的Linalg發射器
    • (成本)這視情況而定。與以前的選項相比,匹配XLA性能的新實現需要通過基準<->優化工作流,這對於某些操作而言可能是一筆可觀的成本。
    • (優點)統一堆棧;社區支持;可移植性;更多優化潛力。

結論:

  • 不要去(2)。 (1)或(3)比(2)好。 (2)的成本要高於(1),因為它需要大量的機械重構。使用(1),我們仍然可以實現使XLA能夠拾取MLIR發射器的目標。這是通過執行LHLO-> LLVM IR->運行舊設備發射器來完成的。
  • ElementalIrEmitter ops適用於(4),但不能遞增。無法通過op進行操作,因為所有元素發射的op都連接到同一圖形中。這項工作還可以用作多個持續使用的力量(xla / service / mlir_gpu,內核生成器Linalg)的統一點。
  • 其他所有操作都用於(1)。作為擴展目標,它們可能會遷移到(3)或(4)。

優先次序

雖然上述所有三個任務都是可並行化的,但是在有限的資源下它們必須被序列化。優先級重點在於完成每個任務的可見結果。

優先級為:任務1(舊發射器的LHLO)>任務2(Thunk)>任務3(MLIR發射器)。

在任務1結束時,XLA的用戶可以生成LHLO(例如內核生成器)並執行它們。編譯格式將不是可序列化的MLIR。

在任務2結束時,LHLO降低為適當的可序列化MLIR。這樣可以進行離線編譯。

在任務3結束時,所有XLA發射器的實現均基於MLIR。

詳細設計

步驟1 :(任務1)完成LHLO並讓舊式發射器接受LHLO

此步驟使所有現有的XLA / GPU發射器都與MLIR操作交互。此步驟是純重構和NFC。

此步驟主要是機械步驟,但值得注意的是,未嵌套的HloComputation和LHLO之間存在以下差異:

  • 每個Hlo指令都可以直接訪問其操作數(數據流DAG)。相反,每個LHLO op只能訪問其操作數緩衝區(ops和緩衝區之間的二分之一)。 LHLO操作必須通過use-def鏈來訪問其操作數操作。
  • 根據經驗,未嵌套的舊發射器幾乎永遠不會訪問其操作數。唯一的例外是kReduce。
  • 未嵌套的舊發射器僅用於獲取切片而訪問BufferAssignment,而不用於訪問諸如dataflow_analysis()或alias_analysis()之類的輔助數據結構。 llvm_ir基於切片信息構建自己的alias_analysis()。

結論是LHLO應該在沒有重大麻煩的情況下就可以正確地入內。

步驟2 :(可選)分析支持

僅當我們開始丟棄某些XLA Thunk邏輯時才需要執行此步驟(請參閱下一步)。

在實際打開任何基於MLIR的發射器之前,我們需要對基於MLIR的發射器進行性能分析。

當前,XLA通過調用StreamExecutor的計時器執行其自身的性能分析。底層計時器在內核啟動之前和之後插入兩個事件,並測量這兩個事件之間的同步時間。

支持MLIR中的概要分析的方法大致有三種:

  • 端到端運行探查器
  • 使用注入的探查器為LHLO中的每個op添加一個探查op。

“端到端”方法對MLIR是透明的,但存在使XLA最初不使用它的問題:探查器(nvprof / ...)收集的庫調用無法輕鬆地與HLO相關聯行動。例如,cuDNN為每個HLO啟動多個內核,很難區分哪個內核對應於哪個HLO。

“注入的探查器”方法要求:

  • LHLO以探查器作為參數。
  • 在每個操作之前和之後插入profile.start / profile.end。
  • 從較低的配置文件開始。{start,end}到C ++實現。

對於MLIR生成的操作,無法輕鬆進行準確的性能分析,因為:

  • MLIR沒有計時器,也不取決於TFRT / StreamExecutor。
  • MLIR不會輕易調用具有復雜參數的C函數。

步驟3 :(任務2)遷移暴徒

請注意,大約有三種類型的重擊:

  • KernelThunk,它啟動一個內核。
  • 控制流重擊,具有主機控制流邏輯(有條件的,同時的,順序的)和啟動主體內核。
  • 庫重擊:cuDNN,cuBLAS,cuFFT,NCCL等

該計劃是:

  • 使Thunks(反序列化)。
  • 幫助將TFRT提高到可以支持這些語義的狀態。
  • 隨著狀態的改善,逐步遷移單個thunk。

這些操作項僅部分訂購。實際的執行順序/工程並行性將隨其進行評估。

步驟4 :(任務3)遷移ElementalIrEmitter

分析完成後,我們就可以完成並調整MLIR中所有基於ElementalIrEmitter的發射器。然後,假設所有這些基於MLIR的發射器都使用單個流,則默認情況下將它們打開。

注意,遷移XLA / CPU的ElementalIrEmitter也是有益的,因為它們共享了大部分代碼。

完成所有基準測試和性能搜尋(TODO:定義性能奇偶校驗)後,我們打開新的基於MLIR的元素發射器,並刪除舊的ElementalIrEmitter。

此步驟還為以後的遷移提供了簡單的融合過渡(嵌套操作)。

第5步:多流支持或刪除

在MLIR中支持它之前,我們不能刪除某些發射器 ,否則我們將其刪除。 MLIR的工作量相對較大,而XLA的工作量則較小。我們應該調查多流XLA / GPU用戶的當前用戶,並在合理的情況下嘗試刪除此功能。

步驟6 :(任務3)遷移設備操作

此步驟將遷移所有未嵌套的操作,然後我們可以刪除所有未嵌套的發射器。

這要求對kCopy和kReduce進行重寫/重構。 kReduce已經進行了大量工作,因此需要完成的實際工作量尚待觀察。