jit(上):tensorflow如何实现即时编译?

author author     2023-03-20     260

关键词:

参考技术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提供。

CUDA 表达式模板和即时编译 (JIT)

...对使用CUDA进行即时(JIT)编译有一些疑问。我已经根据论文实现了一个基于表达式模板的库J.M.Cohen,“使用C++元编程处理设备阵列”,GPU计算宝石-Jade版它似乎工作得相当好。如果我比较矩阵元素运算的计 查看详情

jit即时编译(史上最全)(代码片段)

...后面,水不知道有多深。首先,考察了大家对JVM即时编译组件的 查看详情

javajava即时编译(jit)器原理解析及实践

...代码直接编译成CPU所能理解的机器码来运行。而Java为了实现“一次编译,处处运行”的特性,把编译的过程分成两部分,首先它会先由javac编译成通用的中间形式——字节码,然后 查看详情

jit即时编译

...oogle+用电子邮件发送本页面JIT是justintime的缩写,也就是即时编译编译器。使用即时编译器技术,能够加速Java程序的执行速度。下面,就对该编译器技术做个简单的讲解。首先,我们大家都知道,通常通过javac将... 查看详情

jit即时编译

...文件   class文件有2种方式执行:解释执行、JIT即时编译后执行    通常情况下默认解释器边解释边运行,但对于一些热点代码会首先编译为机器码,缓存起来,后续执行执行缓存的机器码即可,提升执行... 查看详情

jvm虚拟机aot和jit即时编译(代码片段)

文章目录JVM虚拟机AOT和JIT即时编译1、AOT2、JITJVM虚拟机AOT和JIT即时编译1、AOTAOT是AheadOfTime,运行前编译与JIT编译对比JIT吞吐量高,有运行时性能加成,可以跑得更快,并可以做到动态生成代码等,但是相对启动... 查看详情

学习笔记tf062:tensorflow线性代数编译框架xla

XLA(AcceleratedLinearAlgebra),线性代数领域专用编译器(demain-specificcompiler),优化TensorFlow计算。即时(just-in-time,JIT)编译或提前(ahead-of-time,AOT)编译实现XLA,有助于硬件加速。XLA还在试验阶段。https://www.tensorflow.org/versions/master/exper 查看详情

java即时编译jit

...可执行的代码(原生型指令码)JITCompiler(Just-in-timeCompiler)即时编译最早的Java建置方案是由一套转译程式(interpreter),将每个Java指令都转译成对等的微处理器指令,并根据转译后的指令先后次序依序执行,由于一个Java指令可能... 查看详情

转:什么是即时编译(jit)!?openjdkhotspotvm剖析

重点应用程序可以选择一个适当的即时编译器来进行接近机器级的性能优化。分层编译由五层编译构成。分层编译提供了极好的启动性能,并指导编译的下一层编译器提供高性能优化。提供即时编译相关诊断信息的JVM开关。像内... 查看详情

java虚拟机即时编译器技术jit

...术用以提升速度,尤其是与加载器操作相关的,被称为“即时”(Just-In-Time,JIT)编译器的技术。这种技术可以把程序全部或部分翻译成本地机器码(这本来是JVM的工作),程序运行速度因此得以提升。当需要装载某个类时,编译器... 查看详情

Jax、jit 和动态形状:Tensorflow 的回归?

】Jax、jit和动态形状:Tensorflow的回归?【英文标题】:Jax,jitanddynamicshapes:aregressionfromTensorflow?【发布时间】:2021-06-1701:41:48【问题描述】:documentationforJAX说,并非所有JAX代码都可以JIT编译,因为它要求数组形状是静态的并且在... 查看详情

如何通过创建反应应用程序启用 JIT(即时模式)?

】如何通过创建反应应用程序启用JIT(即时模式)?【英文标题】:HowtoenableJIT(Justintimemode)withcreatereactapp?【发布时间】:2021-10-0717:04:40【问题描述】:我尝试自己在create-react-app中设置JIT,但它似乎没有工作,因为样式没有更新... 查看详情

jvm--运行期优化;jit(代码片段)

...拟机标准对JIT的存在没有作出任何规范,这是虚拟机实现的自定义优化技术。HotSpot虚拟机的执行引擎在执行J 查看详情

jvm--运行期优化;jit(代码片段)

...拟机标准对JIT的存在没有作出任何规范,这是虚拟机实现的自定义优化技术。HotSpot虚拟机的执行引擎在执行J 查看详情

juc并发编程--jit即时编译器之锁清除(代码片段)

1.1锁消除导入相关依赖:因为JMH是JDK9自带的,如果是JDK9之前的版本需要加入如下依赖<dependency><groupId>org.openjdk.jmh</groupId><artifactId>jmh-core</artifactId><version>1.23</version>&l 查看详情

重新编译 JIT 函数时的 JAX 挂钩/信息/警告

...使用一个hacky解决方法来获得重新编译的通知。在当前的实现中,跟踪器在需要编译时执 查看详情

什么是jit?(代码片段)

...更常见的是字节码到机器码的翻译,然后直接执行。实现JIT编译器的系统通常会持续分析正在执行的代码,并识别代码中从 查看详情

java之jit(justintime)

...码,并进行各种层次的优化,完成这个任务的编译器称为即时编译器(JIT编译器,不是Java虚拟机内必须的部分)。要了解HotSpot虚拟机内的即时编译器的运作过程,要解决几个 查看详情