開發 XLA 的新後端

本指南適用於希望 XLA 輸出能夠有效率地指定硬體的程式。本指南並非逐步說明,且假設您具備 LLVMBazel 和 XLA 的知識。

XLA 提供抽象介面,可導入新的架構或加速器來建立後端,以執行 XLA 輸出的機器學習程式。與從前端架構 (例如 PyTorch 或 TensorFlow) 實作新硬體的現有作業相比,重新指定 XLA 應該更為簡單且擴充性。

大部分的導入都可分為以下其中一種情況:

  1. XLA 尚未正式支援現有的 CPU 架構,且包含或沒有現有的 LLVM 後端。
  2. 具有現有 LLVM 後端且不具 CPU 特性的硬體。
  3. 非具有 CPU 特性的硬體,沒有現有的 LLVM 後端。

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

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

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

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

情境 2:已有 LLVM 後端且與 CPU 無關的硬體

您可以在現有的 xla::CPUCompilerxla::GPUCompiler 類別上建立新的 xla::Compiler 實作模型,因為這些類別已發出 LLVM IR。視硬體的性質而定,LLVM IR 產生的許多方面可能都必須變更,不過可與現有的後端共用大量程式碼。

常見的例子是 XLA 的 GPU 後端。GPU 後端指定的是類似 CPU 的 ISA,因此產生程式碼的某些面向是 GPU 網域獨有的。其他類型的硬體,例如 Hexagon 等 DSP (具有上游 LLVM 後端) 可以重複使用部分 LLVM IR 排放邏輯,但其他部分會獨一無二。

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

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

  • StreamExecutor:許多裝置都需要使用部分 StreamExecutor 方法。詳情請參閱現有的 StreamExecutor 實作。
  • xla::Compiler:此類別會將 HLO 運算的編譯封裝封裝至 xla::Executable
  • xla::Executable:此類別用於在平台上啟動已編譯的運算。
  • xla::TransferManager:這個類別可讓後端提供平台專屬的機制,以從特定裝置記憶體控制代碼建構 XLA 常值資料。換句話說,這有助於封裝從主機至裝置轉移的資料。