深度學(xué)習(xí)編譯MLIR初步

深度模型的推理引擎

目前深度模型的推理引擎按照實(shí)現(xiàn)方式大體分為兩類：解釋型推理引擎和編譯型推理引擎。

解釋型推理引擎

一般包含模型解析器，模型解釋器，模型優(yōu)化器。

• 模型解析器負(fù)責(zé)讀取和解析模型文件，轉(zhuǎn)換為適用于解釋器處理的內(nèi)存格式；

• 模型優(yōu)化器負(fù)責(zé)將原始模型變換為等價(jià)的、但具有更快的推理速度的模型；

• 模型解釋器分析內(nèi)存格式的模型并接受模型的輸入數(shù)據(jù)，然后根據(jù)模型的結(jié)構(gòu)依次執(zhí)行相應(yīng)的模型內(nèi)部的算子，最后產(chǎn)生模型的輸出。

編譯型推理引擎

編譯型推理引擎一般包含模型解析器和模型編譯器。

• 模型解析器的作用與解釋型推理引擎相同
• 模型編譯器負(fù)責(zé)將模型編譯為計(jì)算設(shè)備（CPU、GPU 等）可直接處理的機(jī)器碼，并且可能在編譯的過程中應(yīng)用各種優(yōu)化方法來提高生成的機(jī)器碼的效率。

編譯型推理引擎的優(yōu)勢

由于機(jī)器碼的模型可以直接被計(jì)算設(shè)備處理，無需額外的解釋器的參與，其消除了解釋器調(diào)度的開銷。相對于解釋型推理引擎，由于生成機(jī)器碼的過程更加靠底層，編譯器有更多的優(yōu)化機(jī)會以達(dá)到更高的執(zhí)行效率。由于現(xiàn)在業(yè)界對于推理引擎的執(zhí)行速度有了更高的需求，編譯型推理引擎也逐漸成為高速推理引擎的發(fā)展方向。編譯型推理引擎有 Apache TVM、oneDNN、PlaidML、TensorFlow XLA、TensorFlow Runtime 等。

中間表示

為了便于優(yōu)化，一般來說推理引擎會把模型轉(zhuǎn)換為中間表示，然后對中間表示進(jìn)行優(yōu)化和變換，最終生成目標(biāo)模型（對于解釋型推理引擎）或目標(biāo)機(jī)器碼（對于編譯型推理引擎）。此外，除了深度學(xué)習(xí)領(lǐng)域，在很早以前編程語言領(lǐng)域就引入了中間表示來做優(yōu)化和變換。而新的編程語言層出不窮，因此就出現(xiàn)了各種各樣的中間表示：不同的推理引擎或者編譯器都會有自己的中間表示和優(yōu)化方案，而每種中間表示和優(yōu)化方案可能都需要從頭實(shí)現(xiàn)，最終可能會導(dǎo)致軟件的碎片化和重復(fù)的開發(fā)工作。

MLIR

聽到這個(gè)名字的時(shí)候，下意識覺得這應(yīng)該是個(gè)關(guān)于機(jī)器學(xué)習(xí)的中間表達(dá)結(jié)構(gòu)，畢竟這年頭聽到ML都會認(rèn)為是Machine Learning，但是MLIR還真不是Machine Learning Intermediate Representation，而是Multi-Level Intermediate Representation。

1. 什么是MLIR

MLIR（Multi-Level Intermediate Representation，多級中間表示）是一種用來構(gòu)建可重用和可擴(kuò)展編譯基礎(chǔ)設(shè)施的新方法。MLIR旨在解決軟件碎片化，改進(jìn)異構(gòu)硬件的編譯，顯著減少構(gòu)建特定領(lǐng)域編譯器的成本，并幫助連接現(xiàn)有的編譯器。

具體來說，MLIR是通過定義一個(gè)通用的中間表示（IR），將在TensorFlow和類似的ML框架中執(zhí)行高性能機(jī)器學(xué)習(xí)模型所需的基礎(chǔ)設(shè)施進(jìn)行統(tǒng)一，包括高性能計(jì)算技術(shù)應(yīng)用或強(qiáng)化學(xué)習(xí)這類搜索算法的集成。MLIR旨在降低開發(fā)新硬件的成本，并提高現(xiàn)有TensorFlow用戶的可用性。

