開發 XLA 的新後端

本指南適用於希望 XLA 輸出以高效率目標硬體為目標的系統工程師。本指南並非逐步說明,而是假設您熟悉 LLVMBazel 和 XLA。

XLA 提供抽象介面,其中新的架構或加速器可以實作這個介面,建立後端來執行 XLA 輸出的機器學習程式。重新指定 XLA 應比針對新硬體導入前端架構 (例如 PyTorch 或 TensorFlow) 的所有現有運算,輕鬆且可擴充。

大多數的導入作業都屬於下列其中一種情況:

  1. 現有 CPU 架構尚未正式支援 XLA,不論現有或沒有現有的 LLVM 後端。
  2. 非 CPU 類的硬體,搭配現有的 LLVM 後端。
  3. 沒有現有 LLVM 後端的非 CPU 類硬體。

情境 1:XLA 尚未正式支援的現有 CPU 架構

在這個情境中,請先查看現有的 XLA CPU 後端。XLA 可使用 LLVM 輕鬆指定不同的 CPU,因為 CPU 的 XLA 後端主要差異在於 LLVM 產生的程式碼。

如果硬體廠商有硬體的 LLVM 後端,就能輕鬆連結後端與使用 XLA 建構的 LLVM。在 JIT 模式下,XLA CPU 後端會發出主機 CPU 的程式碼。對於提前編譯,xla::AotCompilationOptions 可提供 LLVM 三元來設定目標架構。

如果目前沒有 LLVM 後端,但已存在其他類型的程式碼產生器,則應重複使用大部分的現有 CPU 後端。

情境 2:與 CPU 無關的硬體,已具備 LLVM 後端

您可以在現有的 xla::CPUCompilerxla::GPUCompiler 類別上為新的 xla::Compiler 實作建模,因為這些類別已經發出 LLVM IR。視硬體的性質而定,LLVM IR 產生作業的許多部分都可能會變更,但許多程式碼仍可與現有的後端共用。

XLA 的 GPU 後端就是一個很好的範例。GPU 後端是以非 CPU 的 ISA 做為目標,因此其程式碼的某些部分專屬於 GPU 網域。其他類型的硬體 (例如 Hexagon (具有上游 LLVM 後端)) 可以重複使用 LLVM IR 排放邏輯的部分,但其他部分不得重複。

情境 3:類似非 CPU 的硬體,但沒有現有的 LLVM 後端

如果無法使用 LLVM,則最佳做法是為所需的硬體實作 XLA 的新後端。這個選項需要最為費心。需要實作的類別如下:

  • StreamExecutor:對許多裝置而言,不需要所有 StreamExecutor 方法。詳情請參閱現有的 StreamExecutor 實作。
  • xla::Compiler:這個類別會將 HLO 運算的編譯作業封裝為 xla::Executable
  • xla::Executable:這個類別可用於在平台上啟動編譯的運算。
  • xla::TransferManager:這個類別可讓後端提供平台專屬的機制,以便使用特定裝置記憶體控制代碼建構 XLA 常值資料。換句話說,它有助於封裝將資料從主機傳輸到裝置,以及傳輸到裝置之間的傳輸作業。