XLA优化实例

XLA概述

纠错

XLA（加速线性代数）是用于优化TensorFlow计算的线性代数的域特定编译器。结果是在服务器和移动平台上的速度，内存使用率和可移植性得到了改善。最初，大多数用户不会从XLA中看到很大的好处，通过使用即时（JIT）编译或提前编译（AOT）的XLA进行试验，针对新硬件加速器尝试XLA。

XLA框架是实验性和积极的开发。尽管现有操作的语义不太可能发生变化，但预计将增加更多的操作来涵盖重要的用例。

为什么要构建XLA？

XLA与TensorFlow合作有几个目标：

• 提高执行速度。编译子图以减少短暂Ops的执行时间，以消除TensorFlow运行时间的开销，融合流水线操作以减少内存开销，并专用于已知张量形状以允许更积极的恒定传播。
• 改善内存使用。分析和计划内存使用情况，原则上消除许多中间存储缓冲区。
• 减少对自定义操作的依赖。通过改进自动融合低级Ops的性能来消除对许多自定义Ops的需求，以匹配手工融合的自定义Ops的性能。
• 减少移动足迹。通过提前编译子图并发出可以直接链接到另一个应用程序的对象/头文件对来消除TensorFlow运行时。结果可以将移动推断的占用空间减少几个数量级。
• 提高可移植性。为新颖的硬件编写新的后端程序相对容易，此时大部分TensorFlow程序将在该硬件上未修改地运行。这与专门针对新硬件的个体单片Ops的方法形成对比，这需要重写TensorFlow程序以利用这些Ops。

XLA如何工作？

XLA的输入语言称为“HLO IR”，或称为HLO（高级优化程序）。操作语义页面描述了HLO的语义。将HLO视为编译器IR是最方便的。

XLA将HLO中定义的图形（“计算”）编译成各种体系结构的机器指令。XLA是模块化的，很容易插入替代后端，以便定位一些新颖的硬件架构。用于x64和ARM64的CPU后端以及NVIDIA GPU后端均位于TensorFlow源代码树中。

纠错

下图显示了XLA中的编译过程：

 

 

 XLA带有多个与目标无关的优化和分析，如CSE，独立于目标的操作融合以及为计算分配运行时内存的缓冲区分析。

在独立于目标的步骤之后，XLA将HLO计算发送到后端。后端可以执行进一步的HLO级别分析和优化，针对具体目标信息和需求。例如，XLA GPU后端可以执行专用于GPU编程模型的算子融合，并确定如何将计算划分为流。这时，后端也可以模式匹配某些操作或其组合来优化库调用。

下一步是目标特定的代码生成。XLA附带的CPU和GPU后端使用LLVM进行低级IR，优化和代码生成。这些后端以有效的方式发出代表XLA HLO计算所需的LLVM IR，然后调用LLVM从此LLVM IR发出本机代码。

纠错

GPU后端当前通过LLVM NVPTX后端支持NVIDIA GPU; CPU后端支持多个CPU ISA。

支持的平台

XLA目前支持x86-64和NVIDIA GPU上的JIT编译; 以及针对x86-64和ARM的AOT编译。

XLA开发后端

XLA提供了一个抽象接口，新体系结构或加速器可以实现创建后端以运行TensorFlow图形。重新定位XLA应该比实现每个现有的TensorFlow Op用于新硬件更加简单和可扩展。

大多数实现将落入以下情况之一：

1. 现有的CPU体系结构尚未正式由XLA支持，无论是否存在LLVM后端。

2. 具有现有LLVM后端的非CPU类硬件。

3. 没有现有LLVM后端的非CPU类硬件。

注意： LLVM后端可以是官方发布的LLVM后端或内部开发的定制LLVM后端。

情况1：XLA尚未正式支持现有CPU架构

在这种情况下，首先查看现有的XLA CPU后端。通过使用LLVM，XLA可以轻松地将TensorFlow重定向到不同的CPU，因为XLA后端对于CPU的主要区别在于LLVM生成的代码。Google测试XLA for x64和ARM64体系结构。

如果硬件供应商为其硬件提供LLVM后端，则将后端与使用XLA构建的LLVM进行链接很简单。在JIT模式下，XLA CPU后端为主机CPU发出代码。对于提前编译，xla::AotCompilationOptions可以提供一个LLVM三元组来配置目标体系结构。

如果没有现有的LLVM后端，但存在另一种代码生成器，则应该可以重新使用大部分现有的CPU后端。

场景2：具有现有LLVM后端的非CPU类硬件

可以xla::Compiler在现有类xla::CPUCompiler和xla::GPUCompiler类上建立一个新的实现，因为它们已经发出了LLVM IR。根据硬件的性质，许多LLVM IR生成方面可能需要更改，但可以与现有后端共享大量代码。

一个很好的例子就是XLA 的GPU后端。GPU后端以非CPU类ISA为目标，代码生成的某些方面对于GPU域是唯一的。其它类型的硬件，例如Hexagon（具有上游LLVM后端）的DSP可以重新使用部分LLVM IR发射逻辑，但其它部分将是唯一的。

场景3：没有现有LLVM后端的非CPU类硬件

如果无法使用LLVM，最好的选择是为XLA实现所需硬件的新后端。这个选项需要最多的努力。需要实施的类如下：

