TensorFlow技术内幕（八）：模型优化之XLA（下）

上一章我们分析了XLA在TensofFlow中的两种调用方式AOT和JIT，本章分析XLA编译器的实现。

LLVM

提到编译器就不得不提大名鼎鼎的LLVM。LLVM是一个编译器框架，由C++语言编写而成，包括一系列分模块、可重用的编译工具。

LLVM框架的主要组成部分有：

• 前端：负责将源代码转换为一种中间表示

• 优化器：负责优化中间代码

• 后端：生成可执行机器码的模块

图1：LLVM框架结构

LLVM为不同的语言提供了同一种中间表示LLVM IR，这样子如果我们需要开发一种新的语言的时候，我们只需要实现对应的前端模块，如果我们想要支持一种新的硬件，我们只需要实现对应的后端模块，其他部分可以复用。

XLA目录结构

XLA的实现目录是tensorflow/compiler,目录结构如下：

目录名	功能
aot	aot编译相关代码，前面分析的tfcompile_tool代码就在这里
jit	jit编译相关代码，例如xlalaunch节点的OpKenel、XLA相关的计算图重构，都在这里
plugin	此模块看起来还没完成，暂不分析
tests	测试代码
tf2xla	GraphDef转化为XLA Hlo IR代码
xla	xla编译器核心代码，HLO IR转化为LLVM IR以及机器码的生成

XLA编译

XLA也是基于LLVM框架开发的，前端的输入是Graph，前端没有将Graph直接转化为LLVM IR，而是转化为了XLA的自定义的中间表示HLO IR.并且为HLO IR设计了一系列的优化器。经过优化的HLO IR接下来会被转化为LLVM IR。

图2：XLA框架结构

具体来说包含了下列几步：

• 步骤一：由GraphDef创建Graph

• 步骤二：由tensorflow.Graph编译为HLO IR

• 步骤三：分析与优化HLO IR

• 步骤四：由HLO IR转化为llvm IR

• 步骤五：分析与优化llvm IR

• 步骤六：生成特定平台的二进制文件

AOT

AOT编译流程图：

图3：AOT编译流程

对照图2来分析一下AOT编译流程:

• tensorflow.XlaCompiler.CompilerGraph函数将Graph编译成XLA的中间表示xla.UserComputation.

• tensorflow.XlaCompiler.CompilerGraph会创建Executor来执行待编译的Graph，通过绑定设备，为所有节点的创建运算核都是专门设计用来编译的，基类是tensorflow.XlaOpKernel.

• tensorflow.XlaOpKernel的子类需要实现Compile接口，通过调用xla.ComputeBuilder接口，将本节点的运算转化为Xla指令(instruction).

• xla.ComputeBuilder是对xla.Client的调用封装，通过本接口创建的xla指令(instruction)的操作，最终都会通过xla.Client传输到xla.Service.

• xla.Client 和 xla.Service 支持单机模式和分布式模式，实际的编译过程发生在Service端.

• AOT编译中，用到的是 xla.CompileOnlyClient 和 xla.CompileOnlyService，分别是xla.Client和xla.Service的实现类.

• 可以看到，图2中的第一个循环（loop for every node）会为每个node生成一系列xla指令(instruction)，这些指令最终会被加入xla.UserComputation的指令队列里。

• 接下来xla.CompileOnlyClient.CompileAheadOfTime会将xla.UserComputation编译为可执行代码.

• xla.ComputationTracker.BuildHloModule函数会将所有的xla.UserComputation转化为xla.HloComputation，并为之创建xla.HloModule.

• 至此，Graph 到 HLO IR 的转化阶段完成。

• HLO IR进入后续的编译过程，根据平台调用不同平台的具体编译器实现类，这里我们以xla.CpuComiler为例来分析.

• xla.CpuComiler的输入是xla.HloModule，首先会调用RunHloPasses创建HloPassPipeline，添加并运行一系列的HloPass.

• 每一个HloPass都实现了一类HLO指令优化逻辑。通常也是我们比较关心的逻辑所在，包含单不限于图中列举出来的
  xla.AlebraicSimplifier(代数简化),xla.HloConstantFolding（常量折叠）,xla.HloCSE（公共表达式消除）等。

• HloPassPipeline优化HLO IR之后，将创建xla.cpu.IrEmitter，进入图2中的第三个循环处理逻辑(loop for every computation of module)：将xla.HloModule中的每个xla.HloComputation转化为llvm IR表示，并创建对应的llvm.Module.

• 至此，Hlo IR 到 llvm IR的转化阶段完成，后面进入llvm IR的处理阶段。

• 创建xla.cpu.CompilerFunctor将llvm IR转化为最终的可执行机器代码llvm.object.ObjectFile.中间会调用一系列的llvm ir pass对llvm ir进行优化处理。

• 至此，llvm ir到可执行机器码的转化阶段完成。

JIT

JIT编译流程图：

图4：JIT编译流程

JIT对比AOT来说，过程比较类似，略过共同的部分，我们来分析一下：

• JIT调用方式的入口在运算核tensorflow.XlaLocalLaunchOp.Compute，tensorflow.XlaLocalLaunchOp是连接外部Graph的Executor和内部JIT调用的桥梁。

• 如果被调用的计算图缓存不命中，则会调用xla.XlaCompile进行实际的编译。

• 编译过程类似AOT，不同之处主要在于：首先这次调用的Client和Service的实现类是xla.LocalClient和xla.LocalService；其次，llvm ir到机器码的编译过程，这次是通过xla.cpu.SimpleOrcJIT完成的，它将llvm ir编译为可执行代码，并可被立即调用。

• 可执行机器码后续会被封装为xla.LocalExecutale

• 调用xla.LocalExecutable的如后函数Run.