XLA - TensorFlow, compiled

原文链接：https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html
发布时间：2017-03

XLA helps TensorFlow retain its flexibility while eliminating performance concerns.

TF的核心竞争力在于其灵活性、扩展性，可以方便的定义任意数据流图并在多种计算设备（如CPUs、GPUs）上运行。

为了保持灵活性，往往会损失性能。TF单独地优化每个OP，如果一个OP存在高效的实现方式或者每个OP都比较重量级，那问题不大；否则，用户可以用一些低级的OP重新组合成这个OP，不过这种方法无法保证是最高效的方法。

XLA（Accelerated Linear Algebra），一个用于TF的编译器。XLA使用JIT编译技术分析运行时的TF图，专门针对运行时的维度和类型，将多个OP融合在一起并生成对应的高效原生机器代码。

比如softmax可以由TF原生OP组合而成。

softmax = exp(logits) / reduce_sum(exp(logits), dim)

这种方法可能会比较低效，因为这个方法包括了额外的数据移动操作、实例化临时结果，这些其实是softmax OP不需要的。另外，如果在GPUs这种协处理器上，拆分实现可能会引起多次“内核启动”，适得其反。

Tensorflow, augmented with XLA, retains flexibility without sacrificing runtime performance, by analyzing the graph at runtime, fusing ops together and producing efficient machine code for the fused subgraphs.

XLA可以处理整个TF子图并将它们融合到一个高效的循环中，这样仅需要最少的内核启动即可。比如下图：

可以被处理为s[j] = softmax[j](ReLU(bias[j] + matmul_result[j]))。一个融合后的实现可以直接计算最终结果，而无需申请不必要的内存。

On internal benchmarks, XLA shows up to 50% speedups over TensorFlow without XLA on Nvidia GPUs.However, XLA should still be considered experimental, and some benchmarks may experience slowdowns.