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