AI编译器XLA调研
文章目录
- 一、XLA简介
- 二、XLA在TensorFlow中的应用
-
- 2.1 XLA是什么？（tensorflow\compiler\xla）
- 2.2 TensorFlow怎样转化为XLA （tensorflow\compiler\tf2xla）
- 2.3 JIT(just in time) 即时编译 （tensorflow\compiler\jit）
-
- 2.3.1开启jit编译
- 2.3.2自动聚类
- 2.3.3jit实例
- 2.4 AOT (Ahead-of-time compilation)超前编译
- 三、总结
一、XLA简介
XLA(Accelerated Linear Algebra)-加速线性代数，是Google推出的高性能机器学习领域编译器，它可以在不更改源代码的条件下加速Tensorflow模型。TensorFlow在设计时主要考虑灵活性和可拓展性，而灵活性往往和高性能不可兼得，因此TensorFlow的计算性能有所欠缺。而机器学习的运算中99%都是向量乘以矩阵、矩阵乘以矩阵的计算，XLA是专门用来优化这些计算的。
举个例子，运行在GPU上的model_fn函数会顺序调用multiply、add和reduce_sum这三个op，而且multiply，也就是y * z的计算结果会先从GPU拷贝回host，再拷贝到device作为add的input，同样的，add的计算结果也会以相同的方式传递给下一个op。
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
显然，对于整个函数来说，将中间变量在host和device间来回倒腾是没有意义的。因此，如果把函数看作一个op，那在计算中产生的中间结果就不必返回到host，少了数据传输的时间开销，就可以大幅提升运算效率。
这种将多个op融合成一个op的方法就称为fuse，当前fuse的技术路线有：
- 通过手写或codegen工具来开发fused op，例如在上述例子中就可以开发tf.fused_reduce_sum(x, y, z)。它的优点是代码可控性高，易于性能优化，但缺点是程序缺乏灵活性。像Pytorch这种动态图的框架走的就是这条路线，Nvidia的Apex提供有大量fused kernel。
- 通过XLA等AI编译器将python函数编译成fused op。这样做的好处是灵活性强，可以fuse任何计算，弊端则是开发难度大，且性能通常会逊色于手写或codegen kernel。
编译器一般构成
XLA是AI编译器，只要是编译器就逃脱不了几大构成：
传统的编译器通常分为三个部分，前端（frontEnd），优化器（Optimizer）和后端（backEnd）. 在编译过程中，编译器一般构成
XLA是AI编译器，只要是编译器就逃脱不了几大构成：
传统的编译器通常分为三个部分，前端（frontEnd），优化器（Optimizer）和后端（backEnd）. 在编译过程中，前端主要负责词法和语法分析，将源代码转化为抽象语法树；优化器则是在前端的基础上，对得到的中间代码进行优化，使代码更加高效；后端则是将已经优化的中间代码转化为针对各自平台的机器代码。
二、XLA在TensorFlow中的应用
2.1 XLA是什么？（tensorflow\compiler\xla）
现代编译器核心离不开IR（Intermediate Representation），编译过程是将高级语言不断进行变换最终变成机器可以运行的机器指令，整个过程是一个复杂的过程。为了简化编译器架构，复用核心算法，大部分编译器都引入了IR。
上图（图1）演示了XLA Graph 到最终的机器代码生成。
XLA Graph :节点和边缘组成，其中节点是线性代数操作，边缘是数据流，比如tensors。
XLA Graph从线性代数降低到编译的形式，到达LLVM IR ,这种中间表示形态。
然后进行code generator生成实际的机器平台代码。
这张图最重要的是我们将信息输入到XLA图表，然后得到我们的机器代码。
2.2 TensorFlow怎样转化为XLA （tensorflow\compiler\tf2xla）
根据源代码的位置和标题，顾名思义，就是从TensorFlow到XLA。
首先看一下TensorFlow是什么。从下图（图2）可以看出，首先，我们用不同的编程语言来构建TensorFlow项目，python是用的最多的，但是其他语言也是支持的(c++需要先编译)。然后，我们在构建程序的时候，就是要构建图中中间蓝色线框的图表，也就是TensorFlow Graph，所有的tensor都会经过它。最后，对于这张图表，会有底层的C++实际执行的运行时间。如红色框图所示，里边有两个部件，一个是执行器，另一个是内核。执行器的目的是执行TensorFlow Graph，当他遇到每个节点时候，它会查询内核然后找到合适的内核并运行。
例如，如下图（图3）所示。遇到add 、softmax节点，查询内核然后运行。这就是TensorFlow的运行方式。并不需要XLA。
接下来将了解TensorFlow如何以XLA程序运行（怎么转化？）。如下图（图4）所示，红色线框表示我们的TensorFlow Graph，我们的目的是生成蓝色线框的XLA Graph。正如第一张图（图1）所示，只要我们将信息转化为XLA Graph 这种表达方式，XLA就会处理剩下的事情，将信息编译成机器代码。在图中（图4）有一个局部执行器，和普通的TensorFlow runtime中的是一样的。但是在这里边不使用TensorFlow的内核，而是换成了XLA的核心程序（粉色线框）。这个核心程序的工作不是执行操作，而是创建图表，将TensorFlow Graph 转换为 XLA Graph。（/tf2xla/kernels)
转换过程如下图所示（图5），例如，在TensorFlow Graph中，遇到了add节点，然后查询核心转换程序，可以将TensorFlow add 转换为 XLA add，下图展示的是比较简单的一对一呈现。
而对于比较复杂的运算，比如softmax操作，内部存在一系列的运算，查看转换核心程序之后，将其转换为XLA图表中的多重运算，如下图（图6）所示。
那么问题就来了，我们应该怎么在TensorFlow中使用XLA进行加速呢？
2.3 JIT(just in time) 即时编译 （tensorflow\compiler\jit）
个人理解就是，在TensorFlow Graph中，将可以转换的节点进行转换，对于可以排序的多个节点，识别定义为cluster(簇、集群），然后进行编译和运行。
如下图（图7）所示，梯形框中的三个节点都可以被转换为XLA，经过转换之后，就变成了右边的XLA Graph，未经过转换的node（节点）按照TensorFlow方式正常运算，而可以转换的簇经过转换之后（变成了XLA graph）可以直接编译和运行，变成实际的机器代码。并且，以后遇到相同类型的，不需要重复编译，就可以一次一次地执行。
同样的，如果存在多个clusters。如下图（图8）所示，存在多个集群，当TensorFlow执行器运行时，识别出来该集群可以进行转换就会转换为xla graph，然后对下图中的多个集群进行编译并且执行。
2.3.1开启jit编译
第一种方式：全局模式
首先读取整个TensorFlow Graph,然后决定图表中的哪些集群实际上可以编译，最后重写Graph。当具体执行时候，每遇到节点，就会启动JIT，编译成机器代码。自动进行。
第二种方式：手动范围
可以手动控制设置明确什么被编译，什么不被编译。如下图所示，仅仅add 算子被编译。
2.3.2自动聚类
若要在 TensorFlow 模型中开始使用 XLA 并且无需做出任何更改，最简单的方法是启用“自动聚类”，此功能会自动在可以使用 XLA 进行编译和执行的 TensorFlow 函数中查找聚类（连通的子图）。您可以通过设置 TF_XLA_FLAGS 环境变量，在 GPU 上启用自动聚类功能：
$ TF_XLA_FLAGS=–tf_xla_auto_jit=2 path/to/your/tf/program
自动聚类目前已针对 GPU 工作负载进行了优化，但您也可以通过另外使用 --tf_xla_cpu_global_jit 标记在 CPU 上启用它：
$ TF_XLA_FLAGS=“–tf_xla_auto_jit=2 --tf_xla_cpu_global_jit” path/to/your/program
但是自动聚类的方式容易出现波动，不一定出现好的优化效果，建议使用更加清晰的模式，@tf.function(jit_compile=True)。
TensorFlow中大约1500多个OP，不是每个op都可以被编译，当遇到不可编译的OP时，会报错并建议重构代码。
其次，编译要求是静态的，如果形状变动频繁，会出现重新编译的情况，带来额外的延时并导致并不能优化速度。如下图所示，
需要重新编译。此外，对于一些参数，也需要保持固定，不然也会重新编译。
2.3.3jit实例
import tensorflow as tf
import time
@tf.function(jit_compile=True)
def running_example(x,y):
return tf.reduce_mean(tf.multiply(x**2,3)+y)
x = tf.random.uniform((15000,15000))
y = tf.random.uniform((15000,15000))
options = tf.profiler.experimental.ProfilerOptions(host_tracer_level = 2,python_tracer_level = 0,device_tracer_level = 1)
with tf.profiler.experimental.Profile('logdir_jit',options=options):
print(running_example(x,y))
不加入jit的timeline。从图中可以看出有四个算子，总时间为520ms。
加入jit之后，只有一个总得计算过程，时间消耗 250ms。 大大减小了时间消耗。
2.4 AOT (Ahead-of-time compilation)超前编译
和JIT不同的是，在JIT中，当遇到没法编译的node时，自动按照TensorFlow runtime的方式正常运行，仅仅编译可以被编译的node即可。
AOT则需要将整个TensorFlow Graph 编译成为XLA Graph。进入执行阶段后，不再有编译过程发生。好处是可以直接在手机端、服务器端运行。
具体实现方式：通过tfcompile来实现。类似一个包装器。
如下图所示，给定一个TensorFlow Graph，然后给它一个配置文件，告诉它需要提供（输入）和获取（输出）什么。然后把以上东西送入tfcompile，从另一端输出机器平台运行代码。整个编译相当于一个函数，知道输入和输出。
三、总结
XLA与TensorFlow合作有几个目标：
提高执行速度。编译子图以减少短暂Ops的执行时间，以消除TensorFlow运行时间的开销，融合流水线操作以减少内存开销，并专用于已知张量形状以允许更积极的恒定传播。
改善内存使用。分析和计划内存使用情况，原则上消除许多中间存储缓冲区。
减少对自定义操作的依赖。通过改进自动融合低级Ops的性能，消除对许多自定义Ops的需求，匹配手工融合的自定义Ops的性能。
减少移动足迹。通过提前编译子图。发出可以直接链接到另一个应用程序的对象/头文件对来消除TensorFlow运行时。结果可以将移动推断的占用空间。减少几个数量级。
提高可移植性。为新颖的硬件编写新的后端程序相对容易，此时大部分TensorFlow程序。将在该硬件上未修改地运行。与专门针对新硬件的个体单片Ops的方法形成对比，需要重写TensorFlow程序以利用这些Ops。