上一章我们分析了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.
网友评论