深度學習編譯MLIR初步

深度模型的推理引擎

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

解釋型推理引擎

一般包含模型解析器，模型解釋器，模型優化器。

模型解析器負責讀取和解析模型文件，轉換為適用于解釋器處理的內存格式；
模型優化器負責將原始模型變換為等價的、但具有更快的推理速度的模型；
模型解釋器分析內存格式的模型并接受模型的輸入數據，然后根據模型的結構依次執行相應的模型內部的算子，最后產生模型的輸出。

編譯型推理引擎

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

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

編譯型推理引擎的優勢

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

中間表示

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

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框架中執行高性能機器學習模型所需的基礎設施進行統一，包括高性能計算技術應用或強化學習這類搜索算法的集成。MLIR旨在降低開發新硬件的成本，并提高現有TensorFlow用戶的可用性。

在這里插入圖片描述

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

表征數據流圖的能力，包括動態形狀（dynamic shape），用戶可擴展的操作集，TensorFlow變量，等等。
支持典型的圖優化和圖變換，例如TensorFlow Grappler。
可以以一種易于優化的形式來表征ML操作的Kernel。
具備掌控跨Kernel的高性能計算風格（high-performance-computing-style）循環優化（例如fusion，loop interchange，tiling等polyhedral compilation）的能力，以及變換數據內存排布的能力。
代碼生成“下降lowering”轉換，顯式的插入DMA，顯示的緩存cache管理，內存分塊tiling，為適應1D和2D寄存器架構的自動向量化
表達特定后端操作的能力，比如特定加速器的高層操作。
在深度學習圖上已有的量化和其他圖變化。

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

2. MLIR的作用

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

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

在這里插入圖片描述

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

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

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

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

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

在這里插入圖片描述

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

在這里插入圖片描述

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

3. MLIR中的 “方言”：dialect

為什么要有方言

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

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

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

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

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

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

將方言添加至 MLIR

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

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

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

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

dialect是如何工作的

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

在這里插入圖片描述