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

為XLA開發新的後端

本初步指南適用於希望以有效方式輕鬆將TensorFlow重新定位到其硬件的早期採用者。本指南不是分步介紹的,並且假定您具有LLVMBazel和TensorFlow的知識。

XLA提供了一個抽象接口,新架構或加速器可以實現該抽象接口來創建後端以運行TensorFlow圖。與為新硬件實現每個現有的TensorFlow Op相比,重新定向XLA應該明顯更簡單且可擴展。

大多數實現將屬於以下情況之一:

  1. 帶有或不帶有現有LLVM後端的XLA尚未正式支持現有的CPU體系結構。
  2. 具有現有LLVM後端的非CPU類硬件。
  3. 沒有現有LLVM後端的非類CPU硬件。

方案1:XLA尚未正式支持現有的CPU體系結構

在這種情況下, 請先查看現有的XLA CPU後端 。 XLA可以通過使用LLVM輕鬆將TensorFlow重定向到不同的CPU,因為針對CPU的XLA後端之間的主要區別是LLVM生成的代碼。 Google針對X64和ARM64體系結構測試了XLA。

如果硬件供應商為其硬件提供了LLVM後端,則很容易將後端與使用XLA構建的LLVM鏈接。在JIT模式下,XLA CPU後端為主機CPU發出代碼。對於提前編譯, xla::AotCompilationOptions可以提供LLVM三元組來配置目標體系結構。

如果沒有現有的LLVM後端,但是存在另一種代碼生成器,則應該可以重用大多數現有的CPU後端。

方案2:具有現有LLVM後端的非CPU類硬件

可以在現有xla::CPUCompilerxla::GPUCompiler類上為新的xla::Compiler實現xla::GPUCompiler ,因為它們已經發出了LLVM IR。根據硬件的性質,有可能必須更改許多LLVM IR生成方面,但是可以與現有後端共享許多代碼。

一個很好的例子是XLA的GPU後端 。 GPU後端針對非CPU的ISA,因此其代碼生成的某些方面對於GPU域而言是唯一的。其他種類的硬件,例如像Hexagon這樣的DSP(具有上游LLVM後端),可以重用LLVM IR發射邏輯的一部分,但是其他部分將是唯一的。

方案3:沒有現有LLVM後端的非CPU類硬件

如果無法利用LLVM,則最好的選擇是為所需的硬件為XLA實施新的後端。此選項需要最大的努力。需要實現的類如下:

  • StreamExecutor :對於許多設備,並不需要所有方法的StreamExecutor 。有關詳細信息,請參見現有的StreamExecutor實現。
  • xla::Compiler :此類將HLO計算的編譯封裝到xla::Executable
  • xla::Executable :此類用於在平台上啟動編譯的計算。
  • xla::TransferManager :此類使後端能夠提供特定於平台的機制,以從給定的設備內存句柄構造XLA文字數據。換句話說,它有助於封裝從主機到設備再到設備的數據傳輸。