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 的模板匹配的方法（https://mlir.llvm.org/docs/PatternRewriter/）盡可能的完成 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（https://mlir.llvm.org/docs/PatternRewriter/）?進(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?包https://github.com/MegEngine/MegCC/releases)?mgb-importer?和 megcc-opt，觀察經(jīng)過(guò)各個(gè) Pass 的處理 IR 的變化。也可使用?mgb-to-tinynn?工具直接完成模型的編譯過(guò)程，詳見(jiàn)?MegCC 入門文檔：https://github.com/MegEngine/MegCC/blob/main/doc/how-to-use-chinese.md

1、dump 模型（使用 megengine）

2、導(dǎo)入模型

這一步主要將上面 dump 好的 MegEngine 模型 import 到 MegCC 的 MGB IR中，使用的工具是 MegCC 的 release 包中 bin/mgb-importer，執(zhí)行命令：

執(zhí)行完成之后打開(kāi) test_model_mgb_ir.mlir，結(jié)果如下：

這里使用的 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ì)算圖。

3、將 Abstract Kernel IR 加載上代碼，并降低到 Kernel IR

執(zhí)行之后在終端中將輸出：

上面就是最后編譯完成之后的模型：

所有的內(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)已開(kāi)源，倉(cāng)庫(kù)地址：github.com/MegEngine/MegCC，歡迎試用、star、issue。

附：

更多 MegEngine 信息獲取，您可以：查看文檔：https://www.megengine.org.cn/doc/stable/zh/

GitHub 項(xiàng)目： https://github.com/MegEngine

加入 MegEngine 用戶交流 QQ 群：1029741705

歡迎參與 MegEngine 社區(qū)貢獻(xiàn)，成為?Awesome MegEngineer：https://www.megengine.org.cn/community-AMGE，榮譽(yù)證書(shū)、定制禮品享不停。