本文整顿自 2 月 25 号 MegEngine 首期 Meetup 的分享。
作者：王彪 | 旷视框架部异构计算组工程师
视频回放：《JIT in MegEngine》_MegEngine Meetup 第一期

一、背景

什么是天元

旷视天元（MegEngine）是一个深度学习框架，它次要蕴含训练和推理两方面内容。训练侧个别应用 Python 搭建网络；而推理侧思考到产品性能的因素，个别应用 C++ 语言集成天元框架。无论在训练侧还是推理侧，天元都负担着将训练和推理的代码运行到各种计算后端上的工作。目前天元反对的计算后端有 CPU、GPU、ARM 和一些畛域专用的加速器，笼罩了云、端、芯等各个场景。

天元次要有三大特色：

1. 训推一体，不论是训练任务还是推理工作都能够由天元一个框架来实现。

2. 动静联合，天元同时反对动态图和动态图，并且动静之间的转换也十分不便。

3. 多平台的高性能反对。

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

动态图和动态图是绝对的，在动态图下是没有计算图的概念的。但在动态图下，天元会保护一张计算图。如图 2 所示为天元中的计算图示意，图中圆形示意算子（operator），三角形示意输出。在天元框架中，动态图和动态图之间的转换只须要一条简略的语句即可实现，如下方代码所示：

if __name__ == '__main__’: 
    gm = ad.GradManager().attach(model.parameters()) 
    opt = optim.SGD(model.parameters(), lr=0.0125, momentum=0.9, weight_decay=1e-4) 
   <em> # 通过 trace 转换为动态图 </em>
    @trace(symbolic=True) 
    def train(): 
        with gm: 
            logits = model(image) 
            loss = F.loss.cross_entropy(logits, label) 
            gm.backward(loss) 
        opt.step() 
        opt.clear_grad() 
        return loss 
    loss = train() 
    loss.numpy()

什么是 AOT 和 JIT

AOT(Ahead Of Time) 和 JIT(Just In Time) 都是编译中的概念。以传统的 C/C++ 语言为例，咱们写完代码之后，个别会通过编译器编译生成可执行文件，而后再执行该可执行文件取得执行后果。如果咱们将从源代码编译生成可执行文件的过程称为 build 阶段，将执行可执行文件叫做 runtime 阶段的话，JIT 是没有 build 阶段的，它只有 runtime 阶段。JIT 个别被用在解释执行的语言如 Python 中，JIT 会在代码执行的过程中检测热点函数，随后对热点函数进行重编译，下次运行时遇到热点函数则间接执行编译后果即可。这样做能够显著放慢代码执行的速度。

什么是 MLIR

随着各种编程语言的呈现，古代编译器也日趋多样化。特地是近年来随着深度学习的衰亡，深度学习软件框架和 AI 畛域专用硬件呈爆发式增长。一直减少的软件框架和 AI 硬件之间逐步造成了一个越来越大的沟壑，如何将框架层对深度学习模型的形容精准高效的翻译成适应各类硬件的语言成为难点。MLIR(Multi-Level Intermediate Representation) 是一种能够在对立的基础架构下满足多样化需要的混合 IR。MLIR 能够满足包含但不限于以下的需要：

1. 表白数据流图（如动态图模式下的 MegEngine Graph）

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

3. 进行各种算子优化如算子交融（kernel fusion）、循环交融、算子分块和内存格局（memory layout）转换等

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

作为一个专用的 IR，MLIR 具备十分优良的表达能力和可扩展性。MLIR 能够表白图层面的运算，同时能够表白传统编译器中的 IR 信息，也能够示意硬件专用的运算。这种不同属性，不同类型的运算的汇合形成了 MLIR 中的方言（Dialect）。MLIR 还提供方便的机制实现不同方言之间的转换（Lowering Down），因而 MLIR 的一个通用优化将会在多个方面产生收益。接入 MLIR 也将有更大可能享受到它的生态益处，包含性能和扩展性等方面。

二、动机

为什么做

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

element-wise 操作在卷积神经网络中所占的位置不可漠视。如表 1 所示，咱们抉择公开的卷积神经网络训练模型，以纯 device kernel 的执行工夫为基准统计卷积神经网络中 element-wise 操作的重要性。

