本指南適用於希望 XLA 輸出以高效率目標硬體為目標的系統工程師。本指南並非逐步說明,而是假設您熟悉 LLVM、Bazel 和 XLA。
XLA 提供抽象介面,其中新的架構或加速器可以實作這個介面,建立後端來執行 XLA 輸出的機器學習程式。重新指定 XLA 應比針對新硬體導入前端架構 (例如 PyTorch 或 TensorFlow) 的所有現有運算,輕鬆且可擴充。
大多數的導入作業都屬於下列其中一種情況:
- 現有 CPU 架構尚未正式支援 XLA,不論現有或沒有現有的 LLVM 後端。
- 非 CPU 類的硬體,搭配現有的 LLVM 後端。
- 沒有現有 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::CPUCompiler
和 xla::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 常值資料。換句話說,它有助於封裝將資料從主機傳輸到裝置,以及傳輸到裝置之間的傳輸作業。