AI编译器XLA调研

大家好，欢迎来到IT知识分享网。

一、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在TensorFlow中的应用

2.1 XLA是什么？（tensorflow\compiler\xla）

现代编译器核心离不开IR（Intermediate Representation），编译过程是将高级语言不断进行变换最终变成机器可以运行的机器指令，整个过程是一个复杂的过程。为了简化编译器架构，复用核心算法，大部分编译器都引入了IR。

上图（图1）演示了XLA Graph 到最终的机器代码生成。

2.2 TensorFlow怎样转化为XLA （tensorflow\compiler\tf2xla）

根据源代码的位置和标题，顾名思义，就是从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(簇、集群），然后进行编译和运行。

2.3.1开启jit编译

第一种方式：全局模式

2.3.2自动聚类

若要在 TensorFlow 模型中开始使用 XLA 并且无需做出任何更改，最简单的方法是启用“自动聚类”，此功能会自动在可以使用 XLA 进行编译和执行的 TensorFlow 函数中查找聚类（连通的子图）。您可以通过设置 TF_XLA_FLAGS 环境变量，在 GPU 上启用自动聚类功能：

$ TF_XLA_FLAGS=“–tf_xla_auto_jit=2 –tf_xla_cpu_global_jit” path/to/your/program

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(x2,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)超前编译

三、总结

XLA与TensorFlow合作有几个目标：