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

2023-05-30
北京
本文字数：6138 字
阅读完需：约 20 分钟

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 入门文档。

dump 模型（使用 megengine）

import megengine as mgeimport megengine.functional as Fimport megengine.module as Mimport megengine.jit as jit import numpy as np # Define modelclass 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)

复制代码

导入模型

这一步主要将上面 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 都有一个输入和对一个输入，这些输入输出可以通过链接关系构成一张计算图。

将 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，荣誉证书、定制礼品享不停。

发布于: 刚刚阅读数: 3

MegEngineBot

关注

工业级研发的开源深度学习框架-MegEngine 2022-10-21 加入

官网：https://www.megengine.org.cn/ MegEngine 技术交流 QQ 群：1029741705 框架使用相关交流或反馈，欢迎访问论坛：https://discuss.megengine.org.cn/ GitHub ：http://github.com/MegEngine/MegEngine

发布

暂无评论

创作场景

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

附：

MegEngineBot

评论