乐趣区

关于编译器:MegEngine-使用小技巧如何解读-MegCC-编译模型几个阶段-Pass-的作用

jiezi

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 都有一个输出和对一个输出，这些输入输出能够通过链接关系形成一张计算图。

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

退出移动版