TensorFlow XLA優化原理與示例
XLA概述
XLA（加速線性代數）是用於優化TensorFlow計算的線性代數的域特定編譯器。結果是在服務器和移動平台上的速度，內存使用率和可移植性得到了改善。最初，大多數用戶不會從XLA中看到很大的好處，通過使用即時（JIT）編譯或提前編譯（AOT）的XLA進行試驗，針對新硬件加速器嘗試XLA。
XLA框架是實驗性和積極的開發。盡管現有操作的語義不太可能發生變化，但預計將增加更多的操作來涵蓋重要的用例。
構建XLA
XLA與TensorFlow合作有幾個目標：
- 提高執行速度。編譯子圖以減少短暫Ops的執行時間，以消除TensorFlow運行時間的開銷，融合流水線操作以減少內存開銷，並專用於已知張量形狀以允許更積極的恆定傳播。
- 改善內存使用。分析和計划內存使用情況，原則上消除許多中間存儲緩沖區。
- 減少對自定義操作的依賴。通過改進自動融合低級Ops的性能，消除對許多自定義Ops的需求，匹配手工融合的自定義Ops的性能。
- 減少移動足跡。通過提前編譯子圖。發出可以直接鏈接到另一個應用程序的對象/頭文件對來消除TensorFlow運行時。結果可以將移動推斷的占用空間。減少幾個數量級。
- 提高可移植性。為新穎的硬件編寫新的后端程序相對容易，此時大部分TensorFlow程序。將在該硬件上未修改地運行。與專門針對新硬件的個體單片Ops的方法形成對比，需要重寫TensorFlow程序以利用這些Ops。
XLA如何工作？
XLA的輸入語言稱為“HLO IR”，或稱為HLO（高級優化程序）。操作語義頁面描述了HLO的語義。將HLO視為編譯器IR是最方便的。
XLA將HLO中定義的圖形（“計算”）編譯成各種體系結構的機器指令。XLA是模塊化的，很容易插入替代后端，以便定位一些新穎的硬件架構。用於x64和ARM64的CPU后端，以及NVIDIA GPU后端，均位於TensorFlow源代碼樹中。
下圖顯示了XLA中的編譯過程：
XLA帶有多個與目標無關的優化和分析，如CSE，獨立於目標的操作融合，以及為計算分配運行時，內存的緩沖區分析。
在獨立於目標的步驟之后，XLA將HLO計算發送到后端。后端可以執行進一步的HLO級別分析和優化，針對具體目標信息和需求。例如，XLA GPU后端可以執行專用於GPU編程模型的算子融合，並確定如何將計算划分為流。這時，后端也可以模式匹配某些操作或其組合來優化庫調用。
下一步是目標特定的代碼生成。XLA附帶的CPU和GPU后端使用LLVM進行低級IR，優化和代碼生成。這些后端以有效的方式，發出代表XLA HLO計算所需的LLVM IR，然后調用LLVM從此LLVM IR發出本機代碼。
糾錯
GPU后端當前通過LLVM NVPTX后端，支持NVIDIA GPU; CPU后端支持多個CPU ISA。
支持的平台
XLA目前支持x86-64和NVIDIA GPU上的JIT編譯; 以及針對x86-64和ARM的AOT編譯。
XLA開發后端
XLA提供了一個抽象接口，新體系結構或加速器可以實現創建后端，運行TensorFlow圖形。重新定位XLA，應該比實現每個現有的TensorFlow Op用於新硬件，更加簡單和可擴展。
大多數實現將落入以下情況之一：
1. 現有的CPU體系結構，尚未正式由XLA支持，無論是否存在LLVM后端。
2. 具有現有LLVM后端的非CPU類硬件。
3. 沒有現有LLVM后端的非CPU類硬件。
注意： LLVM后端可以是官方發布的LLVM后端，或內部開發的定制LLVM后端。
情況1：XLA尚未正式支持現有CPU架構
在這種情況下，首先查看現有的XLA CPU后端。通過使用LLVM，XLA可以輕松地將TensorFlow重定向到不同的CPU，因為XLA后端對於CPU的主要區別在於LLVM生成的代碼。Google測試XLA for x64和ARM64體系結構。
如果硬件供應商為其硬件提供LLVM后端，則將后端與使用XLA構建的LLVM進行鏈接很簡單。在JIT模式下，XLA CPU后端為主機CPU發出代碼。對於提前編譯，xla::AotCompilationOptions可以提供一個LLVM三元組來配置目標體系結構。
如果沒有現有的LLVM后端，但存在另一種代碼生成器，則應該可以重新使用大部分現有的CPU后端。
場景2：具有現有LLVM后端的非CPU類硬件
可以xla::Compiler在現有類xla::CPUCompiler和xla::GPUCompiler類上建立一個新的實現，因為它們已經發出了LLVM IR。根據硬件的性質，許多LLVM IR生成方面可能需要更改，但可以與現有后端共享大量代碼。
一個很好的例子就是XLA 的GPU后端。GPU后端以非CPU類ISA為目標，代碼生成的某些方面對於GPU域是唯一的。其它類型的硬件，例如Hexagon（具有上游LLVM后端）的DSP，可以重新使用部分LLVM IR發射邏輯，但其它部分將是唯一的。
場景3：沒有現有LLVM后端的非CPU類硬件
如果無法使用LLVM，最好的選擇是為XLA實現所需硬件的新后端。這個選項需要最多的努力。需要實施的類如下：
- StreamExecutor：對於許多設備，並非所有的方法
StreamExecutor都是必需的。詳情請參閱現有的StreamExecutor實施。 - xla :: Compiler：這個類將HLO計算的編譯封裝為一個
xla::Executable。 xla::Executable：該類用於在平台上啟動編譯的計算。xla::TransferManager：該類使后端能夠提供特定於平台的機制，用於從給定的設備內存句柄構造XLA文字數據。換句話說，有助於封裝從主機到設備的數據傳輸並返回。
使用JIT編譯
TensorFlow必須從源代碼編譯為包含XLA。
使用即時（JIT）編譯
TensorFlow / XLA JIT編譯器，通過XLA編譯和運行TensorFlow圖形的一部分。與標准TensorFlow實現相比，這樣做的好處是XLA可以將多個算子（內核融合），融合到少量的編譯內核中。與TensorFlow執行程序一樣，與一次執行算子相比，定位算子可以減少內存帶寬要求並提高性能。
通過XLA運行TensorFlow圖表
有兩種方法通過XLA運行TensorFlow計算，或者通過JIT編譯算子放置在CPU或GPU的設備上，或通過將算子在XLA_CPU或XLA_GPUTensorFlow設備。將算子直接放在TensorFlow XLA設備上強制算子在該設備上運行，主要用於測試。
Note: The XLA CPU backend produces fast single-threaded code (in most cases), but does not yet parallelize as well as the TensorFlow CPU backend. The XLA GPU backend is competitive with the standard TensorFlow implementation, sometimes faster, sometimes slower.
打開JIT編譯
JIT編譯可以在會話級別，打開或手動進行選擇操作。這兩種方法都是零拷貝---在編譯的XLA內核和置於同一設備上的TensorFlow算子之間傳遞數據時，不需要復制數據。
Session
在會話級別打開JIT編譯，導致所有可能的算子，貪婪地編譯成XLA計算。每個XLA計算，編譯為一個或多個內核設備。
受限於一些限制，如果圖中有兩個相鄰的算子，都具有XLA實現，編譯為單個XLA計算。
JIT編譯在會話級別打開，方法是在會話初始化期間，將config 設置global_jit_level為tf.OptimizerOptions.ON_1，傳遞配置。
# Config to turn on JIT compilation
config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(config=config)
Note: Turning on JIT at the session level will not result in operations being compiled for the CPU. JIT compilation for CPU operations must be done via the manual method documented below. This decision was made due to the CPU backend being single-threaded.
手動
JIT編譯也可以為一個或多個算子手動打開。這是通過標記算子以使用屬性進行編譯完成的_XlaCompile=true。最簡單的方法是通過在中tf.contrib.compiler.jit.experimental_jit_scope()定義的范圍tensorflow/contrib/compiler/jit.py。用法示例：
jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
x = tf.placeholder(np.float32)
with jit_scope():
y = tf.add(x, x) # The "add" will be compiled with XLA.
該_XlaCompile屬性目前支持盡力而為。如果無法編譯算子， TensorFlow將默默回退到正常實現。
將算子放置在XLA設備上
通過XLA運行計算的另一種方法，將算子放置在特定的XLA設備上。通常僅用於測試。有效目標是XLA_CPU或XLA_GPU。
with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
output = tf.add(input1, input2)
與標准CPU和GPU設備上的JIT編譯不同，這些設備在將數據傳輸到設備上，或從設備傳輸時將復制數據。額外的副本使XLA和TensorFlow算子，在同一個圖中混合成本很高。
如何在開啟JIT的情況下，訓練MNIST softmax。當前在會話級別的JIT，僅支持GPU。
驗證LD_LIBRARY環境變量或ldconfig，包含$CUDA_ROOT/extras/CUPTI/lib64，其中包含CUDA分析工具界面（CUPTI）的庫。TensorFlow使用CUPTI從GPU中，提取跟蹤信息。
JIT-Just In Time Compilation
先來看XLA如何作用於TensorFlow的計算圖，下面是一張簡單的TensorFlow計算圖。
XLA通過一個TensorFlow的圖優化Pass(MarkForCompilation)，在TensorFlow計算圖中，找到適合JIT編譯的區域。假設XLA僅支持MatMul和Add。
TensorFlow XLA把這個區域定義為一個Cluster，作為一個獨立的JIT編譯單元，在TensorFlow計算圖中通過Node Attribute標示。
然后，另一個TensorFlow的圖優化Pass(EncapsulateSubgraphs)，把cluster轉化成TensorFlow的一個Function子圖。在原圖上用一個Caller節點表示這個Function在原圖的位置。
最后，調用TensorFlow的圖優化Pass(BuildXlaOps)，把Function節點轉化成特殊的Xla節點。
在TensorFlow運行時，運行到XlaCompile時，編譯Xla cluster子圖，然后把編譯完的Executable可執行文件，通過XlaExecutableClosure傳給XlaRun運行。
TF2XLA
TensorFlow運行到XlaCompile節點時。
為了編譯這個Function，通過把TensorFlow子圖所有的節點，翻譯成XLA HLO Instruction虛擬指令的形式表達，整個子圖也由此轉化成XLA HLO Computation。
XLA-HLO
XLA在HLO的圖表達上進行圖優化。聚合可在同一個GPU Kernel中執行的HLO指令。
HLO圖優化前
HLO圖優化后
代碼生成
首先根據虛擬指令分配GPU Stream和顯存。
然后IrEmitter把HLO Graph轉化成由編譯器的中間表達LLVM IR表示的GPU Kernel。LLVM IR如下所示：
source_filename = "cluster_36__XlaCompiledKernel_true__XlaNumConstantArgs_1__XlaNumResourceArgs_0_.36"target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"target triple = "nvptx64-nvidia-cuda" @0 = private unnamed_addr constant [4 x i8] zeroinitializer@1 = private unnamed_addr constant [4 x i8] zeroinitializer@2 = private unnamed_addr constant [4 x i8] zeroinitializer@3 = private unnamed_addr constant [4 x i8] zeroinitializer@4 = private unnamed_addr constant [4 x i8] zeroinitializer@5 = private unnamed_addr constant [4 x i8] zeroinitializer@6 = private unnamed_addr constant [4 x i8] zeroinitializer define void @fusion_1(i8* align 16 dereferenceable(3564544) %alloc2, i8* align 64 dereferenceable(3776) %temp_buf) {entry: %output.invar_address = alloca i64 %output_y.invar_address = alloca i64 %arg0.1.raw = getelementptr inbounds i8, i8* %alloc2, i64 0 %arg0.1.typed = bitcast i8* %arg0.1.raw to [944 x [944 x float]]* %fusion.1.raw = getelementptr inbounds i8, i8* %temp_buf, i64 0 %fusion.1.typed = bitcast i8* %fusion.1.raw to [944 x float]* %0 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x(), !range !4 %thread.id.x = sext i32 %0 to i64 %thread.x = urem i64 %thread.id.x, 944 %thread.y = udiv i64 %thread.id.x, 944 %1 = alloca float %partial_reduction_result.0 = alloca float %2 = load float, float* bitcast ([4 x i8]* @0 to float*) %3 = getelementptr inbounds float, float* %partial_reduction_result.0, i32 0 store float %2, float* %3 %current_output_linear_index_address = alloca i64 %4 = alloca i1 store i1 false, i1* %4 %5 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x(), !range !5 %block.id.x = sext i32 %5 to i64 %6 = udiv i64 %block.id.x, 1 %7 = urem i64 %6, 1 %8 = udiv i64 %block.id.x, 1 %9 = urem i64 %8, 8 %10 = udiv i64 %block.id.x, 8 %block_origin.0 = mul i64 %10, 1 %block_origin.1 = mul i64 %9, 1 ...
由LLVM生成nvPTX（Nvidia定義的虛擬底層指令表達形式）表達，進而由NVCC生成CuBin可執行代碼。PTX如下所示:
.reg .f32 %f<25>; .reg .b32 %r<31>; .reg .b64 %rd<61>; ld.param.u64 %rd27, [fusion_1_param_0]; ld.param.u64 %rd28, [fusion_1_param_1]; cvta.to.global.u64 %rd29, %rd28; cvta.to.global.u64 %rd1, %rd27; cvta.global.u64 %rd2, %rd29; mov.u32 %r3, %tid.x; cvt.u64.u32 %rd3, %r3; mov.u32 %r1, %ctaid.x; setp.eq.s32 %p1, %r1, 7; @%p1 bra LBB0_4; bra.uni LBB0_1;LBB0_4: selp.b64 %rd4, 48, 128, %p1; cvt.u32.u64 %r26, %rd3; shl.b64 %rd47, %rd3, 2; add.s64 %rd48, %rd47, %rd1; add.s64 %rd59, %rd48, 3383296; or.b32 %r27, %r26, 845824; mul.wide.u32 %rd49, %r27, 582368447; shr.u64 %rd50, %rd49, 39; cvt.u32.u64 %r28, %rd50; mul.lo.s32 %r29, %r28, 945; sub.s32 %r2, %r27, %r29; mov.f32 %f23, 0f00000000;...
代碼執行
當TensorFlow運行到XlaRun時，運行由XlaCompile編譯得到的GPU可執行代碼（Cubin或PTX）。
參考鏈接：
https://www.cnblogs.com/wujianming-110117/p/14948227.html
https://zhuanlan.zhihu.com/p/98565435