MLIR通過一種混合的中間表示，實(shí)現(xiàn)在一個(gè)統(tǒng)一的架構(gòu)基礎(chǔ)上支持以下多個(gè)不同的需求：

表征數(shù)據(jù)流圖的能力，包括動態(tài)形狀（dynamic shape），用戶可擴(kuò)展的操作集，TensorFlow變量，等等。

支持典型的圖優(yōu)化和圖變換，例如TensorFlow Grappler。

可以以一種易于優(yōu)化的形式來表征ML操作的Kernel。

具備掌控跨Kernel的高性能計(jì)算風(fēng)格（high-performance-computing-style）循環(huán)優(yōu)化（例如fusion，loop interchange，tiling等polyhedral compilation）的能力，以及變換數(shù)據(jù)內(nèi)存排布的能力。

代碼生成“下降lowering”轉(zhuǎn)換，顯式的插入DMA，顯示的緩存cache管理，內(nèi)存分塊tiling，為適應(yīng)1D和2D寄存器架構(gòu)的自動向量化

表達(dá)特定后端操作的能力，比如特定加速器的高層操作。

在深度學(xué)習(xí)圖上已有的量化和其他圖變化。

MLIR 并非萬能，不支持底層機(jī)器代碼生成的算法（比如寄存器分配和指令調(diào)度），這些更適合于低層優(yōu)化器（如llvm），它也不支持用戶手寫Kernel（不同于CUDA）。MLIR是提供一種框架，使自定義的 DSL（Domain-Specific Language，領(lǐng)域?qū)Ｓ谜Z言） 表達(dá)可以融入到一個(gè)統(tǒng)一的生態(tài)。

2. MLIR的作用

雖然 MLIR 中的 ML 不是Machine Learning，但是Machine Learning確實(shí)是是MLIR的一個(gè)重要應(yīng)用領(lǐng)域。我們接下來主要看一下在機(jī)器學(xué)習(xí)領(lǐng)域，MLIR可以做哪些事情。在了解MLIR是怎么工作之前，我先得弄明白這個(gè)IR在解決什么問題。

說到機(jī)器學(xué)習(xí)，我們就用TensorFlow這個(gè)框架來舉例。我們知道TensorFlow是使用數(shù)據(jù)流圖作為數(shù)據(jù)結(jié)構(gòu)來進(jìn)行各種數(shù)值計(jì)算，要讓這些計(jì)算運(yùn)行在硬件上，我們需要一個(gè)TensorFlow的編譯生態(tài)系統(tǒng)：

如圖中所示，TensorFlow圖 能夠以多種不同的方式運(yùn)行，包括：

• 將其發(fā)送至調(diào)用手寫運(yùn)算內(nèi)核的 TensorFlow 執(zhí)行器
• 將圖轉(zhuǎn)化為 XLA 高級優(yōu)化器 (XLA HLO) 表示，反之，這種表示亦可調(diào)用適合 CPU 或 GPU 的 LLVM 編輯器，或者繼續(xù)使用適合 TPU 的 XLA。（或者將二者結(jié)合！）
• 將圖轉(zhuǎn)化為 TensorRT、nGraph 或另一種適合特定硬件指令集的編譯器格式
• 將圖轉(zhuǎn)化為 TensorFlow Lite 格式，然后在 TensorFlow Lite 運(yùn)行時(shí)內(nèi)部執(zhí)行此圖，或者通過 Android 神經(jīng)網(wǎng)絡(luò) API (NNAPI) 或相關(guān)技術(shù)將其進(jìn)一步轉(zhuǎn)化，以在 GPU 或 DSP 上運(yùn)行

此外，我們有時(shí)甚至?xí)捎酶鼜?fù)雜的途徑，包括在每層中執(zhí)行多輪優(yōu)化。例如，Grappler 框架現(xiàn)在便能優(yōu)化 TensorFlow 中的張量布局和運(yùn)算。

