XLA简介

XLA(Accelerated Linear Algebra)-加速线性代数，是Google推出的高性能机器学习领域编译器，它可以在不更改源代码的条件下加速Tensorflow模型。

TensorFlow在设计时主要考虑灵活性和可拓展性，而灵活性往往和高性能不可兼得，因此TensorFlow的计算性能有所欠缺。XLA就是为了解决这个问题而提出的。

XLA为什么有效

1. 融合可组合算子从而提高性能
   XLA通过对TensorFlow运行时的计算图进行分析，将多个低级算子进行融合从而生成高效的机器码。例如：
   
   因为上面计算图中的许多算子都是逐元素(element-wise)的计算，所以可以融合为一个element-wise的循环计算kernel中。matmul的结果加biases时为element-wise；然后对Add的结果每个进行Relu；再对Relu的结果每个元素取幂运算。
   通过将这些算子融合，可以减少申请这些算子间的中间结果所占用的内存。同时因为将多个kernel融合为了一个kernel，因此减少了加载kernel的时间消耗。
   因此，XLA对模型的性能提升主要来自于将多个连续的element-wise算子融合为一个算子。

2. 提高内存利用率
   通过对内存使用的分析与规划，原则上可以消除许多中间数据的内存占用。

3. 减少模型可执行文件大小
   对于移动设备场景，XLA可以减少TensorFlow模型的执行文件大小。通过AOT（ahead-of-time）编译将整个计算图生成轻量级的机器码，这些机器码实现了计算图中的各个操作。
   在模型运行时，并不需要完整的TensorFlow运行时环境，因此只有用户关心的计算图实际执行的操作被转换编译为设备代码。

4. 减少custom算子依赖
   对于softmax算子，可以在一个kernel中实现，也可以通过多个简单算子组合实现：
   softmax = exp(logits) / reduce_sum(exp(logits), dim)
   通过使用简单算子拼接复杂算子，可以大大减少对custom算子的需求。但是多个简单算子拼接往往需要额外存储多余的中间结果以及加载多个kernel的消耗，因此性能较低。
   而XLA可以将多个简单算子融合为一个算子，从而解决上述问题。

5. 方便支持不同硬件后端
   传统的在TensorFlow中支持一种新的设备需要将所有的kernels（ops）再实现一遍，这无疑需要非常巨大的工作量。而通过XLA为TensorFlow新增硬件后端只需要很小的工作量。因为XLA的算子都是操作原语（低级算子），因此数量少且容易实现，XLA会自动地将TensorFLow计算图中复杂的算子拆解为 primitive算子。

XLA如何工作

XLA的输入语言是"HLO IR"，也可以称为HLO(High Level Optimizer)。XLA将HLO描述的计算图（计算流程）编译为针对各种特定后端的机器指令。

首先，XLA对输入的HLO计算图进行与目标设备无关的优化，如CSE，算子融合，运行时内存分配分析。输出为优化后的HLO计算图。
然后，将HLO计算图发送到后端（Backend），后端结合特定的硬件属性对HLO计算图进行进一步的HLO级优化，例如将某些操作或其组合进行模式匹配从而优化计算库调用。
最后，后端将HLO IR转化为LLVM IR，LLVM再进行低级优化并生成机器码。

参考资料：
[1] https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html
[2] https://www.tensorflow.org/xla/architecture