前文回顾
上一篇文章《TensorFlow XLA 学习笔记（一）：源码编译 + 运行例程》介绍了如何通过源码编译带有 XLA 功能的 TensorFlow 并运行例程，实测开启 XLA 后 GPU 上运行 AlexNet Benchmark 有 2~3 倍性能提升。本篇继续深入探索 XLA 具体做了哪些事情。
打印 VLOG 信息
运行 TensorFlow 代码之前，设置环境变量：
export TF_CPP_MIN_VLOG_LEVEL=2
这时运行程序会打印大量日志信息，通过利用这些日志信息可以对 TensorFlow 大体运行流程做一个简单了解。
1. 注册运算符
在程序开始时，注册了大量 XLA ops 如 Conv2D，MaxPool3D，Acosh 等，这些 ops 对应 Device 有 XLA_CPU_JIT 和 XLA_GPU_JIT 两种，分别运行在 CPU 和 GPU 上。截取片段如下图所示：
注意 ：如果你的 TensorFlow 是直接安装官方预编译好的 whl，那么不会出现以上信息。需要使用源码编译开启 XLA 选项。
2. 图优化(GraphOptimization)
TensorFlow 将计算流以有向图的形式组织，其中每个节点（Node）都表示特定运算，每个边为 Tensor 或 Control Flow。图可以用底层 API 直接构建或使用高级 API 编写算法交给框架自动生成计算图。大部分计算图在运行时效率不高，例如我们常见的 Convolution + Batch Normalization + Relu，在 TensorFlow Graph 中直接映射为：
其中 Batch Normalization 被分解为多个运算符（Mul, Add, Sub, Rsqrt……），这种表示在 GPU 上计算时由于会多次启动核函数，执行效率不高。一种图优化技巧是预先计算出某几个节点的结果，在运行时可省去重复计算过程。下图是采用“fold_batch_norms” 策略优化后的计算图：
该优化策略根据 inference 时 Batch Normalization 参数为常量的特性，将 Batch Normalization 参数全部折合到 Conv2D 运算符中，大幅减少了运算符数目。
图优化相关流程可以在运行日志中搜索“Graph after optimization”关键词定位。
3. 运算符替换
经过图优化，减少了无用节点，接下来 XLA 优化会选择部分节点替换为动态生成的新运算符，名字为“_XlaLaunch” 。
替换前的 timeline：
替换后的 timeline:
运行日志中搜索“_XlaLaunch”关键字可以定位到替换过程：
4. 编译缓存
做完上述运算符替换后，会由 TensorFlow 引擎自动调度执行 _XlaLaunch 运算符对应的核函数，该过程会触发 JIT 编译过程，动态生成运行时代码。
搜索日志关键词“XlaCompilationCache::Compile XLA JIT compilation cache” 可定位到该过程。
XlaCompilationCache 作用是将之前动态生成的代码缓存，如果后面又遇到相同情况，可以直接使用上次生成的代码。
5. HLO 优化
HLO 是 XLA 编译器的中间表示（IR），编译流程如下图所示：
在日志中搜索关键词“Running HLO pass pipeline optimization” 可以看到很多组 HLO 优化过程（一次优化称为一个 pass），以下为其中一个 pass：
可以看到类似汇编语言的中间表示即为 HLO IR。这些运算直接反映了编译器内部表示的程序执行过程。
6. LLVM 生成设备代码
LLVM 负责从 HLO IR 到设备代码生成过程，在 GPU 上对应的设备代码就是 PTX。运行日志中搜索关键词“GpuCompiler::RunBackend - CompileToPtx time” 可定位到该过程。生成的部分 PTX 代码打印如下：
详细指令介绍请参考 Nvidia CUDA 手册《PARALLEL THREAD EXECUTION ISA》。
8. 硬件上运行
生成的 PTX 进一步被 ptxas 工具编译为 cubin 文件，加载到一个或多个 GPU stream 运行，得到结果。
小结
目前针对 TensorFlow XLA 运行机制分析的博客比较少，一方面 TF 是一个非常复杂的系统，涵盖图计算、机器学习算法、编译原理、体系结构等内容，需要深厚的技术积累，另一方面在图文排版、代码高亮显示方面缺乏良好的阅读体验。本文笔者做了不算全面的 XLA 介绍，计划后面结合代码进行各个模块梳理。如果有更好的建议可以后台留言反馈，一起学习ing~
本文使用的运行日志可以后台回复“20180304”获取。