xla

 

https://zhuanlan.zhihu.com/p/124269986

数据流通转换器，xla编译器，Jit，aot

XLA是对tensorflow原有函数及运行时组件的增加

 

数据流图转换器用于将TensorFlow数据流图转换为XLA图，其源代码位于tensorflow/compiler/tf2xla。该模块定义了名为XlaOpKernel的XLA层函数类型，并且实现了一系列线性代数相关的核函数子类。在数据流图执行前，转换器基于数据流图的结构创建XLA图，将原图中那些可通过XLA优化的操作节点替换为XLA层函数。由于数据流图操作同HLO IR存在语义差异，节点的替换可能是一对一的，也可能是一对多的。

 

XLA编译器用于将XLA图编译为二进制文件，其源代码位于tensorflow/compile/xla。该模块基于LLVM技术实现，目前支持的后端设备包括x86 CPU、arm CPU和NV GPU。在编译器中，前端解析并优化XLA图，将其底层化为LLVM IR，后端将LLVM IR编译为设备特征的二进制文件，并执行二进制代码优化。XLA编译器采用客户端-服务器模式设计，以便管理编译过程的生命周期。HLO的实现亦位于该模块。

 

JIT编译机制用于在Tensorflow应用运行时创建执行数据流图操作的二进制代码，它是一套面向异构设备的通用性能优化机制。该模块的源代码位于tensorflow/compiler/jit。模型开发者使用JIT机制时，需要在代码中显式引入JIT API，将会话或操作配置为XLA编译模式。JIT编译机制会对数据流图的可优化部分实施优化，将其中大量细粒度的操作融合为少量粗粒度的专用核函数。这些核函数经由XLA编译器生成高效的二进制代码，能够减少图执行过程中内存分配和上下文切换开销

 

AOT编译机制用于在Tensorflow应用运行前创建集成了模型和运行时的二进制代码，主要适用于手机等移动端的推理性能优化。该代码位于tensorflow/compiler/aot。模型开发者使用AOT机制时，需要提供protocol buffers格式编写的数据流图定义文件及编译配置文件。然后使用bazel构建工具，调用AOT封装工具-tfcompile编译数据流图。AOT编译机制生成的二进制代码包含了模型和必要的运行时逻辑，不再需要完整的Tensorflow的运行时库支持，因此能够减小可执行程序的体积，一并生成的c++文件用于在应用代码中访问模型。

 

AOT的编译机制比较适合移动设备，因此这里仅介绍JIT编译机制，很方便，主要有两种方式：

1，在会话打开打开JIT编译。把所有可能操作符编程成XLA计算。

config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level= tf.OptimizerOptions.ON_1
sess = tf.Session(config=config) 

2，为一个或多个操作符手动打开JIT编译。属性_XlaCompile = true标记编译操作符。

jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
x = tf.placeholder(np.float32)
with jit_scope():
  y = tf.add(x, x)

XLA 使用 xla.compile API，让您可以在部分 TensorFlow 图表上显式调用 XLA 编译器。xla.compile 会接受生成 TensorFlow 计算的 Python 函数，然后连接所生成的计算以供 XLA 编译。xla.compile 还会返回张量列表，其中每个张量对应传入函数构建的计算结果，但现在会立即由 XLA 优化。

 

因此，通过调用 xla.compile，您可以使用 XLA 运行以上由 model_fn 生成的计算，如下所示：

1    from tensorflow.contrib.compiler import xla

2

3    def model_fn(x, y, z):

4        return tf.reduce_sum(x + y * z)

5

6    def create_and_run_graph():

7        with tf.Session() as sess:

8            x = tf.placeholder(tf.float32, name='x')

9            y = tf.placeholder(tf.float32, name='y')

10            z = tf.placeholder(tf.float32, name='z')

11            result = xla.compile(computation=model_fn, inputs=(x, y, z))[0] 

12            # `result` is a normal Tensor (albeit one that is computed by an XLA 

13            # compiled executable) and can be used like any other Tensor. 

14            result = tf.add(result, result)

15            return sess.run(result, feed_dict={ ... })

 

您可以使用命令行标记（或其他任意逻辑）来控制是否由 XLA 编译计算。模型通常按照以下方式调用 xla.compile 以便完成简单实验。

1    if should_use_xla():    

2        result = xla.compile(model_fn, (x, y, z))[0]    

3    else:    

4        result = model_fn(x, y, z) 

四、使用XLA的注意事项

1，所有运算都必须具有可推断的形状XLA 需要能够在给定计算输入的情况下，推断出其编译的所有运算的形状。因此，如果模型函数生成的 Tensor 具有不可推断的形状，则运行时将会出现错误，进而导致运行失败。

2，XLA 必须支持所有运算并非所有TensorFlow 运算都可由 XLA 编译，如果模型中有XLA 不支持的运算，XLA 编译就会失败。例如，XLA 不支持 tf.where 运算，因此如果您的模型函数包含此运算，使用 xla.compile 运行模型时便会失败。XLA 支持的每项 TensorFlow 运算都可在tensorflow/compiler/tf2xla/kernels/ 中调用 REGISTER_XLA_OP，因此您可以使用 grep 来搜索 REGISTER_XLA_OP 宏实例，以查找支持的 TensorFlow 运算列表。

 

 

 

 https://mp.weixin.qq.com/s/MPI9KERDS-Al4DTBDRV04w

 

 

https://sketch2sky.com/2019/09/24/tensorflow-jit-%E6%8A%80%E6%9C%AF%E8%AF%A6%E8%A7%A3/