XLA 架构

XLA(加速线性代数)是一种可优化线性代数的机器学习 (ML) 编译器,能够提高执行速度和内存用量。本页将简要介绍 XLA 编译器的目标和架构。

目标

如今,XLA 支持多种机器学习框架前端(包括 PyTorch、TensorFlow 和 JAX),并且是 OpenXLA 项目的一部分,该项目是一个由领先的机器学习硬件和软件组织合作开发的开源编译器技术生态系统。在创建 OpenXLA 项目之前,XLA 是在 TensorFlow 项目中开发的,但基本目标保持不变:

  • 提高执行速度。编译子图以缩短短期有效的操作的执行时间并消除运行时的开销,融合流水线操作以减少内存开销,并专门处理已知的张量形状以实现更积极的常量传播。

  • 提高内存用量。分析和调度内存用量,消除许多中间存储缓冲区。

  • 降低对自定义操作的依赖。通过提升自动融合的低级别操作的性能,使其与最初手动融合的自定义操作的性能相匹配,从而消除对许多自定义操作的需求。

  • 提高可移植性。为新硬件编写新后端相对容易,以便大部分机器学习模型无需修改即可在该硬件上运行。这与为新硬件专门编写单个单体式运算的方法不同,后者需要重写模型才能利用这些运算。

运作方式

XLA 编译器从 StableHLO 中定义的机器学习框架获取模型图,并将它们编译成适用于各种架构的机器指令。StableHLO 定义了一个版本化操作集(HLO = 高级操作),可在机器学习框架和编译器之间提供可移植层。

通常,将模型图转换为针对目标优化的可执行文件的编译过程包括以下步骤:

  1. XLA 对独立于目标的 SttableHLO 图执行多次内置优化和分析传递,例如用于为计算分配运行时内存的 CSE、独立于目标的操作融合和缓冲区分析。在此优化阶段,XLA 还会将 StableHLO 方言转换为内部 HLO 方言。

  2. XLA 将 HLO 计算发送到后端以进行进一步的 HLO 级优化,这次要注意特定于目标的信息和需求。例如,GPU 后端可以执行对 GPU 编程模型有益的运算融合,并确定如何将计算划分到数据流中。在此阶段,后端还可以通过模式匹配某些操作或其组合与优化的库调用。

  3. 然后,后端会执行特定于目标的代码生成操作。XLA 中包含的 CPU 和 GPU 后端使用 LLVM 进行低层级 IR、优化和代码生成。这些后端发出以高效方式表示 HLO 计算所需的 LLVM IR,然后调用 LLVM 从此 LLVM IR 中发出原生代码。

在此过程中,XA 编译器是模块化的,因为很容易插入到替代后端中,以以某些新颖的硬件架构为目标。GPU 后端目前通过 LLVM NVPTX 后端支持 NVIDIA GPU。CPU 后端支持多个 CPU ISA。