0x1. 前言

最近開(kāi)始做一些MLIR的工作，所以開(kāi)始學(xué)習(xí)MLIR的知識(shí)。這篇筆記是對(duì)MLIR的初步印象，并不深入，適合想初步了解MLIR是什么的同學(xué)閱讀，后面會(huì)繼續(xù)分享MLIR的一些項(xiàng)目。這里要大力感謝中科院的法斯特豪斯（知乎ID）同學(xué)先前的一些分享，給了我入門MLIR的方向。

0x2. 什么是IR？

IR即中間表示（Intermediate Representation），可以看作是一種中介的數(shù)據(jù)格式，便于模型在框架間轉(zhuǎn)換。在深度學(xué)習(xí)中可以表示計(jì)算圖的數(shù)據(jù)結(jié)構(gòu)就可以稱作一種IR，例如ONNX，TorchScript，TVM Relay等等。這里舉幾個(gè)例子介紹一下：

首先，ONNX是微軟和FaceBook提出的一種IR，他持有了一套標(biāo)準(zhǔn)化算子格式。無(wú)論你使用哪種深度學(xué)習(xí)框架（Pytorch，TensorFlow，OneFlow）都可以將計(jì)算圖轉(zhuǎn)換成ONNX進(jìn)行存儲(chǔ)。然后各個(gè)部署框架只需要支持ONNX模型格式就可以簡(jiǎn)單的部署各個(gè)框架訓(xùn)練的模型了，解決了各個(gè)框架之間模型互轉(zhuǎn)的復(fù)雜問(wèn)題。

但ONNX設(shè)計(jì)沒(méi)有考慮到一個(gè)問(wèn)題，那就是各個(gè)框架的算子功能和實(shí)現(xiàn)并不是統(tǒng)一的。ONNX要支持所有框架所有版本的算子實(shí)現(xiàn)是不現(xiàn)實(shí)的，目前ONNX的算子版本已經(jīng)有10代以上，這讓用戶非常痛苦。IR可以類比為計(jì)算機(jī)架構(gòu)的指令集，但我們是肯定不能接受指令集頻繁改動(dòng)的。另外ONNX有一些控制流的算子如If，但支持得也很有限。

其次，TorchScript是Pytorch推出的一種IR，它是用來(lái)解決動(dòng)態(tài)圖模式執(zhí)行代碼速度太慢的問(wèn)題。因?yàn)閯?dòng)態(tài)圖模式在每次執(zhí)行計(jì)算時(shí)都需要重新構(gòu)造計(jì)算圖(define by run)，使得性能和可移植性都比較差。為了解決這個(gè)問(wèn)題，Pytorch引入了即時(shí)編譯（JIT）技術(shù)即TorchScript來(lái)解決這一問(wèn)題。Pytorch早在1.0版本就引入了JIT技術(shù)并開(kāi)放了C++ API，用戶之后就可以使用Python編寫(xiě)的動(dòng)態(tài)圖代碼訓(xùn)練模型然后利用JIT將模型（nn.Module）轉(zhuǎn)換為語(yǔ)言無(wú)關(guān)的模型（TorchScript），使得C++ API可以方便的調(diào)用。并且TorchScript也很好的支持了控制流，即用戶在Python層寫(xiě)的控制流可以在TorchScript模型中保存下來(lái)，是Pytorch主推的IR。

最后，Relay IR是一個(gè)函數(shù)式、可微的、靜態(tài)的、針對(duì)機(jī)器學(xué)習(xí)的領(lǐng)域定制編程語(yǔ)言。Relay IR解決了普通DL框架不支持control flow（或者要借用python 的control flow，典型的比如TorchScript）以及dynamic shape的特點(diǎn)，使用lambda calculus作為基準(zhǔn)IR。

Relay IR可以看成一門編程語(yǔ)言，在靈活性上比ONNX更強(qiáng)。但Relay IR并不是一個(gè)獨(dú)立的IR，它和TVM相互耦合，這使得用戶想使用Relay IR就需要基于TVM進(jìn)行開(kāi)發(fā)，這對(duì)一些用戶來(lái)說(shuō)是不可接受的。

