Developing a new backend for XLA(为XLA开发一个新的后端)

本初步指南适用于希望以有效方式轻松将TensorFlow重新定位到其硬件的早期使用者。本指南并非一步一步地掌握LLVMBazel和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::CPUCompilerxla::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::TransferManager:该类使后端能够提供特定于平台的机制,用于从给定的设备内存句柄构造XLA文字数据。换句话说,它有助于封装从主机到设备的数据传输并返回。

扫码关注云+社区

领取腾讯云代金券