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