這幾個(gè)例子就是想要說(shuō)明，深度學(xué)習(xí)中的IR只是一個(gè)深度學(xué)習(xí)框架，公司甚至是一個(gè)人定義的一種中介數(shù)據(jù)格式，它可以表示深度學(xué)習(xí)中的模型（由算子和數(shù)據(jù)構(gòu)成）那么這種格式就是IR。

0x3. 為什么要引入MLIR？

目前深度學(xué)習(xí)領(lǐng)域的IR數(shù)量眾多，很難有一個(gè)IR可以統(tǒng)一其它的IR，這種百花齊放的局面就造成了一些困境。我認(rèn)為中科院的法斯特豪斯同學(xué)B站視頻舉的例子非常好，建議大家去看一下。這里說(shuō)下我的理解，以TensorFlow Graph為例，它可以直接被轉(zhuǎn)換到TensorRT的IR，nGraph IR，CoreML IR，TensorFlow Lite IR來(lái)直接進(jìn)行部署。或者TensorFlow Graph可以被轉(zhuǎn)為XLA HLO，然后用XLA編譯器來(lái)對(duì)其進(jìn)行Graph級(jí)別的優(yōu)化得到優(yōu)化后的XLA HLO，這個(gè)XLA HLO被喂給XLA編譯器的后端進(jìn)行硬件綁定式優(yōu)化和Codegen。在這個(gè)過(guò)程中主要存在兩個(gè)問(wèn)題。

• 第一，IR的數(shù)量太多，開(kāi)源要維護(hù)這么多套IR，每種IR都有自己的圖優(yōu)化Pass，這些Pass可能實(shí)現(xiàn)的功能是一樣的，但無(wú)法在兩種不同的IR中直接遷移。假設(shè)深度學(xué)習(xí)模型對(duì)應(yīng)的DAG一共有10種圖層優(yōu)化Pass，要是為每種IR都實(shí)現(xiàn)10種圖層優(yōu)化Pass，那工作量是巨大的。
• 第二，如果出現(xiàn)了一種新的IR，開(kāi)發(fā)者想把另外一種IR的圖層優(yōu)化Pass遷移過(guò)來(lái)，但由于這兩種IR語(yǔ)法表示完全不同，除了借鑒優(yōu)化Pass的思路之外，就絲毫不能從另外一種IR的Pass實(shí)現(xiàn)受益了，即互相遷移的難度比較大。此外，如果你想為一個(gè)IR添加一個(gè)Pass，難度也是不小的。舉個(gè)例子你可以嘗試為onnx添加一個(gè)圖優(yōu)化Pass，會(huì)發(fā)現(xiàn)這并不是一件簡(jiǎn)單的事，甚至需要我們?nèi)ポ^為完整的學(xué)習(xí)ONNX源碼。
• 第三，在上面的例子中優(yōu)化后的XLA HLO直接被喂給XLA編譯器后端產(chǎn)生LLVM IR然后Codegen，這個(gè)跨度是非常大的。這里怎么理解呢？我想到了一個(gè)例子。以優(yōu)化GEMM來(lái)看，我們第一天學(xué)會(huì)三重for循環(huán)寫(xiě)一個(gè)naive的矩陣乘程序，然后第二天你就要求我用匯編做一個(gè)優(yōu)化程度比較高的矩陣乘法程序？那我肯定是一臉懵逼的，只能git clone了，當(dāng)然是學(xué)不會(huì)的。但如果你緩和一些，讓我第二天去了解并行，第三天去了解分塊，再給幾天學(xué)習(xí)一下SIMD，再給幾個(gè)月學(xué)習(xí)下匯編，沒(méi)準(zhǔn)一年下來(lái)我就可以真正的用匯編優(yōu)化一個(gè)矩陣乘法了。所以跨度太大最大的問(wèn)題在于，我們這種新手玩家很難參與。我之前分享過(guò)TVM的Codegen流程，雖然看起來(lái)理清了Codegen的調(diào)用鏈，但讓我現(xiàn)在自己去實(shí)現(xiàn)一個(gè)完整的Codegen流程，那我是很難做到的。【從零開(kāi)始學(xué)深度學(xué)習(xí)編譯器】九，TVM的CodeGen流程