通常來說，整個(gè)編譯流程先將TensorFlow的圖轉(zhuǎn)化為XLA HLO，即一種類似高級語言的圖的中間表達(dá)形式，可以基于此進(jìn)行一些 High-Level 的優(yōu)化。接著將XLA HLO翻譯為LLVM IR，使用LLVM編譯到各種硬件的匯編語言，從而運(yùn)行在硬件上進(jìn)行數(shù)值計(jì)算。

下圖的藍(lán)色陰影部分是 基于圖的IR，綠色陰影部分是基于 SSA（static single-assignment，靜態(tài)單一賦值） 的IR，然而這樣的編譯方式的缺點(diǎn)在于構(gòu)建這樣的編譯系統(tǒng)的開銷比較大，每一層的設(shè)計(jì)實(shí)現(xiàn)會有重復(fù)部分，同一個(gè)層次的IR彼此之間雖然相似，但是存在天生的“生殖隔離”，升級優(yōu)化缺乏遷移性，即改變優(yōu)化一個(gè)模塊，并不能惠及到同層次的其他模塊。因此，目前存在的問題就在于各種IR之間轉(zhuǎn)換的效率和可遷移性不高。

對于上述問題，MLIR希望為各種DSL提供一種中間表達(dá)形式，將他們集成為一套生態(tài)系統(tǒng)，使用一種一致性強(qiáng)的方式編譯到特定硬件平臺的匯編語言上。利用這樣的形式，MLIR就可以利用它模塊化、可擴(kuò)展的特點(diǎn)來解決IR之間相互配合的問題。

到此為止，我們大致知道了MLIR的誕生是為了解決什么問題。目前它對我來說還是一個(gè)黑盒子，下面的工作就是要去看看MLIR內(nèi)部究竟是一個(gè)什么樣的結(jié)構(gòu)，看看它是怎么把各層IR整合到一起，又是如何實(shí)現(xiàn)擴(kuò)展功能的。

3. MLIR中的 “方言”：dialect

為什么要有方言

之前我們說到當(dāng)前的編譯結(jié)構(gòu)的問題在于各種IR之間轉(zhuǎn)換的效率和可遷移性不高。MLIR試圖使用一種一致性強(qiáng)的方式，為各種DSL提供一種中間表達(dá)形式，將他們集成為一套生態(tài)系統(tǒng)，編譯到特定硬件平臺的匯編語言上。這樣的目標(biāo)是通過什么手段實(shí)現(xiàn)的呢？

從源程序到目標(biāo)程序，要經(jīng)過一系列的抽象以及分析，通過Lowering Pass來實(shí)現(xiàn)從一個(gè)IR到另一個(gè)IR的轉(zhuǎn)換，這樣的過程中會存在有些操作重復(fù)實(shí)現(xiàn)的情況，也就導(dǎo)致了轉(zhuǎn)換效率低的問題。

這就好比，IR們組成一個(gè)流水線要合起伙來干一個(gè)大買賣，但是互相配合不默契，誰也明白不了對方究竟干了啥，為了保險(xiǎn)起見，每個(gè)IR拿到上一個(gè)IR的產(chǎn)品之后只能多干點(diǎn)活保證不出錯(cuò)，這樣一來效率自然就低了。MLIR面對這種IR群雄割據(jù)的現(xiàn)狀，打算一統(tǒng)天下！打天下容易，守天下難呀，不讓我們用各種IR了，你倒是給我們一條活路呀，怎么才能讓源語言變成匯編語言然后跑在機(jī)器上呀？于是，統(tǒng)一IR的第一步就是要統(tǒng)一“語言”，各個(gè)IR原來配合不默契，誰也理解不了誰，就是因?yàn)椤罢Z言”不通，沒法用統(tǒng)一的“語言”指揮流水線干活。MLIR看準(zhǔn)時(shí)機(jī)拿出了自己的法寶：Dialects！讓各個(gè)IR學(xué)習(xí)Dialects這個(gè)“語言”，這樣一來，不光能指揮流水線高效干活了，還能隨意擴(kuò)展更改分工，從此IR們就可以完美地分工協(xié)作。

