谷歌官博详解XLA：可在保留TF灵活性的同时提升效率
选自Google Blog
机器之心编译
参与：Jane W、吴攀
近日，谷歌开发者博客发布了一篇文章，介绍了用于 TensorFlow 的编译器 XLA（Accelerated Linear Algebra/加速线性代数）的原理和能力。
TensorFlow 的设计目标和核心优势之一是其灵活性。TensorFlow 被设计成一个灵活和可扩展的系统，可用于定义任意数据流图（data flow graph）并使用异构计算设备（如 CPU 和 GPU）以分布式方式有效地执行它们。
但是灵活性通常与性能不能兼得。虽然 TensorFlow 旨在定义任何种类的数据流图，但是由于 TensorFlow 分别优化每个 运算/指令（op），所以使所有图都高效地执行是有挑战性的。当一个具有高效实现的运算存在，或者每个运算都是相对重量级的操作（heavyweight operation）时，一切都很好；否则，用户仍然可以从低级 op 中组合 op，但是这种组合不能保证以最有效的方式运行。
这就是我们开发了 XLA（Accelerated Linear Algebra/加速线性代数）的原因，它是一个用于 TensorFlow 的编译器。XLA 使用 JIT 编译技术来分析用户在运行时（runtime）创建的 TensorFlow 图，专门用于实际运行时的维度和类型，它将多个 op 融合在一起并为它们形成高效的本地机器代码——能用于 CPU、GPU 和自定义加速器（例如谷歌的 TPU）。
融合可组合的 op 以提高性能
例如，考虑 tf.nn.softmax（https://www.tensorflow.org/api_docs/python/tf/nn/softmax）运算。它计算的是其参数的 softmax 激活函数如下：
Softmax 可以被用于原始 TensorFlow 的 op 的组合（指数、减法、元素除法等）：
softmax = exp(logits) / reduce_sum(exp(logits), dim)
由于额外的数据移动和 op 内部临时结果的实体化（materialization），这个过程可能是缓慢的。此外，在像 GPU 这样的协处理器上，这样的分解执行可能导致多个「核启动（kernel launches）」，使其速度更加缓慢。
XLA 是编译调试器的秘密武器，它能帮助 TensorFlow 自动优化原始 op 的组合。有了 XLA 的增强，通过在运行时的过程中分析图、融合多个 op 并为融合子图（subgraph）生成有效的机器代码，TensorFlow 能在保留其灵活性的同时而不牺牲运行时的性能。
例如，如上面所示的 softmax 的分解实现，由 XLA 优化后的速度能与人工优化的复合 op 一样快。
更一般地，XLA 可以获取 TensorFlow 运算的整个子图，并将它们融合到需要最少数量内核启动（kernel launch）的高效循环中。例如：
该图中的许多操作可以融合到单个元素的循环（single element-wise loop）中。例如，考虑将偏差向量（bias vector）的单个元素添加到来自 matmul 函数结果的单个元素中。该添加的结果是可以与 0 比较的单个元素（用于 ReLU）。比较的结果可以指数化并除以所有输入的指数的和，从而产生 softmax 的输出。我们不需要为 matmul、add 和 ReLU 创建内存中的中间数组。
s[j] = softmax[j](ReLU(bias[j] + matmul_result[j]))
融合的实现可以在单个元素的循环中计算最终结果，而不需要分配不必要的内存。在更高级的场景中，这些操作甚至可以融合到矩阵乘法中。
XLA 帮助 TensorFlow 保持其灵活性，同时消除性能问题。
在内部基准（internal benchmark）测试中，相比于没有 XLA 的 TensorFlow，XLA 显示了在 Nvidia GPU 上高达 50％的加速。如预期那样，最大的加速来自含有长序列元素操作的模型，因为 XLA 可以将长序列元素操作融合进高效的循环中。然而，XLA 仍然被认为是实验性的，一些基准可能会经历减速过程。
在 TensorFlow 开发者峰会的演讲中，Chris Leary 和 Todd Wang 描述了 TensorFlow 如何利用 XLA、JIT、AOT 和其它编译技术来最小化执行时间并最大限度地利用计算资源。
Youtube 视频链接：https://www.youtube.com/watch?v=kAOanJczHA0
可执行尺寸缩减的极度专业化
除了改进性能，TensorFlow 模型受益于 XLA 的限制内存环境（如移动设备）的要求，因为 XLA 减少了其提供的可执行尺寸（executable size）。
tfcompile 是利用 XLA 进行提前编译（AOT/ahead-of-time compilation）的工具——将整个图（graph）编译为 XLA，然后形成严格的机器代码以实现图中的 op。加上最小运行时间，该方案提供了相当多的尺寸减小。
这种尺寸减小是通过其静态编译隐含的模型所具有的完全专业化来实现的。当模型运行时，不需要 TensorFlow 运行时的全部性能能和灵活性——只有实现用户感兴趣的实际图的 op 被编译为本地代码。也就是说，由 XLA 的 CPU 后端发出的代码的性能仍然远不是最优的；这部分项目需要更多的工作。
对替代性后端和设备的支持
为了在当前的新型计算设备上执行 TensorFlow 图，必须重新实现用于新设备的所有 TensorFlow 的 op（内核）。支持设备可能是非常重要的工作。
通过设计，XLA 通过添加自定义后端（backend）使支持新设备更容易。由于 TensorFlow 可以指向 XLA，因此可以向 XLA 添加新设备后端，从而使其能够运行 TensorFlow 图。XLA 为新设备提供了一个显著更小的实现界面，因为 XLA 操作仅仅是原始的（回想一下 XLA 独自处理复杂 op 的分解）。我们已在下面的页面中记录了向 XLA 添加自定义后端的过程：https://www.tensorflow.org/versions/master/experimental/xla/developing_new_backend。谷歌使用此机制利用 XLA 配置 TPU。
结论与展望
XLA 仍处于发展的早期阶段。在一些使用案例中，它显示出非常有希望的结果，很显然，TensorFlow 未来可以从这项技术中得到更多益处。我们决定尽早向 TensorFlow Github（https://github.com/tensorflow/tilerflow/tree/master/tensorflow/compiler）发布 XLA，以征求社群的意见，并为各种计算设备优化 TensorFlow 提供方便的界面，以及重新定位 TensorFlow 的运行时和建立模型以在新型硬件上运行。
