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.

猜你喜欢

转载自blog.csdn.net/gaofeipaopaotang/article/details/80703367