上一章我們分析了XLA在TensofFlow中的兩種調(diào)用方式AOT和JIT，本章分析XLA編譯器的實現(xiàn)。

LLVM

提到編譯器就不得不提大名鼎鼎的LLVM。LLVM是一個編譯器框架，由C++語言編寫而成，包括一系列分模塊、可重用的編譯工具。

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

• 前端：負責將源代碼轉(zhuǎn)換為一種中間表示

• 優(yōu)化器：負責優(yōu)化中間代碼

• 后端：生成可執(zhí)行機器碼的模塊

圖1：LLVM框架結構

LLVM為不同的語言提供了同一種中間表示LLVM IR，這樣子如果我們需要開發(fā)一種新的語言的時候，我們只需要實現(xiàn)對應的前端模塊，如果我們想要支持一種新的硬件，我們只需要實現(xiàn)對應的后端模塊，其他部分可以復用。

XLA目錄結構

XLA的實現(xiàn)目錄是tensorflow/compiler,目錄結構如下：

XLA編譯

XLA也是基于LLVM框架開發(fā)的，前端的輸入是Graph，前端沒有將Graph直接轉(zhuǎn)化為LLVM IR，而是轉(zhuǎn)化為了XLA的自定義的中間表示HLO IR.并且為HLO IR設計了一系列的優(yōu)化器。經(jīng)過優(yōu)化的HLO IR接下來會被轉(zhuǎn)化為LLVM IR。

圖2：XLA框架結構

具體來說包含了下列幾步：

• 步驟一：由GraphDef創(chuàng)建Graph

• 步驟二：由tensorflow.Graph編譯為HLO IR

• 步驟三：分析與優(yōu)化HLO IR

• 步驟四：由HLO IR轉(zhuǎn)化為llvm IR

• 步驟五：分析與優(yōu)化llvm IR

• 步驟六：生成特定平臺的二進制文件

AOT

AOT編譯流程圖：

圖3：AOT編譯流程

對照圖2來分析一下AOT編譯流程:

• tensorflow.XlaCompiler.CompilerGraph函數(shù)將Graph編譯成XLA的中間表示xla.UserComputation.

• tensorflow.XlaCompiler.CompilerGraph會創(chuàng)建Executor來執(zhí)行待編譯的Graph，通過綁定設備，為所有節(jié)點的創(chuàng)建運算核都是專門設計用來編譯的，基類是tensorflow.XlaOpKernel.

• tensorflow.XlaOpKernel的子類需要實現(xiàn)Compile接口，通過調(diào)用xla.ComputeBuilder接口，將本節(jié)點的運算轉(zhuǎn)化為Xla指令(instruction).

• xla.ComputeBuilder是對xla.Client的調(diào)用封裝，通過本接口創(chuàng)建的xla指令(instruction)的操作，最終都會通過xla.Client傳輸?shù)絰la.Service.

• xla.Client 和 xla.Service 支持單機模式和分布式模式，實際的編譯過程發(fā)生在Service端.

• AOT編譯中，用到的是 xla.CompileOnlyClient 和 xla.CompileOnlyService，分別是xla.Client和xla.Service的實現(xiàn)類.

• 可以看到，圖2中的第一個循環(huán)（loop for every node）會為每個node生成一系列xla指令(instruction)，這些指令最終會被加入xla.UserComputation的指令隊列里。

• 接下來xla.CompileOnlyClient.CompileAheadOfTime會將xla.UserComputation編譯為可執(zhí)行代碼.

• xla.ComputationTracker.BuildHloModule函數(shù)會將所有的xla.UserComputation轉(zhuǎn)化為xla.HloComputation，并為之創(chuàng)建xla.HloModule.

• 至此，Graph 到 HLO IR 的轉(zhuǎn)化階段完成。

• HLO IR進入后續(xù)的編譯過程，根據(jù)平臺調(diào)用不同平臺的具體編譯器實現(xiàn)類，這里我們以xla.CpuComiler為例來分析.

• xla.CpuComiler的輸入是xla.HloModule，首先會調(diào)用RunHloPasses創(chuàng)建HloPassPipeline，添加并運行一系列的HloPass.

• 每一個HloPass都實現(xiàn)了一類HLO指令優(yōu)化邏輯。通常也是我們比較關心的邏輯所在，包含單不限于圖中列舉出來的
  xla.AlebraicSimplifier(代數(shù)簡化),xla.HloConstantFolding（常量折疊）,xla.HloCSE（公共表達式消除）等。

• HloPassPipeline優(yōu)化HLO IR之后，將創(chuàng)建xla.cpu.IrEmitter，進入圖2中的第三個循環(huán)處理邏輯(loop for every computation of module)：將xla.HloModule中的每個xla.HloComputation轉(zhuǎn)化為llvm IR表示，并創(chuàng)建對應的llvm.Module.

• 至此，Hlo IR 到 llvm IR的轉(zhuǎn)化階段完成，后面進入llvm IR的處理階段。

• 創(chuàng)建xla.cpu.CompilerFunctor將llvm IR轉(zhuǎn)化為最終的可執(zhí)行機器代碼llvm.object.ObjectFile.中間會調(diào)用一系列的llvm ir pass對llvm ir進行優(yōu)化處理。

• 至此，llvm ir到可執(zhí)行機器碼的轉(zhuǎn)化階段完成。

JIT

JIT編譯流程圖：

圖4：JIT編譯流程

JIT對比AOT來說，過程比較類似，略過共同的部分，我們來分析一下：

• JIT調(diào)用方式的入口在運算核tensorflow.XlaLocalLaunchOp.Compute，tensorflow.XlaLocalLaunchOp是連接外部Graph的Executor和內(nèi)部JIT調(diào)用的橋梁。

• 如果被調(diào)用的計算圖緩存不命中，則會調(diào)用xla.XlaCompile進行實際的編譯。

• 編譯過程類似AOT，不同之處主要在于：首先這次調(diào)用的Client和Service的實現(xiàn)類是xla.LocalClient和xla.LocalService；其次，llvm ir到機器碼的編譯過程，這次是通過xla.cpu.SimpleOrcJIT完成的，它將llvm ir編譯為可執(zhí)行代碼，并可被立即調(diào)用。

• 可執(zhí)行機器碼后續(xù)會被封裝為xla.LocalExecutale

• 調(diào)用xla.LocalExecutable的如后函數(shù)Run.