如何寫一個深度學(xué)習(xí)編譯器
作者：曠視 MegEngine 框架開發(fā)師 - 李明鑫
編譯器本質(zhì)上是一種提高開發(fā)效率的工具，將高級語言轉(zhuǎn)換為低級語言（通常是二進制機器碼），使得程序員不需要徒手寫二進制。轉(zhuǎn)換過程中，首要任務(wù)是保證正確性，同時需要進行優(yōu)化以提升程序的運行效率。傳統(tǒng)意義上的編譯器的輸入通常是某種高級語言，輸出是可執(zhí)行程序。在實際工作中接觸到了深度學(xué)習(xí)編譯器開發(fā)，其設(shè)計思想與傳統(tǒng)編譯器非常類似，所以本文以深度學(xué)習(xí)編譯器的開發(fā)、結(jié)合我們實際開發(fā)的深度學(xué)習(xí)編譯器 MegCC 為例，來說明如何寫一個編譯器。 本文主要分為以下兩個部分：
介紹深度學(xué)習(xí)編譯器，重點介紹編譯器中前端和后端的優(yōu)化方法。
以 MegCC 為例介紹如何開發(fā)一個深度學(xué)習(xí)編譯器。
深度學(xué)習(xí)編譯器簡介
與傳統(tǒng)編譯器不同，深度學(xué)習(xí)編譯器的輸入是神經(jīng)網(wǎng)絡(luò)模型、輸出是可運行在不同平臺的表達(dá)了輸入的神經(jīng)網(wǎng)絡(luò)模型的計算過程的可執(zhí)行程序。但深度學(xué)習(xí)編譯器又與傳統(tǒng)編譯器類似，都分為前端和后端，前端負(fù)責(zé)執(zhí)行硬件無關(guān)的優(yōu)化，后端負(fù)責(zé)執(zhí)行硬件相關(guān)的優(yōu)化。對編譯器來說，最重要的兩個概念是 IR(intermediate representation, 中間表示)和 Pass。對于人類來說，抽象是理解復(fù)雜事物的一種重要方式，IR 就是對編譯過程中間產(chǎn)物的抽象，IR 通常有多級，越高級的 IR 越抽象，越低級的 IR 越具體。Pass 定義了如何將高級 IR 逐步 lowering 到低級 IR，并負(fù)責(zé)進行優(yōu)化。下面根據(jù)前端和后端進行分類，介紹優(yōu)化的方法。
前端優(yōu)化方法
前端首先需要根據(jù)輸入的模型構(gòu)建計算圖，生成 high-level IR，然后進行一系列的優(yōu)化。由于優(yōu)化是基于計算圖的，并不涉及具體計算，所以該優(yōu)化是后端無關(guān)的。常見的優(yōu)化手段有可分為三類：node-level optimizations；block-level optimizations； dataflow-level optimizations。
node-level optimizations。節(jié)點層面的優(yōu)化主要是消除一些不必要的節(jié)點以及將某些節(jié)點替換為代價更小的節(jié)點。比如使用矩陣 A 與一個 0 維矩陣相加，則可消除該加法操作。
block-level optimizations。塊層面的優(yōu)化主要有代數(shù)簡化和算子融合。
a. 代數(shù)簡化，例如 A^T 和 B^T 進行矩陣乘，則可使用 B 與 A 矩陣乘之后進行轉(zhuǎn)置進行替換，可節(jié)約一次轉(zhuǎn)置運算。
b. 算子融合是常見的深度學(xué)習(xí)的優(yōu)化手段。算子融合雖然不能減少計算量，但是可以減少訪存量，提高計算訪存比，從而提升性能。dataflow-level optimizations。數(shù)據(jù)流層面的優(yōu)化主要有靜態(tài)內(nèi)存規(guī)劃等。
a. 靜態(tài)內(nèi)存規(guī)劃通過在不發(fā)生內(nèi)存重疊的前提下盡可能復(fù)用內(nèi)存，使得程序運行時所使用的內(nèi)存盡可能小。
后端優(yōu)化方法
后端通用的優(yōu)化有循環(huán)展開、循環(huán)融合、掩蓋訪存等；另外根據(jù)硬件的不同，可使用基于硬件的指令映射、向量化等并行計算以及手工編寫匯編 kernel 等手段進行針對性優(yōu)化。圖 1 展示了常用的后端優(yōu)化方法[1]。
MegCC
接下來就以 MegCC 為例概括介紹一下基于 MLIR 實現(xiàn)一個深度學(xué)習(xí)編譯器，其關(guān)鍵就是如何根據(jù)需求定義一系列 IR，以及定義 Pass 將高級 IR lowering 到低級 IR，同時進行上述優(yōu)化。
MegCC簡介
MegCC 實現(xiàn)的原理是：深度學(xué)習(xí)模型在推理時候，每一個 Operator 都會對應(yīng)一個計算 kernel 并完成計算，所以整個深度學(xué)習(xí)模型在推理時就是一次執(zhí)行所有 Operator 的計算 kernel，執(zhí)行完成之后就可以獲得最終推理的結(jié)果。傳統(tǒng)深度學(xué)習(xí)推理框架在運行時會做以下幾件事情：
計算圖優(yōu)化 ----- 主要和模型相關(guān)。
Kernel 選擇 ----- 為模型的每個 Operator 根據(jù)參數(shù)選擇合適的 Kernel 進行計算。
內(nèi)存分配 ----- 由模型以及模型中每個 Operator 執(zhí)行的 Kernel 決定內(nèi)存分配的大小。
執(zhí)行每個 Operator 的 Kernel ----- 和推理的數(shù)據(jù)強相關(guān)。
在上述傳統(tǒng)深度學(xué)習(xí)推理需要完成的事情中，圖優(yōu)化，Kernel 選擇，內(nèi)存分配都是只和訓(xùn)練好的模型相關(guān)和推理時候的輸入數(shù)據(jù)不相關(guān)，因此這些工作都可以放在模型編譯時完成，運行時僅僅執(zhí)行每一個 Operator 的 Kernel 就可以完成推理。MegCC 就是將上面圖優(yōu)化，Kernel 選擇，內(nèi)存分配都放在 MegCC 的編譯階段完成，將 Operator 的 Kernel 計算才放到 Runtime 中進行計算，這樣有以下優(yōu)勢：
Runtime 非常輕量，比起傳統(tǒng)的推理框架小一個數(shù)量級，因為 Runtime 只包含了模型中所必須的 Kernel，不相關(guān)的不會被編譯進去。
提升性能，因為 Runtime 只做 kernel 計算，所以避免了不必要的開銷。
Kernel 性能優(yōu)化，因為每一個 Kernel 都是針對每一個 Operator 定制的，因此可以根據(jù) Operator 的參數(shù)進行更加深入的優(yōu)化。
解決 Operator fuse 之后的算子長尾問題，比如對 conv 之后融合的 activation 的種類和數(shù)量沒有限制，可以支持更多的 fuse，也不造成 Runtime 的大小有明顯的改變。
另外 MegCC 的 runtime 使用純 C 實現(xiàn)，可以輕松移植到其他的嵌入式芯片中。
MegCC 主要包含兩部分，一部分是 compiler 部分，另外一部分是 runtime 部分，下面重點介紹與編譯相關(guān)的 compiler 部分。
MegCC compiler
Compiler 主要流程是：
依賴 MegEngine (我司開源深度學(xué)習(xí)框架)進行模型的導(dǎo)入和靜態(tài)圖優(yōu)化（block-level optimizations，算子融合等）。
將優(yōu)化后的模型轉(zhuǎn)換為基于 mlir 自定義的 MGB IR。
MGB IR 經(jīng)過一系列 pass 經(jīng)過 Abstract Kernel IR 最終轉(zhuǎn)換到 Kernel IR。
將 Kernel IR 導(dǎo)出為 runtime model 和 runtime kernel，供 MegCC 的 runtime 部分使用。
MegCC 中的 IR
MegCC 基于 MLIR 定義了一系列的 IR。MLIR 的 IR 定義需要用戶定義 Dialect（詳見官方文檔），然后由 TableGen 在程序編譯階段轉(zhuǎn)換成 C++ 表示。
MGB IR：定義為和 MegEngine 中 Operator 一一對應(yīng)，是 MegCC 導(dǎo)入進 mlir 系統(tǒng)的入口 IR，它包含了每個 Opr 的類型以及這個 Opr 對應(yīng)的參數(shù)，其每一個輸入輸出變量都是 Tensor，并且是單賦值(SSA)的。詳見?GitHub MegCC MGB IR。
Abstract Kernel IR：抽象 Kernel 層 IR，主要上面 MGB IR 通過轉(zhuǎn)換之后得到，該 IR 中的輸入輸出已經(jīng) lowering 到?Buffer?了，因此不再是 SSA，另外 Opr 的屬性也由 MegEngine 中定義的枚舉值，轉(zhuǎn)變成為了字符串。詳見?GitHub MegCC Abstract Kernel IR。
Kernel IR：表示已經(jīng)生成 Kernel 之后的IR形式，其已經(jīng)沒有 Opr 的概念，整個計算圖通過一個個對應(yīng)的 Kernel 鏈接在一起，Opr 的參數(shù)等都固化在了定義好的 Kernel 中。詳見?GitHub MegCC Kernel IR。
MegCC 中主要的 Pass
MGBToKernelPass：這個 Pass 主要將 MGB IR 轉(zhuǎn)換為 Abstract Kernel IR，轉(zhuǎn)換過程中主要完成幾件事情：
將 MGB IR 中的所有輸入輸出 Tensor 類型轉(zhuǎn)換為 Buffer 類型。
將 MGB IR 中的所有枚舉參數(shù)轉(zhuǎn)換為對應(yīng)的字符，這樣 Abstract Kernel IR 就可以完全和 MegEngine 解耦。
將一些內(nèi)存搬運相關(guān)的 Opr 全部轉(zhuǎn)換為 Relayout，如：Concat，SetSubtensor 等 Opr（node-level optimizations）。
將判斷 Opr 是靜態(tài) shape 還是動態(tài) shape，動態(tài) shape 就是輸入 tensor 的 shape 需要依賴輸入的值才能計算出來的，如：輸出一個 tensor 中所有大于 1 的數(shù)。如果是靜態(tài) shape 直接轉(zhuǎn)換到 Abstract Kernel IR，如果是動態(tài) shape 直接轉(zhuǎn)換到 Kernel IR 的 Instruction 中。
MGBFuseKernelPass：應(yīng)用在 MGB IR 上，基于?mlir 的模板匹配的方法盡可能的完成 kernel 的融合，比如連續(xù)兩個 typecvt 合并成為一個 typecvt 等（block-level optimizations，算子融合）。
MemoryForwardingPass：將遍歷 Abstract Kernel IR 所有可能不用計算，直接 share 輸入內(nèi)存的 Opr，如果這些 Opr 確實不用計算，則直接 forward memory，如果這些 Opr 需要進行內(nèi)存搬運，則會用 Relayout Opr 替換原來的 Opr（node-level optimizations）。
KernelMaterializationPass：將所有 Abstract Kernel IR 都裝載上真正 Kernel code 并轉(zhuǎn)化為 KernelCall，然后添加對應(yīng)的 KernelDef。KernelCall 和 KernelDef 之間通過?symbol?進行匹配。
StaticMemoryPlanningPass：將所有靜態(tài) shape 的 memref 進行內(nèi)存規(guī)劃，內(nèi)存規(guī)劃算法使用改進的 MegEngine 的內(nèi)存規(guī)劃算法--PushDown 算法，能夠極大程度的壓縮運行時內(nèi)存使用量。同時將 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 編譯之后的模型，以及生成計算模型所需的 Kernel 文件。 下面以一個簡單的模型為例，使用 MegCC 的輔助工具(下載?Release?包) mgb-importer 和 megcc-opt，觀察經(jīng)過各個 Pass 的處理 IR 的變化。也可使用 mgb-to-tinynn 工具直接完成模型的編譯過程，詳見?MegCC 入門文檔。
dump 模型（使用 megengine）
2. importer 模型到 MGB IR
可以看到，在 importer 的過程中，乘法運算和加法運算被融合成了"FUSE_MUL_ADD3"。
3. MGBToKernelPass、MemoryForwardingPass 和 StaticMemoryPlanningPass
經(jīng)過上面幾個 Pass，MGB IR 被轉(zhuǎn)換為了 Kernel IR 并進行了內(nèi)存規(guī)劃。感興趣的話可以更細(xì)粒度地看每個 Pass 做的事情，使用 megcc-opt 的參數(shù)控制使用哪些 Pass。
Kernel 生成
MegCC Compiler 會為模型中的每個 Operator 生成一個對應(yīng)的 Kernel 來完成計算。 目前 MegCC 中大多數(shù) Kernel 為人工優(yōu)化并提前寫好的 Kernel 模板，這些模板會根據(jù)具體的 Operator 參數(shù)生成對應(yīng)的 Kernel。大多數(shù)為人工優(yōu)化的 Kernel 的原因是：目前在 CPU 上不搜參的情況下，mlir 生成的 Kernel 性能和手寫的 Kernel 還有一定的距離，但是自動生成 Kernel 的方法長期來看是比較可取的。
MegCC 現(xiàn)已開源，倉庫地址：https://github.com/MegEngine/MegCC，歡迎試用、star、issue。
附：
更多 MegEngine 信息獲取，您可以查看：
文檔：https://www.megengine.org.cn/doc/stable/zh/?
GitHub 項目：https://github.com/MegEngine，或加入 MegEngine 用戶交流 QQ 群：1029741705
歡迎參與 MegEngine 社區(qū)貢獻(xiàn)，成為?Awesome MegEngineer：
?https://www.megengine.org.cn/community-AMGE
榮譽證書、定制禮品享不停。
參考
^The Deep Learning Compiler: A Comprehensive Survey. MINGZHEN LI, YI LIU, etc. 2020.
