MegCC 是一個(gè)真真實(shí)實(shí)的深度學(xué)習(xí)模型編譯器赡艰,具備極其輕量的 Runtime 二進(jìn)制體積嵌纲,高性能蒲跨,方便移植满着,極低內(nèi)存使用以及快啟動(dòng)等核心特點(diǎn)准谚。用戶可在 MLIR 上進(jìn)行計(jì)算圖優(yōu)化挫剑,內(nèi)存規(guī)劃,最后通過(guò)預(yù)先寫好的 code 模版進(jìn)行代碼生成柱衔。
MegCC 中主要的 Pass
- MGBToKernelPass:這個(gè) Pass 主要將 MGB IR 轉(zhuǎn)換為 Abstract Kernel IR樊破,轉(zhuǎn)換過(guò)程中主要完成幾件事情:
- 將 MGB IR 中的所有輸入輸出 Tensor 類型轉(zhuǎn)換為 Buffer 類型。
- 將 MGB IR 中的所有枚舉參數(shù)轉(zhuǎn)換為對(duì)應(yīng)的字符唆铐,這樣 Abstract Kernel IR 就可以完全和 MegEngine 解耦哲戚。
- 將一些內(nèi)存搬運(yùn)相關(guān)的 Opr 全部轉(zhuǎn)換為 Relayout,如:Concat艾岂,SetSubtensor 等 Opr(node-level optimizations)顺少。
- 將判斷 Opr 是靜態(tài) shape 還是動(dòng)態(tài) shape,動(dòng)態(tài) shape 就是輸入 tensor 的 shape 需要依賴輸入的值才能計(jì)算出來(lái)的王浴,如:輸出一個(gè) tensor 中所有大于 1 的數(shù)脆炎。如果是靜態(tài) shape 直接轉(zhuǎn)換到 Abstract Kernel IR,如果是動(dòng)態(tài) shape 直接轉(zhuǎn)換到 Kernel IR 的 Instruction 中氓辣。
- MGBFuseKernelPass:應(yīng)用在 MGB IR 上秒裕,基于 mlir 的模板匹配的方法盡可能的完成 kernel 的融合,比如連續(xù)兩個(gè) typecvt 合并成為一個(gè) typecvt 等(block-level optimizations钞啸,算子融合)几蜻。
- MemoryForwardingPass:將遍歷 Abstract Kernel IR 所有可能不用計(jì)算喇潘,直接 share 輸入內(nèi)存的 Opr,如果這些 Opr 確實(shí)不用計(jì)算梭稚,則直接 forward memory颖低,如果這些 Opr 需要進(jìn)行內(nèi)存搬運(yùn),則會(huì)用 Relayout Opr 替換原來(lái)的 Opr(node-level optimizations)哨毁。KernelMaterializationPass:將所有 Abstract Kernel IR 都裝載上真正 Kernel code 并轉(zhuǎn)化為 KernelCall枫甲,然后添加對(duì)應(yīng)的 KernelDef源武。KernelCall 和 KernelDef 之間通過(guò) symbol 進(jìn)行匹配扼褪。
- StaticMemoryPlanningPass:將所有靜態(tài) shape 的 memref 進(jìn)行內(nèi)存規(guī)劃,內(nèi)存規(guī)劃算法使用改進(jìn)的 MegEngine 的內(nèi)存規(guī)劃算法--PushDown 算法粱栖,能夠極大程度的壓縮運(yùn)行時(shí)內(nèi)存使用量话浇。同時(shí)將 mlir 的 memref.Alloc 替換為 Kernel IR 的 MemPlan,MemPlan 中主要記錄了內(nèi)存規(guī)劃的一整塊 memref 以及該 Tensor 在規(guī)劃的內(nèi)存中的偏移量(dataflow-level optimizations闹究,靜態(tài)內(nèi)存規(guī)劃)幔崖。
上面的 Pass 就完成模型的圖優(yōu)化、內(nèi)存規(guī)劃以及 Kernel 生成渣淤,上文提到的后端優(yōu)化即在 Kernel 生成階段體現(xiàn)赏寇,目前 MegCC 主要使用人工優(yōu)化的 Kernel 模版。最終可以根據(jù) Runtime 中定義的模型格式 dump 編譯之后的模型价认,以及生成計(jì)算模型所需的 Kernel 文件嗅定。 下面以一個(gè)簡(jiǎn)單的模型為例,使用 MegCC 的輔助工具(下載 Release 包) mgb-importer 和 megcc-opt用踩,觀察經(jīng)過(guò)各個(gè) Pass 的處理 IR 的變化渠退。也可使用 mgb-to-tinynn 工具直接完成模型的編譯過(guò)程,詳見(jiàn) MegCC 入門文檔脐彩。
- 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)
- 導(dǎo)入模型
這一步主要將上面 dump 好的 MegEngine 模型 import 到 MegCC 的 MGB IR中碎乃,使用的工具是 MegCC 的 release 包中 bin/mgb-importer,執(zhí)行命令:
./bin/mgb-importer test_model.mge test_model_mgb_ir.mlir
執(zhí)行完成之后打開 test_model_mgb_ir.mlir惠奸,結(jié)果如下:
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 結(jié)構(gòu)梅誓,參考 LLVM 的 IR 模塊組。從上面的 IR 可以清楚的看到整個(gè)模型變成了一個(gè) mlir 的模塊佛南,其中模型的入口變成了一個(gè) func梗掰,還有如下變化:
參數(shù)全部轉(zhuǎn)換為 MGB.ParamStorage,并使用 MGB.ParamProvider 在 func 中作為接口訪問(wèn)共虑,MGB.ParamStorage 并 MGB.ParamProvider 通過(guò) sym_name 連接在一起愧怜,如上面 const{5}[0] 這個(gè)字符就是一個(gè)符號(hào)。
這個(gè) test_model.mge 變成了名字為 test_model_mgb_ir 的 func 類型妈拌,這個(gè) func 的參數(shù)就是整個(gè) test_model.mge 的輸入Tensor拥坛,這里是:%arg0: tensor<1x1x10x10xf32> {mgb.func_arg_name = "data"}蓬蝶。
test_model.mge 中的所有算子一一對(duì)應(yīng)的轉(zhuǎn)換為 MGB IR,如: MGB.ConvBias猜惋,MGB.MatrixMul 等丸氛。
在mlir中每個(gè) op 都有一個(gè)輸入和對(duì)一個(gè)輸入,這些輸入輸出可以通過(guò)鏈接關(guān)系構(gòu)成一張計(jì)算圖著摔。
- 將 Abstract Kernel IR 加載上代碼缓窜,并降低到 Kernel IR
./bin/megcc-opt --MGB-to-Kernel --memory-forwarding --static-memory-planning --kernel-materialization test_model_mgb_ir.mlir
執(zhí)行之后在終端中將輸出:
#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>
}
}
上面就是最后編譯完成之后的模型:
所有的內(nèi)核都以 Kernel.KernelDef 字串形式進(jìn)行定義,在后面將以 Kernel.KernelCall 字串形式進(jìn)行調(diào)用谍咆,所有的 Kernel.KernelDef 都是以字串形式存在的純 C 代碼
Kernel.KernelDef 和 Kernel.KernelCall 之間使用符號(hào)進(jìn)行對(duì)應(yīng)禾锤,如上面的 kernel_conv2d_3x3_NCHW_DENSE_p1x1_s1x1_d1x1_f32f32f32f32_bias_RELU 字符。
所有的內(nèi)存資源都是以 Kernel.MemPlan 的形式進(jìn)行申請(qǐng)摹察,
所有運(yùn)算符的參數(shù)都在 Kernel.KernelCall 以字符串或者其字符的形式傳遞給具體的內(nèi)核
每一個(gè)memref都確定了一個(gè)地圖來(lái)指定其在內(nèi)存計(jì)劃中的訪問(wèn)列表恩掷。
將上面的Kernel IR按照Runtime確定的模型格式進(jìn)行序列化以及將對(duì)應(yīng)的代碼串寫到xxx.c文件中,就完成了整個(gè)模型的編譯過(guò)程供嚎。
MegCC 中大多數(shù) Kernel 為人工優(yōu)化并提前寫好的 Kernel 模板黄娘,這些模板會(huì)根據(jù)具體的 Operator 參數(shù)生成對(duì)應(yīng)的 Kernel。大多數(shù)為人工優(yōu)化的 Kernel 的原因是:目前在 CPU 上不搜參的情況下克滴,mlir 生成的 Kernel 性能和手寫的 Kernel 還有一定的距離逼争,但是自動(dòng)生成 Kernel 的方法長(zhǎng)期來(lái)看是比較可取的。
MegCC 現(xiàn)已開源劝赔,倉(cāng)庫(kù)地址:github.com/MegEngine/MegCC誓焦,歡迎試用、star望忆、issue罩阵。