JIT（上）：Tensorflow如何实现即时编译？

Posted 2023-03-20

参考技术A Tensorflow的JIT（just-in-time）是指在运行 @tf.function 修饰的python函数时，由 jit 、 tf2xla 和 XLA 一起完成一系列如子图构造、子图优化、图编译和图执行等操作。编译后的可执行程序-- executable 会存放到cache中，供再次调用时直接获取执行。JIT的好处在 开篇 已经讲过了，这里不再赘述。

JIT的流程可以概括为：Tensorflow子图构造/优化，graph -> HLO，编译/执行，合并计算结果到Tensorflow图这四部分。本文只涉及图编译和图执行。

函数中ops在子图构造阶段被包裹进一个cluster node，并替换成 xla_compile 和 xla_run 这两op，而 XlaCompileOp 和 XlaRunOp 就是它们的 OpKernel ，分别用于图编译和执行。

XlaCompileOp通过 XlaCompilationCache 获取或编译executable，并将其封装成 XlaExecutableClosure ，并缓存在 XlaExecutableClosureStore 。 XlaRunOp 用从XlaCompileOp传递来的key在cache中查找并执行executable。

从编译流程图可以看到，XLA的编译结果会缓存到XlaCompilationCache，后续调用可以根据 signature 在cache中查找executable。

函数的 signature 是由 BuildSignature(function, args) 根据函数和arguments生成的。即使是同一个函数，只要input tensors不同，signature也会不一样，这就是 power() 被编译两次的原因：第三次函数调用时，由于无法通过signature在cache中找到executable而触发编译。

signature表示唯一的计算图：只要函数中的ops序列和arguments（type/shape）是确定的，那么计算图也是确定的。

编译之前需要通过 tf2xla 将图转换成XLA支持的语言 HLO 。tf2xla为每个Tensorflow op创建了生成HLO的 XlaOp ，因此，只要执行该Tensorflow子图，就可以生成具有相同的拓扑排序的HLO -- XlaComputation 。

XlaComputation （HLO）可以认为是一个运行在device上的纯函数，它的input/output会伴随着host-to-device（H2D）和device-to-host（D2H）的数据传输。

我们知道，Tensorflow图中的input tensor有两种： tf.Placeholder 和 tf.Variable ，前者每个step都会将新data发送到device，而后者是模型参数，它们会常驻内存，只在store/load checkpoint才会有H2D/D2H。

而纯函数的定义是：

不管是input还是output，虽然variable和其他argument一样存在于HLO的参数列表和返回值列表中，但它们实际上是常驻于device的，不需要也不应该H2D/D2H。

因此，HLO在编译时还需要通过 argument_input_indices 、 resource_input_indices 和 resource_update_to_input_index 等options来区分arguments和variables。

此外，如果有input是常数，为了避免无谓的H2D开销，可以把它固化到函数内部。同理，对于常数output，它没必要出现在函数中，可以直接定义在 XlaCompilationResult 的output buffer。

XlaCompilationResult是 Graph -> HLO 的output，它封装了HLO以及上述部分metadata、buffers。

XlaCompileOp会把编译好的executable、metadata、input/output buffers、options等统统封装进一个closure -- XlaExecutableClosure ，并将其缓存在 XlaExecutableClosureStore 供 XlaRunOp 获取。

XlaRunOp 可以通过一个数字字符串key（从0开始累加）从cache中查找并执行XlaExecutableClosure，这个key由XlaCompileOp提供。

浅析 JIT 即时编译技术

即时编译回顾

HotSpot 虚拟机执行 Java 程序时，先通过解释器对代码解释执行，发现某个方法或代码块执行比较频繁后，对热点代码进行编译，编译后生成与本地平台相关的机器码，再去执行机器码获得较高的运行效率。必要时，也会通过逆优化从即时编译回到解释执行，如编译器遇到罕见陷阱的情况。

 

 在 Java 虚拟机规范中，并未要求虚拟机必须实现即时编译，但即时编译在主流的虚拟机中都有实现，后文所有讨论都基于 HotSpot 虚拟机。

 

