MLIR中間表示與編譯
概述
目前深度模型的推理引擎按照實現方式大體分為兩類：
解
解釋型推理引擎：
- 一般包含模型解析器，模型解釋器，模型優化器。
- 模型解析器負責讀取和解析模型文件，轉換為適用於解釋器處理的內存格式；
- 模型優化器負責將原始模型變換為等價的、但具有更快的推理速度的模型；
- 模型解釋器分析內存格式的模型並接受模型的輸入數據，然后根據模型的結構依次執行相應的模型內部的算子，最后產生模型的輸出。
編譯型推理引擎：
- 一般包含模型解析器和模型編譯器。模型解析器的作用與解釋型推理引擎相同；
- 模型編譯器負責將模型編譯為計算設備（CPU、GPU 等）可直接處理的機器碼，並且可能在編譯的過程中應用各種優化方法來提高生成的機器碼的效率。
- 由於機器碼的模型可以直接被計算設備處理，無需額外的解釋器的參與，其消除了解釋器調度的開銷。
- 相對於解釋型推理引擎，由於生成機器碼的過程更加靠底層，編譯器有更多的優化機會以達到更高的執行效率。
由於現在業界對於推理引擎的執行速度有了更高的需求，編譯型推理引擎也逐漸成為高速推理引擎的發展方向。編譯型推理引擎有 Apache TVM、oneDNN、PlaidML、TensorFlow XLA、TensorFlow Runtime 等。
為了便於優化，一般來說推理引擎會把模型轉換為中間表示，然后對中間表示進行優化和變換，最終生成目標模型（對於解釋型推理引擎）或目標機器碼（對於編譯型推理引擎）。此外，除了深度學習領域，在很早以前編程語言領域就引入了中間表示來做優化和變換。而新的編程語言層出不窮，因此就出現了各種各樣的中間表示：
不同的推理引擎或者編譯器都會有自己的中間表示和優化方案，而每種中間表示和優化方案可能都需要從頭實現，最終可能會導致軟件的碎片化和重復的開發工作。
MLIR 簡介
MLIR（Multi-Level Intermediate Representation）是一種新型的用於構建可復用和可擴展的編譯器的框架。MLIR 旨在解決軟件碎片化、改善異構硬件的編譯、降低構建領域特定編譯器的成本，以及幫助將現有的編譯器連接到一起。
MLIR 旨在成為一種在統一的基礎架構中，支持多種不同需求的混合中間表示，例如：
表示數據流圖（例如在 TensorFlow 中）的能力，包括動態性狀、用戶可擴展的算子生態系統、TensorFlow 變量等。
在這些圖中進行優化和變換（例如在 Grappler 中）。
適合優化的形式的機器學習算子內核的表示。
能夠承載跨內核的高性能計算風格的循環優化（融合、循環交換、分塊等），並能夠變換數據的內存布局。
代碼生成“下降”變換，例如 DMA 插入、顯式緩存管理、內存分塊，以及 1 維和 2 維寄存器架構的向量化。
表示目標特定算子的能力，例如加速器特定的高層算子。
在深度學習圖中的做的量化和其它圖變換。
MLIR 是一種支持硬件特定算子的通用中間表示。因此，對圍繞 MLIR 的基礎架構進行的任何投入（例如在編譯器 pass 上的工作）都將產生良好的回報；許多目標都可以使用該基礎架構，並從中受益。
盡管 MLIR 是一種強大的框架，也有一些非目標。MLIR 不試圖去支持底層機器碼生成算法（如寄存器分配和指令調度）。這些更適合於底層優化器（例如 LLVM）。此外，MLIR 也不意圖成為最終用戶寫算子內核的源語言（類似於 CUDA 和 C++）。另一方面，MLIR 提供了用於表示此類領域特定語言並將其集成到生態系統中的支柱。
MLIR 在構建時受益於從構建其它中間表示（LLVM IR、XLA HLO 和 Swift SIL）的過程中獲得的經驗。MLIR 框架鼓勵現存的最佳實踐，例如：編寫和維護中間表示規范、構建中間表示驗證器、提供將 MLIR 文件轉儲和解析為文本的功能、使用 FileCheck 工具編寫詳盡的單元測試、以及以一組可以以新的方式組合的模塊化庫的形式構建基礎框架。
其它的經驗教訓也已經整合到了設計中。例如，LLVM 有一個不明顯的設計錯誤，其會阻止多線程編譯器同時處理 LLVM 模塊中的多個函數。MLIR 通過限制 SSA 作用域來減少使用-定義鏈，並用顯式的符號引用代替跨函數引用來解決這些問題。
MLIR 方言（Dialect）
MLIR 通過“方言”來定義不同層次的中間表示，每一個方言都有自己唯一的名字空間。開發者可以創建自定義方言，並在方言內部定義算子、類型和屬性，以及語義。MLIR 推薦使用方言來對 MLIR 進行擴展。有這樣一個統一的中間表示框架，降低了開發新的編譯器的成本。除了可以使用 C++ 語言，對方言進行定義之外，MLIR 也提供了一種聲明式的方式來定義方言，即用戶通過編寫 TableGen 格式的文件來定義方言，然后使用 TableGen 工具來生成對應的 C++ 頭文件和源文件，以及對應的文檔。MLIR 也推薦使用這種聲明式的的方式來定義方言。此外，MLIR 也提供了一個框架用於在方言之間或者方言內部進行轉換。
為了方便開發，MLIR 也內置了一些方言可供直接使用：
acc
affine
async
avx512
gpu
linalg
llvm
nvvm
omp
pdl
pdl_interp
quant
rocdl
scf
shape
spv
std
vector
MLIR 使用“算子”來描述不同層次的抽象和計算。MLIR 中的算子也是可擴展的，用戶可以創建自定義的算子並規定其語義。例如目標無關算子、仿射算子和目標特定算子。MLIR 也支持用戶通過聲明式的方式（TableGen）來創建自定義算子。
MLIR 中的每個值都有其對應的“類型”，MLIR 內置了一些原始類型（比如整數）和聚合類型（張量和內存緩沖區）。MLIR 的類型系統也允許用戶對其進行擴展，創建自定義的類型以及規定其語義。
此外在 MLIR 中，用戶可以通過指定算子的“屬性”的值來控制算子的行為。算子可以定義自身的屬性，比如卷積算子的 stride 屬性等。
方言的變換
在 MLIR 中定義算子的時候可以定義其規范化的行為，比如將 x + 2 和 2 + x 統一規范化為 x + 2，以便后續的優化過程更為方便地進行。MLIR 以一種貪婪地策略，不斷地應用規范化變換，直到中間表示收斂為止。
在 MLIR 中進行方言內部或方言之間的轉換時，用戶首先要定義一個轉換目標。轉換目標規定了生成的目標中可以出現哪些算子。然后用戶需要指定一組重寫模式，這些重寫模式定義了算子之間的轉換關系。最后框架根據用戶指定的轉換目標和重寫模式執行轉換。這個轉換過程會自動檢測轉換方式，例如如果指定了 A → B 和 B → C 的重寫模式，框架會自動完成 A → C 的轉換過程。MLIR 也支持用戶通過聲明式的方式（TableGen），創建自定義的重寫模式。當轉換的方言之間有着不同的類型系統，用戶可以使用類型轉換器，完成類型之間的轉換。
MLIR 的用戶
ONNX MLIR：將 ONNX 格式的深度學習網絡模型轉換為能在不同二進制執行的二進制格式。
PlaidML：一個開源的張量編譯器，允許在不同的硬件平台上運行深度學習模型。
TensorFlow：TensorFlow 項目的 XLA 和 TensorFlow Lite 模型轉換器用到了 MLIR。
TensorFlow Runtime：一種新的 TensorFlow 運行時。
Verona：一種新的研究型的編程語言，用於探索並發所有權。其提供了一個可以與所有權無縫集成新的並發模型。
附錄：編譯和安裝 MLIR
下載 MLIR
MLIR 是 LLVM 項目的子項目，要編譯 MLIR，首先獲取 LLVM 的源代碼。
LLVM 的源碼可從 GitHub 獲取：
git clone https://github.com/llvm/llvm-project.git
用戶也可以直接下載源碼包：https://github.com/llvm/llvm-project/releases。
假定 LLVM 的源碼目錄為 $LLVM_SRC。
編譯 MLIR
首先用戶需要指定一個路徑用於存放編譯中間產物，假定其路徑為 $LLVM_BUILD。然后使用下列命令對 LLVM 進行配置：
cmake -S "$LLVM_SRC" -B "$LLVM_BUILD" -DLLVM_ENABLE_PROJECTS=mlir -DCMAKE_BUILD_TYPE=Release
默認情況下，LLVM 禁用了異常處理和運行時類型信息。如果應用程序需要依賴這些功能，可指定在配置時指定 LLVM_ENABLE_EH 和 LLVM_ENABLE_RTTI CMake 變量的值為 ON：
cmake -S "$LLVM_SRC" -B "$LLVM_BUILD" -DLLVM_ENABLE_PROJECTS=mlir -DLLVM_ENABLE_EH=ON -DLLVM_ENABLE_RTTI=ON -DCMAKE_BUILD_TYPE=Release
更多的 LLVM 配置參數參見 https://llvm.org/docs/CMake.html。
執行完配置過程后使用下列命令執行編譯：
cmake --build "$LLVM_BUILD"
安裝 MLIR
使用如下命令將 LLVM 安裝到 /usr/local 目錄：
cmake --install "$LLVM_BUILD"
如果想指定另外一個安裝目錄，例如 $INSTALL_DIR，可以使用 --prefix 命令行參數來指定：
cmake --install "$LLVM_BUILD" --prefix "$INSTALL_DIR"
在 CMake 項目中使用 MLIR
用戶可以在 CMake 項目文件中使用下列語句添加查找 MLIR 依賴：
find_package(MLIR REQUIRED CONFIG)
如果 MLIR 被安裝到了系統目錄（比如 /、/usr、/usr/local 等），CMake 無需額外的配置就能找到 MLIR；如果 MLIR 被安裝到了非系統目錄，可以在 CMake 的配置過程通過 CMake 的 MLIR_DIR 變量來指定 MLIR 的安裝位置：
cmake "$MY_PROJECT_DIR" -DMLIR_DIR="$INSTALL_DIR"
成功之后用戶可以直接使用 MLIR 的庫作為編譯目標的依賴：
add_executable(my-executable main.cpp)
target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})
target_link_libraries(my-executable PRIVATE MLIRIR)
其中 MLIR_INCLUDE_DIRS 是自動生成的變量，其指向 MLIR 的包含目錄。
在使用 CMake 定義可執行文件目標時，如果 LLVM 禁用了運行時類型信息，那么依賴於 LLVM 的可執行文件目標也需要禁用運行時類型信息，否則可能會編譯失敗。LLVM 提供了一個 CMake 幫助函數 llvm_update_compile_flags 可以自動完成這個配置。這個函數定義在 LLVM 提供的 AddLLVM.cmake 文件中。用戶可以使用下列語句導入 AddLLVM.cmake 文件：
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
include(AddLLVM)
導入 AddLLVM.cmake 文件后就可以對編譯目標進行配置了：
llvm_update_compile_flags(my-executable)
完整的 CMake 項目文件示例如下：
cmake_minimum_required(VERSION 3.15)
project(my-executable)
find_package(MLIR REQUIRED CONFIG)
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
include(AddLLVM)
add_executable(my-executable main.cpp)
target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})
target_link_libraries(my-executable PRIVATE MLIRIR)
llvm_update_compile_flags(my-executable)
結論
MLIR 是一種新型的編譯器框架，設計從已有的編譯器的實現中吸取了經驗和教訓，包括了中間表示的定義、轉換以及優化等功能，極大地方便了新的編譯器的開發和調試工作。同時，MLIR 也包含了很多現成的工具可直接使用（batteries included）。MLIR 包攬了編譯器設計中的通用部分，使得編譯器的開發人員，可以專注於核心的語義分析、中間表示的設計和變換，以此降低開發成本，提高開發效率和提高成品質量。
參考鏈接
MLIR 主頁：https://mlir.llvm.org/。
MLIR 語言參考：https://mlir.llvm.org/docs/LangRef/。