为 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 字面量数据。换句话说,它有助于封装从主机到设备之间的数据传输。