为 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:没有现有 LLVM 后端的非 CPU 类硬件

如果无法使用 LLVM,那么最好的选择就是为所需的硬件实现新的 XLA 后端。此选项需要的工作量最大。需要实现的类如下所示:

  • StreamExecutor:对于许多设备来说,并非需要 StreamExecutor 的所有方法。如需了解详情,请参阅现有的 StreamExecutor 实现。
  • xla::Compiler:此类将 HLO 计算的编译封装到 xla::Executable 中。
  • xla::Executable:此类用于在平台上启动编译后的计算。
  • xla::TransferManager:此类使后端能够提供特定于平台的机制,用于根据给定的设备内存句柄构建 XLA 字面量数据。换句话说,它有助于封装数据从主机到设备之间的传输。