為區(qū)分不同的硬件與軟件受眾，MLIR 提供 “方言”，其中包括：

• TensorFlow IR，代表 TensorFlow 圖中可能存在的一切
• XLA HLO IR，旨在利用 XLA 的編譯功能（輸出到 TPU 等）
• 實(shí)驗(yàn)性仿射方言，側(cè)重于多面表示與優(yōu)化
• LLVM IR，與 LLVM 自我表示之間存在 1:1 映射，可使 MLIR 通過 LLVM 發(fā)出 GPU 與 CPU 代碼
• TensorFlow Lite，將會轉(zhuǎn)換以在移動平臺上運(yùn)行代碼

每種方言均由一組存在不變性的已定義操作組成，如：“這是一個(gè)二進(jìn)制運(yùn)算符，輸入與輸出擁有相同類型。”

將方言添加至 MLIR

MLIR 沒有眾所周知的固定或內(nèi)置的操作列表（無 “內(nèi)聯(lián)函數(shù)”）。方言可完全定義自定義類型，即 MLIR 如何對 LLVM IR 類型系統(tǒng)（擁有一流匯總）、域抽象（對量化類型等經(jīng)機(jī)器學(xué)習(xí) (ML) 優(yōu)化的加速器有著重要意義），乃至未來的 Swift 或 Clang 類型系統(tǒng)（圍繞 Swift 或 Clang 聲明節(jié)點(diǎn)而構(gòu)建）進(jìn)行建模。

如果我們想要連接新的低級編譯器，則需要創(chuàng)建新方言，以及 TensorFlow 圖方言與我們的方言之間的降階。如此一來，硬件及編譯器制造商便可暢通無阻。我們甚至可以在同一個(gè)模型中定位不同級別的方言；高級優(yōu)化器將保留 IR 中不熟悉的部分，并等待較低級別的優(yōu)化器來處理此類部分。

對于編譯器研究者和框架制造者，則可以借助 MLIR 在每個(gè)級別進(jìn)行轉(zhuǎn)換，甚至是在 IR 中定義自己的操作和抽象，從而針對試圖解決的問題領(lǐng)域構(gòu)建最佳模型。由此看來，MLIR 比 LLVM 更像是純編譯器基礎(chǔ)設(shè)施。

雖然 MLIR 充當(dāng) ML 的編譯器，但我們也看到，MLIR 同樣支持在編譯器內(nèi)部使用機(jī)器學(xué)習(xí)技術(shù)！這一點(diǎn)尤為重要，因?yàn)樵谶M(jìn)行擴(kuò)展時(shí)，開發(fā)數(shù)字庫的工程師無法跟上 ML 模型或硬件的多樣化速度。MLIR 的擴(kuò)展性有助于探索代碼降階策略，并在抽象之間執(zhí)行逐步降階。

dialect是如何工作的

dialects是將所有的IR放在了同一個(gè)命名空間中，分別對每個(gè)IR定義對應(yīng)的產(chǎn)生式以及綁定相應(yīng)的操作，從而生成一個(gè)MLIR的模型。整個(gè)的編譯過程，從源語言生成AST（Abstract Syntax Tree，抽象語法樹），借助Dialects遍歷AST，產(chǎn)生MLIR的表達(dá)式，此處可為多層IR通過Lowering Pass依次進(jìn)行分析，最后經(jīng)過MLIR分析器，生成目標(biāo)語言。

這里只是簡單介紹一下MLIR中dialect的工作機(jī)制，參考知乎@法斯特豪斯的理解，之后會更加詳細(xì)的進(jìn)行介紹并增加實(shí)例，如有錯(cuò)誤，歡迎討論。

Ref

https://zhuanlan.zhihu.com/p/101879367

https://zhuanlan.zhihu.com/p/102212806

https://www.sohu.com/a/307133340_670669

https://blog.csdn.net/wujianing_110117/article/details/119312999

總結(jié)

以上是生活随笔為你收集整理的深度学习编译：MLIR初步的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。