深度學習編譯MLIR初步

深度模型的推理引擎

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

解釋型推理引擎

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

• 模型解析器負責讀取和解析模型文件，轉換為適用于解釋器處理的內存格式；

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

• 模型解釋器分析內存格式的模型并接受模型的輸入數據，然后根據模型的結構依次執(zhí)行相應的模型內部的算子，最后產生模型的輸出。

編譯型推理引擎

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

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

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

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

中間表示

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

MLIR

聽到這個名字的時候，下意識覺得這應該是個關于機器學習的中間表達結構，畢竟這年頭聽到ML都會認為是Machine Learning，但是MLIR還真不是Machine Learning Intermediate Representation，而是Multi-Level Intermediate Representation。

1. 什么是MLIR

MLIR（Multi-Level Intermediate Representation，多級中間表示）是一種用來構建可重用和可擴展編譯基礎設施的新方法。MLIR旨在解決軟件碎片化，改進異構硬件的編譯，顯著減少構建特定領域編譯器的成本，并幫助連接現有的編譯器。

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

MLIR通過一種混合的中間表示，實現在一個統(tǒng)一的架構基礎上支持以下多個不同的需求：

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

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

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

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

代碼生成“下降lowering”轉換，顯式的插入DMA，顯示的緩存cache管理，內存分塊tiling，為適應1D和2D寄存器架構的自動向量化

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

在深度學習圖上已有的量化和其他圖變化。

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

2. MLIR的作用

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

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

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

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

此外，我們有時甚至會采用更復雜的途徑，包括在每層中執(zhí)行多輪優(yōu)化。例如，Grappler 框架現在便能優(yōu)化 TensorFlow 中的張量布局和運算。

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

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

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

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

3. MLIR中的 “方言”：dialect

為什么要有方言

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

從源程序到目標程序，要經過一系列的抽象以及分析，通過Lowering Pass來實現從一個IR到另一個IR的轉換，這樣的過程中會存在有些操作重復實現的情況，也就導致了轉換效率低的問題。

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

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

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

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

將方言添加至 MLIR

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

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

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

雖然 MLIR 充當 ML 的編譯器，但我們也看到，MLIR 同樣支持在編譯器內部使用機器學習技術！這一點尤為重要，因為在進行擴展時，開發(fā)數字庫的工程師無法跟上 ML 模型或硬件的多樣化速度。MLIR 的擴展性有助于探索代碼降階策略，并在抽象之間執(zhí)行逐步降階。

dialect是如何工作的

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

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

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

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。