JIT in MegEngine

旷视天元（MegEngine）是一个深度学习框架，它主要包含训练和推理两方面内容。训练侧一般使用Python搭建网络；而推理侧考虑到产品性能的因素，一般使用C++语言集成天元框架。无论在训练侧还是推理侧，天元都担负着将训练和推理的代码运行到各种计算后端上的任务。目前天元支持的计算后端有CPU、GPU、ARM和一些领域专用的加速器，覆盖了云、端、芯等各个场景。

天元主要有三大特征：

1.训推一体，不管是训练任务还是推理任务都可以由天元一个框架来完成。

2.动静结合，天元同时支持动态图和静态图，并且动静之间的转换也非常方便。

3.多平台的高性能支持。

图1.天元架构

如图1所示，我们可以看到天元提供了Python和C++两种接口。在图表示上分为动态图和静态图。运算层组件包括自动求导器、图优化和图编译等。天元的运行时模块包括内存管理和计算调度，其中内存管理包括静态内存管理和动态内存管理，以及亚线性内存优化技术。计算内核层包含了天元支持的所有计算后端，我们后续会开源出更多的计算后端。除此之外，天元还包含了一个高性能异构通信库，它一般会在多机多卡的场景下被用到。

图2.计算图

动态图和静态图是相对的，在动态图下是没有计算图的概念的。但在静态图下，天元会维护一张计算图。如图2所示为天元中的计算图表示，图中圆形表示算子（operator），三角形表示输入。在天元框架中，动态图和静态图之间的转换只需要一条简单的语句即可完成，如下代码所示：

动态图和静态图的转换

if__name__=='__main__’:gm=().attach(())opt=((),lr=0.0125,momentum=0.9,weight_decay=1e-4)#通过trace转换为静态图@trace(symbolic=True)deftrain():withgm:logits=model(image)loss=_entropy(logits,label)(loss)()_grad()returnlossloss=train()()

复制代码

AOT(AheadOfTime)和JIT(JustInTime)都是编译中的概念。以传统的C/C++语言为例，我们写完代码之后，一般会通过编译器编译生成可执行文件，然后再执行该可执行文件获得执行结果。如果我们将从源代码编译生成可执行文件的过程称为build阶段，将执行可执行文件叫做runtime阶段的话，JIT是没有build阶段的，它只有runtime阶段。JIT一般被用在解释执行的语言如Python中，JIT会在代码执行的过程中检测热点函数，随后对热点函数进行重编译，下次运行时遇到热点函数则直接执行编译结果即可。这样做可以显著加快代码执行的速度。

随着各种编程语言的出现，现代编译器也日趋多样化。特别是近年来随着深度学习的兴起，深度学习软件框架和AI领域专用硬件呈爆发式增长。不断增加的软件框架和AI硬件之间逐渐形成了一个越来越大的沟壑，如何将框架层对深度学习模型的描述精准高效的翻译成适应各类硬件的语言成为难点。MLIR(Multi-LevelIntermediateRepresentation)是一种可以在统一的基础架构下满足多样化需求的混合IR。MLIR可以满足包括但不限于以下的需求：

1.表达数据流图（如静态图模式下的MegEngineGraph）

2.表达对该图做的优化和变换操作

3.进行各种算子优化如算子融合（kernelfusion）、循环融合、算子分块和内存格式（memorylayout）转换等

4.自动代码生成、显式缓存管理、自动向量化

作为一个公用的IR，MLIR具有非常优秀的表达能力和可扩展性。MLIR可以表达图层面的运算，同时可以表达传统编译器中的IR信息，也可以表示硬件专用的运算。这种不同属性，不同类型的运算的集合构成了MLIR中的方言（Dialect）。MLIR还提供方便的机制实现不同方言之间的转换（LoweringDown），因此MLIR的一个通用优化将会在多个方面产生收益。接入MLIR也将有更大可能享受到它的生态好处，包括性能和扩展性等方面。

众所周知，深度学习模型中有很多element-wise操作，例如加减乘除算术运算和神经网络中的激活函数一般都是element-wise操作。天元将element-wise操作分为一元操作、二元操作和多元操作。一元操作主要有RELU、ABS、SIN和COS等等；二元操作有加法、减法、乘法和除法以及MAX等；多元操作有FUSE-MUL-ADD3和FUSE-MUL-ADD4等，它们分别计算的是“ab+c”以及“ab+c*d”。

表1卷积神经网络中的element-wise操作

model

batchsize

element-wisecomputation/totalcomputation(%)

element-wise****time/totaltime(%)

resnet50

1

0.102054232

4.6

8

0.102085366

10.8

16

0.102080177

11.9

mobilenetV2

1

0.703333333

4.1

8

0.704592902

8.9

16

0.704627697

11.9

vgg16

1

0.02927408

5.8

8

0.029277172

7.1

16

0.029342568

9.4

element-wise操作在卷积神经网络中所占的地位不可忽视。如表1所示，我们选择公开的卷积神经网络训练模型，以纯devicekernel的执行时间为基准统计卷积神经网络中的element-wise操作的重要性。

