Using JIT Compilation（使用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

结果如下所示：