MegEngine 使用小技巧：如何解读 MegCC 编译模型几个阶段 Pass 的作用

MegCC 是一个真真实实的深度学习模型编译器，具备极其轻量的 Runtime 二进制体积，高性能，方便移植，极低内存使用以及快启动等核心特点。用户可在 MLIR 上进行计算图优化，内存规划，最后通过预先写好的 code 模版进行代码生成。

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 as mge
import megengine.functional as F
import megengine.module as M
import megengine.jit as jit
 
import numpy as np
 
# Define model
class ConvNet(M.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = M.Conv2d(1, 4, 3, padding=1)
        self.pool = M.MaxPool2d(2, 2)
        self.classifier = M.Linear(100, 5)
        self.relu = M.ReLU()
 
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = F.flatten(x, 1)
        x = self.classifier(x)
        return x
 
 
model = ConvNet()
 
@jit.trace(symbolic=True, capture_as_const=True)
def fun(data, *, net):
    pred = net(data)
    return pred
 
data = mge.Tensor(np.random.random([1, 1, 10, 10]).astype(np.float32))
 
fun(data, net=model)
fun.dump("test_model.mge", arg_names=["data"], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)

2. 导入模型

这一步主要将上面 dump 好的 MegEngine 模型 import 到 MegCC 的 MGB IR中，使用的工具是 MegCC 的 release 包中 bin/mgb-importer，执行命令：

./bin/mgb-importer test_model.mge test_model_mgb_ir.mlir

执行完成之后打开 test_model_mgb_ir.mlir，结果如下：

module {
  "MGB.ParamStorage"() {sym_name = "const{5}[0]", sym_visibility = "private", type = tensor<5xf32>, user_count = 1 : i32, value = dense<0.000000e+00> : tensor<5xf32>} : () -> ()
  "MGB.ParamStorage"() {sym_name = "const{1,4,1,1}[2]", sym_visibility = "private", type = tensor<1x4x1x1xf32>, user_count = 1 : i32, value = dense<0.000000e+00> : tensor<1x4x1x1xf32>} : () -> ()
  "MGB.ParamStorage"() {sym_name = "const{4,1,3,3}[6]", sym_visibility = "private", type = tensor<4x1x3x3xf32>, user_count = 1 : i32, value = dense<[[[[0.163880527, 0.566941559, 0.108093813], [-0.159407943, -0.3#
  "MGB.ParamStorage"() {sym_name = "const{5,100}[30]", sym_visibility = "private", type = tensor<5x100xf32>, user_count = 1 : i32, value = dense<"0x30394EBDE0DF49BEE368773D456F2B3E67A0FCBD9FC3683B3BF4B3BDCAD5B13#
  func @test_model_mgb_ir(%arg0: tensor<1x1x10x10xf32> {mgb.func_arg_name = "data"}) -> (tensor<1x5xf32> {mgb.func_result_name = "classifier.ADD"}) {
    %0 = "MGB.ParamProvider"() {name = @"const{5,100}[30]"} : () -> tensor<5x100xf32>
    %1 = "MGB.ParamProvider"() {name = @"const{4,1,3,3}[6]"} : () -> tensor<4x1x3x3xf32>
    %2 = "MGB.ParamProvider"() {name = @"const{1,4,1,1}[2]"} : () -> tensor<1x4x1x1xf32>
    %3 = "MGB.ParamProvider"() {name = @"const{5}[0]"} : () -> tensor<5xf32>
    %4 = "MGB.Reshape"(%arg0) {axis = 7 : i32} : (tensor<1x1x10x10xf32>) -> tensor<1x1x10x10xf32>
    %5 = "MGB.ConvBias"(%4, %1, %2) {compute_mode = 0 : i32, dilate_h = 1 : ui32, dilate_w = 1 : ui32, dtype = 0 : i32, format = 0 : i32, mode = 0 : i32, nonlineMode = 1 : i32, pad_h = 1 : ui32, pad_w = 1 : ui32#
    %6 = "MGB.Pooling"(%5) {format = 0 : i32, mode = 0 : i32, pad_h = 0 : ui32, pad_w = 0 : ui32, stride_h = 2 : ui32, stride_w = 2 : ui32, window_h = 2 : ui32, window_w = 2 : ui32} : (tensor<1x4x10x10xf32>) -> #
    %7 = "MGB.Reshape"(%6) {axis = 7 : i32} : (tensor<1x4x5x5xf32>) -> tensor<1x100xf32>
    %8 = "MGB.MatrixMul"(%7, %0) {compute_mode = 0 : i32, format = 0 : i32, strategy = 1 : i32, transposeA = false, transposeB = true, workspace_limit = 18446744073709551615 : ui64} : (tensor<1x100xf32>, tensor<#
    %9 = "MGB.Elemwise"(%3, %8) {mode = 16 : i32} : (tensor<5xf32>, tensor<1x5xf32>) -> tensor<1x5xf32>
    return %9 : tensor<1x5xf32>
  }
}

这里使用的 LLVM 的 IR 结构，参考 LLVM 的 IR 模块组。从上面的 IR 可以清楚的看到整个模型变成了一个 mlir 的模块，其中模型的入口变成了一个 func，还有如下变化：

参数全部转换为 MGB.ParamStorage，并使用 MGB.ParamProvider 在 func 中作为接口访问，MGB.ParamStorage 并 MGB.ParamProvider 通过 sym_name 连接在一起，如上面 const{5}[0] 这个字符就是一个符号。
这个 test_model.mge 变成了名字为 test_model_mgb_ir 的 func 类型，这个 func 的参数就是整个 test_model.mge 的输入Tensor，这里是：%arg0: tensor<1x1x10x10xf32> {mgb.func_arg_name = "data"}。
test_model.mge 中的所有算子一一对应的转换为 MGB IR，如： MGB.ConvBias，MGB.MatrixMul 等。
在mlir中每个 op 都有一个输入和对一个输入，这些输入输出可以通过链接关系构成一张计算图。

3. 将 Abstract Kernel IR 加载上代码，并降低到 Kernel IR

./bin/megcc-opt --MGB-to-Kernel --memory-forwarding --static-memory-planning --kernel-materialization test_model_mgb_ir.mlir

执行之后在终端中将输出：

#map0 = affine_map<(d0, d1) -> (d0 * 5 + d1 + 20)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0 * 100 + d1 * 100 + d2 * 10 + d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0 * 400 + d1 * 100 + d2 * 10 + d3)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0 * 100 + d1 * 25 + d2 * 5 + d3 + 1600)>
#map4 = affine_map<(d0, d1) -> (d0 * 100 + d1 + 1600)>
#map5 = affine_map<(d0, d1) -> (d0 * 5 + d1)>
module {
  "Kernel.KernelDef"() {body = "\0A#include <stdbool.h>....", sym_name = "kernel_conv2d_3x3_NCHW_DENSE_p1x1_s1x1_d1x1_f32f32f32f32_bias_RELU"} : () -> ()
  "Kernel.KernelDef"() {body = "\0A#include <stdbool.h>\0A\0A...", sym_name = "kernel_pooling_MAX_NCHW_p0x0_s2x2_w2x2_f32f32"} : () -> ()
  "Kernel.KernelDef"() {body = "#include <string.h>\0...", sym_name = "naive_kernel_gevmnt"} : () -> ()
  "Kernel.KernelDef"() {body = "\0A                #include \22gi_float.h\22\0A ...)", sym_name = "GI_kernel_elementwise_ADD_binary_VEC_VEC_f32f32f32"} : () -> ()
  "Kernel.WeightStorage"() {sym_name = "const{5}[0]", type = tensor<5xf32>, user_count = 1 : i32, value = dense<0.000000e+00> : tensor<5xf32>} : () -> ()
  "Kernel.WeightStorage"() {sym_name = "const{1,4,1,1}[2]", type = tensor<1x4x1x1xf32>, user_count = 1 : i32, value = dense<0.000000e+00> : tensor<1x4x1x1xf32>} : () -> ()
  "Kernel.WeightStorage"() {sym_name = "const{4,1,3,3}[6]", type = tensor<4x1x3x3xf32>, user_count = 1 : i32, value = dense<[[[[0.163880527, 0.566941559, 0.108093813], ...]]]> : tensor<4x1x3x3xf32>} : () -> ()
  "Kernel.WeightStorage"() {sym_name = "const{5,100}[30]", type = tensor<5x100xf32>, user_count = 1 : i32, value = dense<"0x30394EBDE0DF49BEE3687..."> : tensor<5x100xf32>} : () -> ()
  func @test_model_mgb_ir(%arg0: memref<1x1x10x10xf32> {mgb.func_arg_name = "data"}, %arg1: memref<2000xi8> {mgb.func_arg_name = "kGlobalBuffer"}) -> (memref<1x5xf32, #map0> {mgb.func_result_name = "classifier.ADD"}) {
    %0 = "Kernel.GetWeight"() {name = @"const{5,100}[30]"} : () -> memref<5x100xf32>
    %1 = "Kernel.GetWeight"() {name = @"const{4,1,3,3}[6]"} : () -> memref<4x1x3x3xf32>
    %2 = "Kernel.GetWeight"() {name = @"const{1,4,1,1}[2]"} : () -> memref<1x4x1x1xf32>
    %3 = "Kernel.GetWeight"() {name = @"const{5}[0]"} : () -> memref<5xf32>
    %4 = "Kernel.Reshape"(%arg0) {axis = 7 : i32, determined = true} : (memref<1x1x10x10xf32>) -> memref<1x1x10x10xf32, #map1>
    %5 = "Kernel.MemPlan"(%arg1) : (memref<2000xi8>) -> memref<1x4x10x10xf32, #map2>
    "Kernel.KernelCall"(%4, %1, %2, %5) {attrMap = {compute_mode = "DEFAULT", dilate_h = 1 : ui32, dilate_w = 1 : ui32, format = "NCHW", kernel_h = 3 : i32, kernel_w = 3 : i32, mode = "CROSS_CORRELATION", nonlineMode = "RELU", operand_segment_sizes = dense<[1, 1, 1, 0, 1]> : vector<5xi32>, pad_h = 1 : ui32, pad_w = 1 : ui32, sparse = "DENSE", strategy = 1 : i32, stride_h = 1 : ui32, stride_w = 1 : ui32, workspace_limit = 18446744073709551615 : ui64}, callee = @kernel_conv2d_3x3_NCHW_DENSE_p1x1_s1x1_d1x1_f32f32f32f32_bias_RELU, dynamic_shape = false, operand_segment_sizes = dense<[3, 1, 0]> : vector<3xi32>} : (memref<1x1x10x10xf32, #map1>, memref<4x1x3x3xf32>, memref<1x4x1x1xf32>, memref<1x4x10x10xf32, #map2>) -> ()
    %6 = "Kernel.MemPlan"(%arg1) : (memref<2000xi8>) -> memref<1x4x5x5xf32, #map3>
    "Kernel.KernelCall"(%5, %6) {attrMap = {format = "NCHW", mode = "MAX", pad_h = 0 : ui32, pad_w = 0 : ui32, stride_h = 2 : ui32, stride_w = 2 : ui32, window_h = 2 : ui32, window_w = 2 : ui32}, callee = @kernel_pooling_MAX_NCHW_p0x0_s2x2_w2x2_f32f32, dynamic_shape = false, operand_segment_sizes = dense<[1, 1, 0]> : vector<3xi32>} : (memref<1x4x10x10xf32, #map2>, memref<1x4x5x5xf32, #map3>) -> ()
    %7 = "Kernel.MemPlan"(%arg1) : (memref<2000xi8>) -> memref<1x100xf32, #map4>
    %8 = "Kernel.MemPlan"(%arg1) : (memref<2000xi8>) -> memref<1x5xf32, #map5>
    "Kernel.KernelCall"(%7, %0, %8) {attrMap = {compute_mode = "DEFAULT", format = "DEFAULT", transposeA = false, transposeB = true}, callee = @naive_kernel_gevmnt, dynamic_shape = false, operand_segment_sizes = dense<[2, 1, 0]> : vector<3xi32>} : (memref<1x100xf32, #map4>, memref<5x100xf32>, memref<1x5xf32, #map5>) -> ()
    %9 = "Kernel.MemPlan"(%arg1) : (memref<2000xi8>) -> memref<1x5xf32, #map0>
    "Kernel.KernelCall"(%3, %8, %9) {attrMap = {}, callee = @GI_kernel_elementwise_ADD_binary_VEC_VEC_f32f32f32, dynamic_shape = false, operand_segment_sizes = dense<[2, 1, 0]> : vector<3xi32>} : (memref<5xf32>, memref<1x5xf32, #map5>, memref<1x5xf32, #map0>) -> ()
    return %9 : memref<1x5xf32, #map0>
  }
}

上面就是最后编译完成之后的模型：

所有的内核都以 Kernel.KernelDef 字串形式进行定义，在后面将以 Kernel.KernelCall 字串形式进行调用，所有的 Kernel.KernelDef 都是以字串形式存在的纯 C 代码

Kernel.KernelDef 和 Kernel.KernelCall 之间使用符号进行对应，如上面的 kernel_conv2d_3x3_NCHW_DENSE_p1x1_s1x1_d1x1_f32f32f32f32_bias_RELU 字符。
所有的内存资源都是以 Kernel.MemPlan 的形式进行申请，
所有运算符的参数都在 Kernel.KernelCall 以字符串或者其字符的形式传递给具体的内核
每一个memref都确定了一个地图来指定其在内存计划中的访问列表。
将上面的Kernel IR按照Runtime确定的模型格式进行序列化以及将对应的代码串写到xxx.c文件中，就完成了整个模型的编译过程。

MegCC 中大多数 Kernel 为人工优化并提前写好的 Kernel 模板，这些模板会根据具体的 Operator 参数生成对应的 Kernel。大多数为人工优化的 Kernel 的原因是：目前在 CPU 上不搜参的情况下，mlir 生成的 Kernel 性能和手写的 Kernel 还有一定的距离，但是自动生成 Kernel 的方法长期来看是比较可取的。
MegCC 现已开源，仓库地址：github.com/MegEngine/MegCC，欢迎试用、star、issue。

附：

更多 MegEngine 信息获取，您可以：查看文档和 GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。