首先可以清晰的看到，element-wise的计算量的占比相比于运行时间占比要低1-2个数量级。它的计算量占的非常少，但是它的运行时间占比非常多，这个结论是比较反直觉的。并且随着batchsize的增加，这个现象也越来越明显。这是因为element-wise操作计算量较低但是访存量较高，即计算访存比较低，是一种典型的访存受限(memorybound)的操作。以“a+b”为例，我们首先要将a读到内存中，再将b读到内存中，做完一次加法之后，我们将结果c再写到内存中。整个过程要经过两次读和一次写才能完成一次计算，所以它的计算反应访存比非常低。针对访存受限的操作，优化计算时间实际上是没有没有太多的意义的，而应该集中精力优化访存，访存优化的常见的优化手段是融合(fusion)。如果我们能将网络中连在一起的element-wise操作融合成一个算子，则将减少element-wise操作的访存量，增加计算访存比从而加速网络的整体性能。

卷积神经网络有两个鲜明的特征。一个是静态图模式下的模型训练过程中模型的结构一般是不会变的跑；另一个是在模型训练的过程中，一般会经过很多个iter/min-batch，不同的iter/min-batch之间输入张量形状（tensorshape）一般也不会变。基于卷积神经网络的这两个特征，我们决定应用JIT技术，原因如下：

1.只需要在首次运行的时候编译一次，随后的不同iter/mini-batch可以重用第一次编译出来的结果

2.JIT具有较强的可移植性，因为它在运行时获取平台信息，然后生成可以在该平台运行的代码

3.JIT可以解决element-wise模式组合爆炸的问题

我们通过Element-wiseFusion可以把多个element-wise操作融合成一个，减少了算子数量也就减少了算子之间的读写次数。如图3所示计算图算的是“a*b+c”，它需要4次读，2次写。4次读分别是乘法在读a和b两个输入，乘法其实还要写一个隐藏的输出，加法会读乘法的输出作为输入，以及加法读c作为输入。两次写分别是乘法和加法对它们结果的两次写操作，总共加起来是4次读，2次写。

我们将其融合成一个算子FUSE_MUL_ADD3，由于天元现在已经支持FUSE_MUL_ADD3这个element-wise模式，所以我们可以直接做模型手术将计算图从图3左侧形式转到图3右侧形式。对于融合之后的计算图，我们只需要3次读和1次写就可以完成等价计算，相比于融合前减少了1次读和1次写操作。

图3融合优化减少访存次数

我们无法预测用户将搭出来怎样的一张计算图，考虑图4所示的计算图，其中element-wise的个数和顺序都不固定，显然我们不可能提前将各种element-wise模式的组合都写进天元的。在这种情况下，天元会创建一个虚拟的算子来表示整个可被融合的子图。有了虚拟算子的存在，接下来我们还要解决两个问题，一个是用虚拟算子替换原始计算图中可以被融合的子图，这个工作会在图优化阶段做；另一个是我们要动态生成虚拟算子的代码并执行。如果我们解决了这两个问题，我们就解决了整个问题。

图4子图融合优化

为了将一张计算图中的可被融合的子图融合成一个算子，天元将进行检测（detection）和融合（fusion）两步操作，如下步骤1-3属于检测，步骤4则属于融合：

1.对原始计算图进行检测后生成internalgraphgenerator，一个internalgraphgenerator对应一个唯一的子图

2.internalgraphgenerator稍后会生成internalgraph

3.由internalgraph创建JITExcutor算子

4.将JITExcutor写回原始的计算图

检测算法的主要功能是找出可以被融合的子图。为了方便描述，设G是计算图，opr是图G中的算子，var是opr的输入和输出。检测算法的输入是原始的计算图G，输出是一个哈希表M，表中存放的是检测出的可被融合子图的输出var（记作point）与其对应的internalgraphgenerator。算法步骤如下：

1.按照逆拓扑序列遍历图G中的算子opr

2.如果opr不是Elemwise/PowC/TypeCvt/Reduce/Dimshuffle/JITExecutor，返回步骤1

3.如果opr的input/output数据类型不是float32/float16，返回步骤1

4.process_opr(opr)

5.转到步骤1

图5process_opr流程图

拓扑序列要求所有的父节点要先于它的子节点被访问到，与之对应的，逆拓扑序列就是所有的子节点要先于它的父节点被访问到。算法第1步中我们之所以按照逆拓扑序列遍历计算图，是因为要保证遍历到某个opr时，它的子节点都已经被遍历到了。这样算法可以查看该opr的所有的子节点是不是都在同一张子图中，如果是，那么当前opr就有很大的可能也在该子图中。算法的第2步和第3步实际上说明了天元中的JIT的限制。目前天元JIT仅支持Elemwise/PowC/TypeCvt/Reduce/Dimshuffle这几种opr，而且只支持输入输出是float32/float16的数据类型。第4步详细流程如图5所示。需要注意的是算法会经过如下三个判断语句：