• StreamExecutor：对于许多设备，并非所有的方法StreamExecutor都是必需的。详情请参阅现有的StreamExecutor实施。
• xla :: Compiler：这个类将HLO计算的编译封装为一个xla::Executable。
• xla::Executable：该类用于在平台上启动编译的计算。
• xla::TransferManager：该类使后端能够提供特定于平台的机制，用于从给定的设备内存句柄构造XLA文字数据。换句话说，它有助于封装从主机到设备的数据传输并返回。

使用JIT编译

TensorFlow必须从源代码编译为包含XLA。

为什么要使用即时（JIT）编译？

TensorFlow / XLA JIT编译器通过XLA编译和运行TensorFlow图形的一部分。与标准TensorFlow实现相比，这样做的好处是XLA可以将多个运算符（内核融合）融合到少量的编译内核中。与TensorFlow执行程序一样，与一次执行操作员相比，定位操作员可以减少内存带宽要求并提高性能。

通过XLA运行TensorFlow图表

有两种方法通过XLA运行TensorFlow计算，或者通过JIT编译操作员放置在CPU或GPU的设备上，或通过将操作员在XLA_CPU或XLA_GPUTensorFlow设备。将算子直接放在TensorFlow XLA设备上强制算子在该设备上运行，主要用于测试。

Note: The XLA CPU backend produces fast single-threaded code (in most cases), but does not yet parallelize as well as the TensorFlow CPU backend. The XLA GPU backend is competitive with the standard TensorFlow implementation, sometimes faster, sometimes slower.

打开JIT编译

JIT编译可以在会话级别打开或手动进行选择操作。这两种方法都是零拷贝---在编译的XLA内核和置于同一设备上的TensorFlow操作符之间传递数据时，不需要复制数据。

Session

在会话级别打开JIT编译会导致所有可能的操作符被贪婪地编译成XLA计算。每个XLA计算将被编译为一个或多个内核设备。

受限于一些限制，如果图中有两个相邻的运算符都具有XLA实现，那么它们将被编译为单个XLA计算。

JIT编译在会话级别打开，方法是在会话初始化期间将config 设置global_jit_level为tf.OptimizerOptions.ON_1并传递配置。

# Config to turn on JIT compilation

config = tf.ConfigProto()

config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

sess = tf.Session(config=config)

Note: Turning on JIT at the session level will not result in operations being compiled for the CPU. JIT compilation for CPU operations must be done via the manual method documented below. This decision was made due to the CPU backend being single-threaded.

手动

JIT编译也可以为一个或多个操作员手动打开。这是通过标记操作符以使用属性进行编译完成的_XlaCompile=true。最简单的方法是通过在中tf.contrib.compiler.jit.experimental_jit_scope()定义的范围tensorflow/contrib/compiler/jit.py。用法示例：

jit_scope = tf.contrib.compiler.jit.experimental_jit_scope

x = tf.placeholder(np.float32)

with jit_scope():

  y = tf.add(x, x)  # The "add" will be compiled with XLA.

该_XlaCompile属性目前支持尽力而为。如果无法编译运算符，则TensorFlow将默默回退到正常实现。

将操作员放置在XLA设备上

通过XLA运行计算的另一种方法是将操作员放置在特定的XLA设备上。此方法通常仅用于测试。有效目标是XLA_CPU或XLA_GPU。

with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):

  output = tf.add(input1, input2)

与标准CPU和GPU设备上的JIT编译不同，这些设备在将数据传输到设备上或从设备传输时将复制数据。额外的副本使XLA和TensorFlow操作符在同一个图中混合成本很高。

如何在开启JIT的情况下训练MNIST softmax。当前在会话级别的JIT，仅支持GPU。

验证LD_LIBRARY环境变量或ldconfig包含$CUDA_ROOT/extras/CUPTI/lib64，其中包含CUDA分析工具界面（CUPTI）的库。TensorFlow使用CUPTI从GPU中提取跟踪信息。

步骤＃1：准备示例脚本

将mnist_softmax_xla.py下载或移动到TensorFlow源代码树之外的文件夹中。

步骤＃2：在没有XLA的情况下运行

执行python脚本来训练没有XLA的模型。

python mnist_softmax_xla.py --xla=''

使用Chrome Trace Event Profiler（浏览至chrome：// tracing），打开脚本完成时创建的时间线文件：timeline.ctf.json。呈现的时间线应与下图类似，并标注多个绿色框MatMul，可能跨多个CPU。

 

步骤＃3使用XLA运行

执行python脚本以XLA训练模型，并通过输出XLA图形的环境变量打开XLA的调试功能。

纠错

TF_XLA_FLAGS=--xla_generate_hlo_graph=.* python mnist_softmax_xla.py

打开创建的时间轴文件（timeline.ctf.json）。呈现的时间线应该与下面的图片类似，并标注一个长条_XlaLaunch。

 

 

 要了解正在发生的事情_XlaLaunch，请查看控制台输出以获取类似于以下内容的语句：

computation cluster_0[_XlaCompiledKernel=true,_XlaNumConstantArgs=1].v82 [CPU:

pipeline start, before inline]: /tmp/hlo_graph_0.dot

控制台语句指向hlo_graph_xx.dot包含由XLA创建的图表信息的文件的位置。XLA用来融合Ops的过程可以通过hlo_graph_0.dot连续观察每个图表来看到。

纠错

要将.dot文件渲染为png，请安装GraphViz并运行：

dot -Tpng hlo_graph_80.dot -o hlo_graph_80.png

结果如下所示：