上一章我们分析了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.
作者:Jony0917
链接:https://www.jianshu.com/p/cee16080b5be
。