TensorFlow XLA 初探

在此分享一下个人在学习XLA时的一点理解和心得~

XLA（加速线性代数）是一种针对特定领域的线性代数编译器，能够加快 TensorFlow 模型的运行速度，而且可能完全不需要更改源代码。它可以提高运行速度并改进内存用量。

什么是XLA？[1]

XLA全称Accelerated Linear Algebra（加速线性代数），是在特定领域下的线性代数编译器。

TensorFlow生态中包含许多的编译器和优化器，使得代码可以在不同的软硬件堆栈上运行。

TensorFlow MLIR 的一些组件

上图[2]是一个概览，其中XLA本质上是一种深度学习编译器，作为编译器，XLA负责对Tensorflow前端定义的计算图进行优化：将图转化为 XLA 高级优化器 (XLA HLO) 表示，反之，这种表示亦可调用适合 CPU 或 GPU 的 LLVM 编辑器，或者继续使用适合 TPU 的 XLA。（或者将二者结合！）

有着以下的几个目的:

1. 提高执行速度：编译子图（本质上是将多个短期Operator进行融合），缩短执行时间
2. 提高内存使用率：分析内存的使用。内存带宽通常是硬件加速器上最稀缺的资源，因此消除内存操作是提高性能的最佳方法之一。（参考下面的例子）
3. 减少对自定义Operator的依赖：这点本质上是和第一点相近的。因为XLA内嵌了对低级Operators的自动融合处理，因此不需要用户自定义新的将这些Operators合起来的Operator。减少对自定义的依赖。
4. 减少移动端的库的大小
5. 提高可移植性：作为一个中间层，为新硬件编写新的后端相对容易。类似起到IR的作用，提供每一个Operator在针对硬件的实现。

加速的大概原理：

输入：TensorFlow计算图

输出：为给定模型生成计算内核（交由LLVM IR，最后是后端硬件计算）

E.g. 乘法、加法、加法

def model_fn(x, y, z):
  return tf.reduce_sum(x+y*z)

TF对应的计算流图

XLA会将上述的三次计算融合到一个子图中，只需要启动一次GPU内核就能得到计算结果；其次这种融合操作不会将 y∗z 和 x+y∗z 的中间值写到内存中。

如何实现的加速/优化？

几个基本概念

• TensorFlow IR，代表 TensorFlow 图中可能存在的一切
• XLA HLO IR，旨在利用 XLA 的编译功能（输出到 TPU 等）
• LLVM IR，与 LLVM 自我表示之间存在 1:1 映射，可使 MLIR 通过 LLVM 发出 GPU 与 CPU 代码
• TensorFlow Lite，将会转换以在移动平台上运行代码

XLA的输入语言为 HLO IR 或者简称 HLO。本身HLO的表示，就可以被认为是一种编译器的IR。HLO的语义官网上有详细的枚举，其实就是HLO instructions。

个人理解，不一定正确：可以将TensorFlow的计算流图认为是前端，XLA是中端，其中XLA的后端对接了LLVM，并在不同的硬件上能够迁移。

基于抽象的HLO语义，可以做编译的优化，主要是以下几个步骤：

1. XLA有几个与目标硬件无关的优化，过程与硬件后端无关的（target-independent）算子融合，以及用于为计算而在运行时做内存分配的缓冲区分析。
2. 硬件目标无关的步骤之后，XLA将HLO计算发送到后端去做进一步的HLO级优化，HLO级优化会针对硬件目标特定信息和需求，如XLA GPU后端的算子融合、确定如何将计算划分为流。此阶段，后端还可以将某些算子或其组合与优化的库调用进行模式匹配。
3. 下一步是特定于目标硬件的代码生成。CPU包括XLA使用的GPU后端会使用LLVM生成低级IR、优化和代码生成。硬件后端以有效方式发出了表示XLA HLO计算所必需的LLVM IR，然后调用LLVM从LLVM IR发出原生代码。

XLA的编译过程

目前已经支持的后端有 x64和ARM64的CPU后端 以及 NVIDIA GPU后端，TPU后端应该也支持了。

目前XLA的功能主要是体现在两个方面：

即时编译（JIT，Just-In-Time）：建议大家可以参考下如下的链接，提供了一个直观的例子表示如何使用XLA优化TF中的图。

编译流程：TensorFlow Graph （tf2xla）---> XLA HLO Graph ---> Optimized XLA HLO ---> LLVM IR （硬件执行)

注：这里很多时候，可以直接leverage LLVM社区的工作来完成LLVM IR到最终硬件可执行码的映射和优化工作。

姜曦楠：[腾讯机智]TensorFlow XLA工作原理124 赞同 · 10 评论文章正在上传…重新上传取消

超前编译（AOT，Ahead-Of-Time）：需要用到tfcompile，这是一个独立的工具，提前将TensorFlow图编译成可执行代码，可以避免运行时的开销。典型用力是将图编译成移动设备的可执行代码。

XLA HLO 的语义

在下坂奇：TensorFlow XLA HLO 语义5 赞同 · 0 评论文章

参考

1. ^TensorFlow 官网 https://www.tensorflow.org/xla?hl=zh-cn
2. ^MLIR MLIR 发布：全新的中介码与编译器框架 - 程序员大本营