本指南适用于希望 XLA 高效输出针对其硬件的程序的系统工程师。本指南并非分步指南,假定您具备 LLVM、Bazel 和 XLA 方面的知识。
XLA 提供了一个抽象接口,新架构或加速器可以实现该接口来创建后端,以运行由 XLA 输出的机器学习程序。与为新硬件实现来自前端框架(如 PyTorch 或 TensorFlow)的每个现有操作相比,重访定位 XLA 应更简单、更可伸缩。
大多数实现都属于以下某种情况:
- XLA 尚未正式支持的现有 CPU 架构(无论是否具有 LLVM 后端)。
- 具有现有 LLVM 后端的非 CPU 型硬件。
- 类似于 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::CPUCompiler
和 xla::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 字面量数据。换句话说,它有助于封装数据从主机到设备之间的传输。