首先能够清晰的看到，element-wise 的计算量的占比相比于运行工夫占比要低 1-2 个数量级。它的计算量占的非常少，然而它的运行工夫占比十分多，这个论断是比拟反直觉的。并且随着 batch size 的减少，这个景象也越来越显著。这是因为 element-wise 操作计算量较低然而访存量较高，即计算访存比拟低，是一种典型的访存受限 (memory bound) 的操作。以“a+b”为例，咱们首先要将 a 读到内存中，再将 b 读到内存中，做完一次加法之后，咱们将后果 c 再写到内存中。整个过程要通过两次读和一次写能力实现一次计算，所以它的计算反馈访存比非常低。针对访存受限的操作，优化计算工夫实际上是没有没有太多的意义的，而应该集中精力优化访存，访存优化的常见的优化伎俩是交融 (fusion)。如果咱们能将网络中连在一起的 element-wise 操作交融成一个算子，则将缩小 element-wise 操作的访存量，减少计算访存比从而减速网络的整体性能。

为什么用 JIT 做

卷积神经网络有两个显明的特色。一个是动态图模式下的模型训练过程中模型的构造个别是不会变的跑；另一个是在模型训练的过程中，个别会通过很多个 iter/min-batch，不同的 iter/min-batch 之间输出张量形态（tensor shape）个别也不会变。基于卷积神经网络的这两个特色，咱们决定利用 JIT 技术，起因如下：

1. 只须要在首次运行的时候编译一次，随后的不同 iter/mini-batch 能够重用第一次编译进去的后果。

2.JIT 具备较强的可移植性，因为它在运行时获取平台信息，而后生成能够在该平台运行的代码。

3.JIT 能够解决 element-wise 模式组合爆炸的问题。

三、技术计划

咱们通过 Element-wise Fusion 能够把多个 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 次写操作。

咱们无奈预测用户将搭进去怎么的一张计算图，思考图 4 所示的计算图，其中 element-wise 的个数和程序都不固定，显然咱们不可能提前将各种 element-wise 模式的组合都写进天元的。在这种状况下，天元会创立一个虚构的算子来示意整个可被交融的子图。有了虚构算子的存在，接下来咱们还要解决两个问题，一个是用虚构算子替换原始计算图中能够被交融的子图，这个工作会在图优化阶段做；另一个是咱们要动静生成虚构算子的代码并执行。如果咱们解决了这两个问题，咱们就解决了整个问题。

图优化

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

1. 对原始计算图进行检测后生成 internal graph generator，一个 internal graph generator 对应一个惟一的子图

2.internal graph generator 稍后会生成 internal graph

3. 由 internal graph 创立 JITExcutor 算子

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

检测

检测算法的次要性能是找出能够被交融的子图。为了不便形容，设 G 是计算图，opr 是图 G 中的算子，var 是 opr 的输出和输入。检测算法的输出是原始的计算图 G，输入是一个哈希表 M，表中寄存的是检测出的可被交融子图的输入 var（记作 endpoint）与其对应的 internal graph generator。算法步骤如下：

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

拓扑序列要求所有的父节点要先于它的子节点被拜访到，与之对应的，逆拓扑序列就是所有的子节点要先于它的父节点被拜访到。算法第 1 步中咱们之所以依照逆拓扑序列遍历计算图，是因为要保障遍历到某个 opr 时，它的子节点都曾经被遍历到了。这样算法能够查看该 opr 的所有的子节点是不是都在同一张子图中，如果是，那么以后 opr 就有很大的可能也在该子图中。算法的第 2 步和第 3 步实际上阐明了天元中的 JIT 的限度。目前天元 JIT 仅反对 Elemwise/PowC/TypeCvt/Reduce/Dimshuffle 这几种 opr，而且只反对输入输出是 float32/float16 的数据类型。第 4 步具体流程如图 5 所示。须要留神的是算法会通过如下三个判断语句：

1. 该 opr 的子节点是不是都曾经在以后的这张子图中了？

2. 该 opr 的输入的计算节点（compute node）是不是跟子图匹配？天元反对跨计算节点的计算图，例如计算图中一些 opr 能够运行在 CPU 上，一些 opr 能够运行在 GPU 上。但目前天元不反对跨计算节点交融。

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

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

交融

交融算法的次要性能是将检测进去的子图交融成一个算子。交融算法的输出是原始的计算图和检测算法输入的那张哈希表 M，它的输入是通过交融的计算图 G‘。算法流程如下：

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

2. 若 opr 的输出 var 不是 endpoint, 返回步骤 1

3. 从 M 中拿到 var 对应的 internal graph generator, 生成 internal graph

