Developing a new backend for XLA（为XLA开发一个新的后端）

为XLA开发新的后端

本初步指南适用于希望以有效方式轻松将TensorFlow重新定位到其硬件的早期使用者。本指南并非一步一步地掌握LLVM，Bazel和TensorFlow的知识。

XLA提供了一个抽象接口，新体系结构或加速器可以实现创建后端以运行TensorFlow图形。重新定位XLA应该比实现每个现有的TensorFlow Op用于新硬件更加简单和可扩展。

大多数实现将落入以下情况之一：

1. 现有的CPU体系结构尚未正式由XLA支持，无论是否存在LLVM后端。

2. 具有现有LLVM后端的非CPU类硬件。

3. 没有现有LLVM后端的非CPU类硬件。

注意： LLVM后端可以是官方发布的LLVM后端或内部开发的定制LLVM后端。

情况1：XLA尚未正式支持现有CPU架构

在这种情况下，首先查看现有的XLA CPU后端。通过使用LLVM，XLA可以轻松地将TensorFlow重定向到不同的CPU，因为XLA后端对于CPU的主要区别在于LLVM生成的代码。Google测试XLA for x64和ARM64体系结构。

如果硬件供应商为其硬件提供LLVM后端，则将后端与使用XLA构建的LLVM进行链接很简单。在JIT模式下，XLA CPU后端为主机CPU发出代码。对于提前编译，xla::AotCompilationOptions可以提供一个LLVM三元组来配置目标体系结构。

如果没有现有的LLVM后端，但存在另一种代码生成器，则应该可以重新使用大部分现有的CPU后端。

场景2：具有现有LLVM后端的非CPU类硬件

可以xla::Compiler在现有类xla::CPUCompiler和xla::GPUCompiler类上建立一个新的实现，因为它们已经发出了LLVM IR。根据硬件的性质，许多LLVM IR生成方面可能需要更改，但可以与现有后端共享大量代码。

一个很好的例子就是XLA 的GPU后端。GPU后端以非CPU类ISA为目标，因此其代码生成的某些方面对于GPU域是唯一的。其他类型的硬件，例如Hexagon（具有上游LLVM后端）的DSP可以重新使用部分LLVM IR发射逻辑，但其他部分将是唯一的。

场景3：没有现有LLVM后端的非CPU类硬件

如果无法使用LLVM，那么最好的选择是为XLA实现所需硬件的新后端。这个选项需要最多的努力。需要实施的类如下：

  • StreamExecutor：对于许多设备，并非所有的方法StreamExecutor都是必需的。详情请参阅现有的StreamExecutor实施。

  • xla :: Compiler：这个类将HLO计算的编译封装为一个xla::Executable。

  • xla::Executable：该类用于在平台上启动编译的计算。

  • xla::TransferManager：该类使后端能够提供特定于平台的机制，用于从给定的设备内存句柄构造XLA文字数据。换句话说，它有助于封装从主机到设备的数据传输并返回。