針對(duì)上面的問(wèn)題，MLIR（Multi-Level Intermediate Representation）被提出。MLIR是由LLVM團(tuán)隊(duì)開(kāi)發(fā)和維護(hù)的一套編譯器基礎(chǔ)設(shè)施，它強(qiáng)調(diào)工具鏈的可重用性和可擴(kuò)展性。下面我們具體分析一下：

針對(duì)第一個(gè)問(wèn)題和第二個(gè)問(wèn)題，造成這些深度學(xué)習(xí)領(lǐng)域IR的優(yōu)化Pass不能統(tǒng)一的原因就是因?yàn)樗鼈儧](méi)有一個(gè)統(tǒng)一的表示，互轉(zhuǎn)的難度高。因此MLIR提出了Dialect，我們可以將其理解為各種IR需要學(xué)習(xí)的語(yǔ)言，一旦某種IR學(xué)會(huì)這種語(yǔ)言，就可以基于這種語(yǔ)言將其重寫(xiě)為MLIR。Dialect將所有IR都放在了同一個(gè)命名空間里面，分別對(duì)每個(gè)IR定義對(duì)應(yīng)的產(chǎn)生式以及綁定對(duì)應(yīng)的操作，從而生成MLIR模型。關(guān)于Dialect我們后面會(huì)細(xì)講，這篇文章先提一下，它是MLIR的核心組件之一。

針對(duì)第三個(gè)問(wèn)題，怎么解決IR跨度大的問(wèn)題？MLIR通過(guò)Dialect抽象出了多種不同級(jí)別的MLIR，下面展示官方提供的一些MLIR IR抽象，我們可以看到Dialect是對(duì)某一類IR或者一些數(shù)據(jù)結(jié)構(gòu)相關(guān)操作進(jìn)行抽象，比如llvm dialect就是對(duì)LLVM IR的抽象，tensor dialect就是對(duì)Tensor這種數(shù)據(jù)結(jié)構(gòu)和操作進(jìn)行抽象：

除了這些，各種深度學(xué)習(xí)框架都在接入MLIR，比如TensorFlow，Pytorch，OneFlow以及ONNX等等，大家都能在github找到對(duì)應(yīng)工程。

抽象了多個(gè)級(jí)別的IR好處是什么呢？這就要結(jié)合MLIR的編譯流程來(lái)看，MLIR的編譯流程大致如下：

對(duì)于一個(gè)源程序，首先經(jīng)過(guò)語(yǔ)法樹(shù)分析，然后通過(guò)Dialect將其下降為MLIR表達(dá)式，再經(jīng)MLIR分析器得到目標(biāo)程序。注意這個(gè)目標(biāo)程序不一定是可運(yùn)行的程序。比如假設(shè)第一次的目標(biāo)程序是C語(yǔ)言程序，那么它可以作為下一次編譯流程的源程序，通過(guò)Dialect下降為L(zhǎng)LVM MLIR。這個(gè)LLVM MLIR即可以被MLIR中的JIT執(zhí)行，也可以通過(guò)Dialect繼續(xù)下降，下降到三地址碼IR對(duì)應(yīng)的MLIR，再被MLIR分析器解析獲得可執(zhí)行的機(jī)器碼。

因此MLIR這個(gè)多級(jí)別的下降過(guò)程就類似于我們剛才介紹的可以漸進(jìn)式學(xué)習(xí)，解決了IR到之間跨度太大的問(wèn)題。比如我們不熟悉LLVM IR之后的層次，沒(méi)有關(guān)系，我們交給LLVM編譯器，我們?nèi)ネ瓿汕懊婺遣糠值腄ialect實(shí)現(xiàn)就可以了。

0x4. 總結(jié)

這篇文章簡(jiǎn)單聊了一下對(duì)MLIR的粗淺理解，歡迎大家批評(píng)指正。