4. 从 internal graph 创立 JITExecutor

5. 写回原始的计算图 G

6. 转到步骤 1 步骤 2 中如果一个 opr 的输出 var 不是 endpoint 则示意它是一个子图中的两头节点而不是子图的输入节点。步骤 3 中从 internal graph generator 到 internal graph 须要将子图的输出 var 替换为一个新的 opr JITPlaceholder。JITPlaceholder 中会存诸如子图的输出程序这些额定信息，因为某些 element-wise 操作是对输出程序敏感的。例如 a 对 b 取余和 b 对 a 取余显然具备不同的语义。

图 7 即为通过交融算法之后的计算图，截止到目前为止，咱们曾经实现了图优化方面的所有工作。

图编译

通过图优化之后，咱们胜利的将计算图中可被交融的子图交融成为一个新的算子，剩下的工作就是为这个新的算子生成代码了。JITExecutor 算子的运行时代码非常简单，先判断一下以后的可执行对象是不是曾经存在，如果不存在则先编译出一个可执行对象，如已存在则间接运行。这段代码在运行时才会被执行到，所以称之为 JIT。以后天元反对三种 JIT 编译器后端，别离是 NVRTC（反对英伟达 GPU），Halide 和 MLIR。其中后两个编译后端反对的平台泛滥，然而 MLIR 具备更优良的表达能力和扩展性，所以咱们接下来以 MLIR 为例介绍代码生成、编译和执行的过程。

要想应用 MLIR 作为编译后端，首先咱们须要定义和实现天元本人的方言（MGE Dialect），随后咱们将 MGE Dialect 转换到 MLIR 既有的 Dialect 上，接下来的绝大部分工作都能够复用 MLIR 中的根底组件和工具实现。图 8 形容了 CPU 和 GPU 上大略的执行流程。

天元首先将 JITExecutor 算子外部的 internal graph 翻译成 MGE Dialect。在 CPU 上，MGE Dialect 会先 Lowering 到 Affine Dialect 上，而后会通过 LLVM 的组件 Lowering 到 LLVM Dialect 上，LLVM Dialect 能够被间接翻译成 LLVM IR。在这一步之后，其余优化工作都能够间接复用 LLVM 的根底组件。最初天元应用 MLIR ExecutionEngine 执行 LLVM IR 生成的代码。在 GPU 上，天元会先将 MGE Dialect Lowering 到 GPU Dialect 上，随后 Lowering 到 NVVM Dialect，NVVM 会被翻译成 PTX 汇编。最初通过英伟达提供的 CUmodule 和 CUfunction 两个机制运行。

四、试验和剖析

首先参考这篇文档 在天元中开启 JIT 反对。本次试验选了 resnet50, mobilenetV2 和 vgg16 三个业界宽泛应用的模型，batch size 别离设置了 1, 8 和 16。测试硬件环境为 NVIDIA T4，软件环境为 MegEngine v1.2.0。

由图 9 可知，和不关上 JIT 反对相比，关上 JIT 反对后 resnet50 最高能够取得 16% 的减速比，mobilenet V2 则能取得 6% 到 7% 的减速比，而 vgg16 其实上没有显著减速成果。这是因为 vgg16 模型很大，能够被优化的 element-wise 操作比拟少。JIT 的优化成果跟具体的模型是有严密关系的。

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

五、总结和瞻望

本篇文章介绍了天元应用 JIT 实现将任意多个相邻的 element-wise 算子交融成一个算子的优化。咱们在 T4 上用 MegEngine v1.2.0 试验，相比于优化前，resnet 50 最高能够取得 16% 的减速比。

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

1. 将 JIT 编译的后果先离线保留，线上间接将线下编译好的可执行对象读进内存。这种做法能够解决线上第一次运行慢的问题，但它可能会损失一部分可移植性，因为在一种设施上编译产生的可执行对象个别不能适配所有线上设施。2.JIT 反对更多的算子。3.JIT 反对更多的数据类型，天元 JIT 优化临时只反对 float32/float16 这两种数据类型。4. 动态图 JIT，也就是传统意义上的检测热点代码，重编译后再执行。

欢送小伙伴退出咱们 MegEngine 旷视天元开发者交换 QQ 群：1029741705
框架应用相干交换或反馈，欢送拜访论坛：http://discuss.megengine.org.cn；
GitHub 我的项目地址：http://github.com/MegEngine/M…