即时编译器

HotSpot 虚拟机中有两个即时编译器，分别为 Client Compiler 和 Server Compiler，简称为 C1 编译器和 C2 编译器。C1 编译器占用内存小，进行局部优化，代码执行效率比 C2 编译器低，适用于客户端应用； C2 编译器占用内存大，进行全局优化，代码执行效率比 C1 编译器高，适用于服务端应用。在 JDK 1.7 以后的 HotSpot 虚拟机中，采用了被称为分层编译的策略，启动时解释执行提高启动速度，然后 C1 编译获取较好的编译速度，再 C2 编译获取较好的编译质量。

默认情况下，虚拟机会自动选择合适的编译器与解释器一起配合工作，用户也可以在启动参数里使用 “-client” 或 “-server” 来强制要求虚拟机使用某一种编译器。编译器与解释器一起配合工作的模式，被称为混合模式。用户也可以在启动参数里使用 “-Xint” 来要求虚拟机只使用解释器，使用 “-Xcomp” 来要求虚拟机只使用编译器。

被多次调用的方法和被多次执行的循环体，就是即时编译器关注的热点代码。在 HotSpot 中，采用了基于计数器的热点探测技术，为每个方法定义了两个计数器：方法调用计数器、回边计数器。

默认情况下，如果一段时间内方法调用计数器的值没有超过虚拟机设置的阈值，则在垃圾回收时计数器会热度衰减，数值减少 1/2，此时间范围又被称为半衰期。所以方法调用计数器中存储的实际上并不是方法调用的绝对次数。用户可以调整半衰期的值，甚至可以关闭热度衰减。

回边计数器则统计了循环体执行的绝对次数，它的阈值可以由方法调用计数器的阈值计算出来。如果回边计数器发生溢出，也会把方法调用计数器调整为溢出。

调用方法时，如果两个计数器之和超过了方法调用计数器的阈值，就会提交方法的编译请求。循环体执行时，如果两个计数器之和超过了回边计数器的阈值，也是编译代码块所在的方法，因为此时方法正在执行中，又被称为栈上替换，即替换方法时方法的栈帧还在栈上。

默认情况下，编译器在后台进行编译，即调用热点代码发现需要编译时，先以解释方式继续执行，向编译器发送编译请求，编译结束后下次调用时才使用编译的本地代码。用户也可以通过启动参数来要求虚拟机禁止后台编译，编译结束前热点代码的执行处于阻塞状态。

 

优化技术

HotSpot 虚拟机使用了很多种优化技术，这里只简单介绍其中的几种，完整的优化技术介绍可以参考官网内容。

公共子表达式消除：

如果一个表达式已经进行过计算，并且在下次用到之前依赖的变量没有变化，即表达式的计算结果不会发生变化，则在下次使用这个表达式时直接使用计算的结果。

数组边界检查消除：

在 Java 中访问数组时，会自动进行边界检查来防止数组下标越界。但是对于某些情况并不需要每次访问都去检查，如在一个循环中遍历数组元素，如果虚拟机能够确定下标不会发生越界并且优化确实能够提高运行速度，则虚拟机会去除每次访问的下标检查。

方法内联：

对于可以内联的方法，直接复制到调用者代码中，减少方法调用次数和性能消耗。

逃逸分析：

方法中定义的一个对象，如果会被其他方法访问则称为方法逃逸，如果会被其他线程访问则称为线程逃逸。对于不能逃逸的对象，HotSpot 虚拟机采用了栈上分配、同步消除、标量替换等方法进行优化。

 

每周 3 篇学习笔记或技术总结，面向有一定基础的 Java 程序员，内容涉及 Java 进阶、虚拟机、MySQL、NoSQL、分布式计算、开源框架等多个领域。关注作者或微信公众号 backend-develop 第一时间获取最新内容。

浅析 JIT 即时编译技术 | 后端开发那点事儿