1.该opr的子节点是不是都已经在当前的这张子图中了？

2.该opr的输出的计算节点（computenode）是不是跟子图匹配？天元支持跨计算节点的计算图，例如计算图中一些opr可以运行在CPU上，一些opr可以运行在GPU上。但目前天元不支持跨计算节点融合。

3.该opr的输出的shape是不是跟子图匹配？因为最终生成的代码本质上是一个大的循环，循环的维度就是opr输出的shape，所以如果shape不匹配是不能被融合的。

图6检测算法检测出的可被融合的子图

图6中虚线框出来的即为检测算法检测出的两个可被融合的子图。

融合算法的主要功能是将检测出来的子图融合成一个算子。融合算法的输入是原始的计算图和检测算法输出的那张哈希表M，它的输出是经过融合的计算图G‘。算法流程如下：

1.按照拓扑序列遍历图G中的算子opr

2.若opr的输入var不是point,返回步骤1

3.从M中拿到var对应的internalgraphgenerator,生成internalgraph

4.从internalgraph创建JITExecutor

5.写回原始的计算图G

6.转到步骤1

步骤2中如果一个opr的输入var不是point则表示它是一个子图中的中间节点而不是子图的输出节点。步骤3中从internalgraphgenerator到internalgraph需要将子图的输入var替换为一个新的oprJITPlaceholder。JITPlaceholder中会存诸如子图的输入顺序这些额外信息，因为某些element-wise操作是对输入顺序敏感的。例如a对b取余和b对a取余显然具有不同的语义。

图7融合后的计算图

图7即为经过融合算法之后的计算图，截止到目前为止，我们已经完成了图优化方面的所有工作。

经过图优化之后，我们成功的将计算图中可被融合的子图融合成为一个新的算子，剩下的工作就是为这个新的算子生成代码了。JITExecutor算子的运行时代码非常简单，先判断一下当前的可执行对象是不是已经存在，如果不存在则先编译出一个可执行对象，如已存在则直接运行。这段代码在运行时才会被执行到，所以称之为JIT。当前天元支持三种JIT编译器后端，分别是NVRTC（支持英伟达GPU），Halide和MLIR。其中后两个编译后端支持的平台众多，但是MLIR具有更优秀的表达能力和扩展性，所以我们接下来以MLIR为例介绍代码生成、编译和执行的过程。

要想使用MLIR作为编译后端，首先我们需要定义和实现天元自己的方言（MGEDialect），随后我们将MGEDialect转换到MLIR既有的Dialect上，接下来的绝大部分工作都可以复用MLIR中的基础组件和工具完成。图8描述了CPU和GPU上大概的执行流程。

图8JIT编译器工作流

天元首先将JITExecutor算子内部的internalgraph翻译成MGEDialect。在CPU上，MGEDialect会先Lowering到AffineDialect上，然后会通过LLVM的组件Lowering到LLVMDialect上，LLVMDialect可以被直接翻译成LLVMIR。在这一步之后，其他优化工作都可以直接复用LLVM的基础组件。最后天元使用MLIRExecutionEngine执行LLVMIR生成的代码。在GPU上，天元会先将MGEDialectLowering到GPUDialect上，随后Lowering到NVVMDialect，NVVM会被翻译成PTX汇编。最后通过英伟达提供的CUmodule和CUfunction两个机制运行。

首先参考这篇文档在天元中开启JIT支持。本次实验选了resnet50,mobilenetV2和vgg16三个业界广泛使用的模型，batchsize分别设置了1,8和16。测试硬件环境为NVIDIAT4，软件环境为。

图9打开JIT相比于不开JIT的加速比

由图9可知，和不打开JIT支持相比，打开JIT支持后resnet50最高可以获得16%的加速比，mobilenetV2则能获得6%到7%的加速比，而vgg16其实上没有明显加速效果。这是因为vgg16模型很大，可以被优化的element-wise操作比较少。JIT的优化效果跟具体的模型是有紧密关系的。

图10JIT编译耗时

如果打开了JIT支持，那么天元首次运行的时候会有一次JIT编译的过程。JIT编译耗时跟具体的编译的后端以及模型有关，如图10所示resnet50耗时2.7毫秒，mobilenetV2耗时3.9毫秒。

本篇文章介绍了天元使用JIT实现将任意多个相邻的element-wise算子融合成一个算子的优化。我们在T4上用实验，相比于优化前，resnet50最高可以获得16%的加速比。

以此为基，展望未来我们可能做的事情如下：

1.将JIT编译的结果先离线保存，线上直接将线下编译好的可执行对象读进内存。这种做法可以解决线上第一次运行慢的问题，但它可能会损失一部分可移植性，因为在一种设备上编译产生的可执行对象一般不能适配所有线上设备。

2.JIT支持更多的算子

3.JIT支持更多的数据类型，天元JIT优化暂时只支持float32/float16这两种数据类型。

4.动态图JIT，也就是传统意义上的检测热点代码，重编译后再执行。