目前社区已经有多个移动端深度学习推理框架，如：NCNN、MNN... 这些推理引擎都给社区的用户带来了在移动端上部署深度学习非常多的便利，但是他们也都有一个共性问题：随着不断地迭代以及性能优化，运行时库会逐渐的增大，特别是在不同算子 fuse 的时候，会导致非常多的长尾算子，这就会增大我们 App 或者 SDK 的体积， 特别是在一些安全执行环境，类似 TEE 下，SDK 体积大小要求更加苛刻。。
为了解决这个问题，由MegEngine 团队开源的 MegCC 创新使用模型预编译的方案，生成模型推理必要的代码，去除掉了和模型推理无关的代码，因此极大程度上减少了推理引擎的体积。主要方法是：
将传统框架运行时的计算图优化、Kernel 选择、内存分配都移到编译时，从而最大程度上减少了 Runtime 时的二进制体积大小，并根据模型信息做进一步的性能优化。
该方案有以下优点：
·随着框架的迭代将不会使得推理引擎的体积增大
·很多的算子融合可以在编译时根据模型信息生成对应的 code
·模型编译时可以获得整个计算图的信息，这样可以进一步进行极致的性能优化
·可以吸收社区在代码生成方面的经验用于为 MegCC 生成 code
·Kernel 性能优化，因为每一个 Kernel 都是针对每一个 Operator 定制的，因此可以根据 Operator 的参数进行更加深入的优化。
·解决 Operator fuse 之后的算子长尾问题，比如对 conv 之后融合的 activation 的种类和数量没有限制，可以支持更多的 fuse，也不造成 Runtime 的大小有明显的改变。
·另外 MegCC 的 runtime 使用纯 C 实现，可以轻松移植到其他的嵌入式芯片中。
不同于传统推理框架，MegCC 是一个真真实实的深度学习模型编译器，具备极其轻量的 Runtime 二进制体积，高性能，方便移植，极低内存使用以及快启动等核心特点。用户可在 MLIR 上进行计算图优化，内存规划，最后通过预先写好的 code 模版进行代码生成。目前，MegCC 已支持 Arm64，Armv7，x86，risc-v 以及单片机平台。
近日，MegCC 持续升级，最新版除了针对Arm64优化，更有新工具，新体验，并针对用户使用体验以及模型推理性能进行全面提升，主要的提升包括：
1.新增 Benchmark 工具， 用于用于快速 Benchmark 常用模型的推理性能并可视化；
2.新增 Kernel C 代码导出工具，方便用户定制化获取算子 Kernel， 方便迁移与复用；
3.优化 NN Kernel 性能, 保持推理 SDK 性能先进；
4.支持第三方 NPU loader，方便 NPU 相关应用迁移。
使用方法及效果
使用 MegCC 完成模型部署只需要完成以下 3 步：
·模型编译：编译 MegEngine 模型，生成运行这个模型对应的 Kernel 以及优化之后的模型。
·Runtime编译：这个阶段会将 Runtime 和上一步中生成的 Kernel 一起编译成一个 SDK。
·集成到应用中：调用上一步编译的 SDK 的接口进行推理。
以 YOLOX 模型为例，运行效果如下图：
从图中可见，MegCC 生成的推理程序在保证推理性能良好（模型测速结果为 670ms）的情况下，其大小可以小到 95KB。
持续优化中的模型性能：
上图中部分模型 MegCC 略微慢于 MegEngine 的原因是: MegEngine 有完善的算法搜索逻辑，部分场景选出来的算法优于 MegCC， 后续版本的 MegCC 会补齐这部分工作。
详细操作文档：MegCC/how-to-use-chinese.md at main · MegEngine/MegCC
手把手操作教程：【挑战 100KB 可执行程序高性能推理 YOLOX 模型】 https://www.bilibili.com/video/BV1tg411B7dx/?
原理解析
作为一个基于 MLIR 实现的深度学习编译器，MegCC 的实现关键就是如何根据需求定义一系列 IR 以及 Pass ，将高级 IR 逐步 lowering 到低级 IR，同时进行算子融合等优化。
MegCC 实现的原理是：深度学习模型在推理时候，每一个 Operator 都会对应一个计算 kernel 并完成计算，所以整个深度学习模型在推理时就是一次执行所有 Operator 的计算 kernel，执行完成之后就可以获得最终推理的结果。传统深度学习推理框架在运行时会做以下几件事情：
·计算图优化 ----- 主要和模型相关。
·Kernel 选择 ----- 为模型的每个 Operator 根据参数选择合适的 Kernel 进行计算。
·内存分配 ----- 由模型以及模型中每个 Operator 执行的 Kernel 决定内存分配的大小。
·执行每个 Operator 的 Kernel ----- 和推理的数据强相关。
在上述传统深度学习推理需要完成的事情中，图优化，Kernel 选择，内存分配都是只和训练好的模型相关和推理时候的输入数据不相关，因此这些工作都可以放在模型编译时完成，运行时仅仅执行每一个 Operator 的 Kernel 就可以完成推理。MegCC 就是将上面图优化，Kernel 选择，内存分配都放在 MegCC 的编译阶段完成，将 Operator 的 Kernel 计算才放到 Runtime 中进行计算。
MegCC 主要包含两部分，一部分是 compiler 部分，另外一部分是 runtime 部分，下面重点介绍与编译相关的 compiler 部分。
MegCC compiler
Compiler 主要流程是：
1.依赖 MegEngine 进行模型的导入和静态图优化（block-level optimizations，算子融合等）。
2.将优化后的模型转换为基于 mlir 自定义的 MGB IR。
3.MGB IR 经过一系列 pass 经过 Abstract Kernel IR 最终转换到 Kernel IR。
4.将 Kernel IR 导出为 runtime model 和 runtime kernel，供 MegCC 的 runtime 部分使用。
MegCC compiler 流程
MegCC 中的 IR
MegCC 基于 MLIR 定义了一系列的 IR。MLIR 的 IR 定义需要用户定义 Dialect（详见官方文档），然后由 TableGen 在程序编译阶段转换成 C++ 表示。
·MGB IR：定义和 MegEngine 中 Operator 一一对应，是 MegCC 导入进 mlir 系统的入口 IR，它包含了每个 Opr 的类型以及这个 Opr 对应的参数，其每一个输入输出变量都是 Tensor，并且是单赋值(SSA)的。详见 GitHub MegCC MGB IR。
·Abstract Kernel IR：抽象 Kernel 层 IR，主要上面 MGB IR 通过转换之后得到，该 IR 中的输入输出已经 lowering 到 Buffer 了，因此不再是 SSA，另外 Opr 的属性也由 MegEngine 中定义的枚举值，转变成为了字符串。详见 GitHub MegCC Abstract Kernel IR。
·Kernel IR：表示已经生成 Kernel 之后的IR形式，其已经没有 Opr 的概念，整个计算图通过一个个对应的 Kernel 链接在一起，Opr 的参数等都固化在了定义好的 Kernel 中。详见 GitHub MegCC Kernel IR。
MegCC 中主要的 Pass
- MGBToKernelPass：这个 Pass 主要将 MGB IR 转换为 Abstract Kernel IR，转换过程中主要完成几件事情：
· 将 MGB IR 中的所有输入输出 Tensor 类型转换为 Buffer 类型。
· 将 MGB IR 中的所有枚举参数转换为对应的字符，这样 Abstract Kernel IR 就可以完全和 MegEngine 解耦。
· 将一些内存搬运相关的 Opr 全部转换为 Relayout，如：Concat，SetSubtensor 等 Opr（node-level optimizations）。
· 将判断 Opr 是静态 shape 还是动态 shape，动态 shape 就是输入 tensor 的 shape 需要依赖输入的值才能计算出来的，如：输出一个 tensor 中所有大于 1 的数。如果是静态 shape 直接转换到 Abstract Kernel IR，如果是动态 shape 直接转换到 Kernel IR 的 Instruction 中。
- MGBFuseKernelPass：应用在 MGB IR 上，基于 mlir 的模板匹配的方法尽可能的完成 kernel 的融合，比如连续两个 typecvt 合并成为一个 typecvt 等（block-level optimizations，算子融合）。
- MemoryForwardingPass：将遍历 Abstract Kernel IR 所有可能不用计算，直接 share 输入内存的 Opr，如果这些 Opr 确实不用计算，则直接 forward memory，如果这些 Opr 需要进行内存搬运，则会用 Relayout Opr 替换原来的 Opr（node-level optimizations）。
- KernelMaterializationPass：将所有 Abstract Kernel IR 都装载上真正 Kernel code 并转化为 KernelCall，然后添加对应的 KernelDef。KernelCall 和 KernelDef 之间通过 symbol 进行匹配。
- StaticMemoryPlanningPass：将所有静态 shape 的 memref 进行内存规划，内存规划算法使用改进的 MegEngine 的内存规划算法--PushDown 算法，能够极大程度的压缩运行时内存使用量。同时将 mlir 的 memref.Alloc 替换为 Kernel IR 的 MemPlan，MemPlan 中主要记录了内存规划的一整块 memref 以及该 Tensor 在规划的内存中的偏移量（dataflow-level optimizations，静态内存规划）。
上面的 Pass 就完成模型的图优化、内存规划以及 Kernel 生成，上文提到的后端优化即在 Kernel 生成阶段体现，目前 MegCC 主要使用人工优化的 Kernel 模版。最终可以根据 Runtime 中定义的模型格式 dump 编译之后的模型，以及生成计算模型所需的 Kernel 文件。 下面以一个简单的模型为例，使用 MegCC 的辅助工具(下载 Release 包) mgb-importer 和 megcc-opt，观察经过各个 Pass 的处理 IR 的变化。也可使用 mgb-to-tinynn 工具直接完成模型的编译过程，详见 MegCC 入门文档。
1、dump 模型（使用 megengine）
import megengine.functional as F import megengine.module as M import megengine.optimizer as optim from megengine import jit import megengine import numpy as np class MulAddNet(M.Module): def __init__(self): super().__init__() def forward(self, input): x = input * 2. x = x + 1.5 return x model = MulAddNet() model.eval() @jit.trace(symbolic=True, capture_as_const=True) def infer_func(data, *, model): pred = model(data) return pred data = megengine.Tensor([[1., 2.], [3., 4.]]) output = infer_func(data, model=model) print(output)
2、importer 模型到 MGB IR
./bin/mgb-importer MulAdd.mge mulAdd.mlir
cat mulAdd.mlir
output:
module {
"MGB.ParamStorage"() {sym_name = "const<2>[2]", sym_visibility = "private", type = tensor<1xf32>, user_count = 1 : i32, value = dense<2.000000e+00> : tensor<1xf32>} : () -> ()
"MGB.ParamStorage"() {sym_name = "const<1.5>[4]", sym_visibility = "private", type = tensor<1xf32>, user_count = 1 : i32, value = dense<1.500000e+00> : tensor<1xf32>} : () -> ()
func @mulAdd(%arg0: tensor<2x2xf32> {mgb.func_arg_name = "data"}) -> (tensor<2x2xf32> {mgb.func_result_name = "FUSE_MUL_ADD3(const<2>[2],data,const<1.5>[4])[14]"}) {
%0 = "MGB.Reshape"(%arg0) {axis = 7 : i32} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%1 = "MGB.ParamProvider"() {name = @"const<1.5>[4]"} : () -> tensor<1xf32>
%2 = "MGB.ParamProvider"() {name = @"const<2>[2]"} : () -> tensor<1xf32>
%3 = "MGB.Elemwise"(%2, %0, %1) {mode = 35 : i32} : (tensor<1xf32>, tensor<2x2xf32>, tensor<1xf32>) -> tensor<2x2xf32>
return%3 : tensor<2x2xf32>
}
}
可以看到，在 importer 的过程中，乘法运算和加法运算被融合成了"FUSE_MUL_ADD3"。
3、MGBToKernelPass、MemoryForwardingPass 和 StaticMemoryPlanningPass
./bin/megcc-opt --MGB-to-Kernel --memory-forwarding --static-memory-planning mulAdd.mlir > mulAdd_final.mlir
cat mulAdd_final.mlir
output:
#map = affine_map<(d0, d1) -> (d0 * 2+ d1)>
module {
"Kernel.WeightStorage"() {sym_name = "const<2>[2]", type = tensor<1xf32>, user_count = 1: i32, value = dense<2.000000e+00> : tensor<1xf32>} : () -> ()
"Kernel.WeightStorage"() {sym_name = "const<1.5>[4]", type = tensor<1xf32>, user_count = 1: i32, value = dense<1.500000e+00> : tensor<1xf32>} : () -> ()
func @mulAdd(%arg0: memref<2x2xf32> {mgb.func_arg_name = "data"}, %arg1: memref<16xi8> {mgb.func_arg_name = "kGlobalBuffer"}) -> (memref<2x2xf32, #map> {mgb.func_result_name = "FUSE_MUL_ADD3(const<2>[2],data,const<1.5>[4])[14]"}) {
%0= "Kernel.Reshape"(%arg0) {axis = 7: i32, determined = true} : (memref<2x2xf32>) -> memref<2x2xf32, #map>
%1= "Kernel.GetWeight"() {name = @"const<1.5>[4]"} : () -> memref<1xf32>
%2= "Kernel.GetWeight"() {name = @"const<2>[2]"} : () -> memref<1xf32>
%3= "Kernel.MemPlan"(%arg1) : (memref<16xi8>) -> memref<2x2xf32, #map>
"Kernel.FUSE_MUL_ADD3"(%2, %0, %1, %3) : (memref<1xf32>, memref<2x2xf32, #map>, memref<1xf32>, memref<2x2xf32, #map>) -> ()
return%3: memref<2x2xf32, #map>
}
}
经过上面几个 Pass，MGB IR 被转换为了 Kernel IR 并进行了内存规划。感兴趣的话可以更细粒度地看每个 Pass 做的事情，使用 megcc-opt 的参数控制使用哪些 Pass。
Kernel 生成
MegCC Compiler 会为模型中的每个 Operator 生成一个对应的 Kernel 来完成计算。 目前 MegCC 中大多数 Kernel 为人工优化并提前写好的 Kernel 模板，这些模板会根据具体的 Operator 参数生成对应的 Kernel。大多数为人工优化的 Kernel 的原因是：目前在 CPU 上不搜参的情况下，mlir 生成的 Kernel 性能和手写的 Kernel 还有一定的距离，但是自动生成 Kernel 的方法长期来看是比较可取的。MegCC中也做了部分尝试，细节请参考 MegCC 自动代码生成。
MegCC 现已开源，未来也将不断进阶，仓库地址： https://github.com/MegEngine/MegCC，欢迎试用、star、issue。
附：
更多 MegEngine 信息获取，您可以：查看文档、和GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。
