GEMM(General Matrix Multiplication)即通用矩陣乘法運算，由於其計算行為具有一定的複雜性以及規律性，是編譯演算法研究的絕佳場景。MLIR是近期非常熱門的一個編譯器軟體框架，是工業界及科研界研究的一個熱點，其提供了一套靈活的軟體基礎設施，對中間表示式（IR）及其相互之間的轉換進行規範的管理，是一個非常友好的編譯器開發平臺[1][2]。本文即是分析在MLIR框架下，實現GEMM最佳化的內容，以及對MLIR在這一方面的實現優勢的討論。
矩陣乘法運算，由於其過程會包含大量的乘加操作，並且伴隨大量的資料讀寫，因而如何充分利用好底層硬體的儲存資源以及計算資源，是編譯器對其效能最佳化的關鍵。目前，已有的一些最佳化策略主要包括：
當前的處理器效能主要受限於記憶體牆，即計算速度要大於資料儲存的速度。為了打破記憶體牆的約束，各類硬體包括CPU及其他專用處理器，會設定不同層次的儲存單元，而這些不同層級的儲存單元往往大小以及讀寫速度不同，一般越靠近計算單元的儲存其儲存容量也越小但訪問的速度也越快。如果可以將計算過程的資料區域性化分塊，而這些分塊的資料可以獨立完成計算，那麼分塊的資料就可以放在層次化的儲存中，然後透過不同儲存間建立Ping-Pong的資料傳輸方式，將資料儲存與計算解耦，從而可以有效得隱藏儲存牆的問題，提高計算效率。矩陣運算就有這種特點，因而可以透過矩陣分塊來加速運算，如下圖1所示，假設有兩層儲存，將輸入矩陣A和B，以及輸出矩陣C，根據儲存大小劃分成相應的小塊，即m->mc，n->nc，k->kc，每次將Ac(mc, kc), Bc(kc,nc), Cc(mc, nc)送入到離計算單元更近的儲存模組內，完成區域性的計算後再進行下一次的計算。
在不同的底層硬體中，由於儲存的層次以及不同層次的儲存的容量大小不一樣，分塊的大小也會不一樣。比如，文章[3]中對CPU而言，（Ac, Bc, Cc）劃塊的大小與cache大小一致，而為了充分利用register的資源，還會對（Ac, Bc, Cc）再進一步細劃塊成（Ar, Br, Cr），其尺寸大小與暫存器的數量一致。
向量化的操作，主要是利用硬體的向量化指令或者SIMD（單指令多數）指令的特性，實現一個指令週期對多個值操作的能力。如下圖2所示，透過將4個數據組成向量，利用處理器可以處理4個元素的新向量的計算能力，可以將4個指令週期的處理時間，壓縮成1個指令週期的處理時間，從而極大提高運算處理能力。
由於矩陣乘法有多層迴圈構成，如果底層硬體有一定的並行化能力，包括多執行緒多程序處理能力等，那麼可以對迴圈進行適當展開，從而提高計算的並行度，增加併發執行的機會。如下圖3所示，將一個次數為1024的迴圈，展開成256次迴圈，新的迴圈內又包含4條可以並行執行的展開計算，如果處理器能夠並行處理迴圈內部的展開計算，那麼透過對原來的迴圈展開，可以獲得接近4倍的效能提升。
矩陣乘法的運算也包括其他的最佳化策略，比如資料重排等，但總體而言，各類編譯器都是利用這些策略，充分利用硬體的儲存及計算資源，達到最佳的運算效率。一些傳統的計算庫，如：OpenBLAS, BLIS, MKL等，開發時間長，效能也有比較優秀的表現。
MLIR基於多層中間表示的方言（Dialect）思想，提供了一整套完善的編譯器基礎框架，可以幫助開發者快速實現編譯策略想法的編譯器。本文主要參考論文[4]，分析GEMM運算在MLIR中的實現，對應的硬體Target是因特爾i7-8700K處理器，每個核包含有32/256KB L1/L2 Cache以及多核共享的12MB L3 Cache，處理器支援AVX-2指令（256bit),最佳化目標是一個2088x2048xf64與2048x2048xf64的矩陣乘。
首先，其在高層次的Dialect上定義了一個矩陣運算的運算元，這個運算元的引數包含了輸入矩陣（A，B）以及輸出矩陣（C），同時為這個運算元添加了tile/unroll 的尺寸等屬性。如下圖4所示，其中（M_C, K_C, M_R, N_R）屬於Tile尺寸，K_U屬於Unroll的大小。這裡面（M_C, K_C）的選擇是使得M_CxK_C大小的A矩陣塊能夠在L2 cache中複用，（K_C, N_R）的選擇是使得K_CxN_R大小的B矩陣塊能夠在L1 cache中複用，（M_R, N_R）的選擇是使得M_RxN_R大小的輸出矩陣塊能夠在CPU Register中複用，這些值是根據硬體計算或者tunning出來的，在這裡面的測試取了一個經驗值。這些屬性可以協助轉換到更低一層的運算元的策略實現，而選擇哪些屬性，則是跟編譯演算法以及編譯的底層硬體物件有關，這些屬性也是協助轉換成下一層跟硬體更貼近的中間表示的實現，因而可以根據實際需要，靈活使用。
其次，MLIR的特點就是透過統一的多層中間表示，來實現對運算元的層層低層化（lower）到具體的硬體目標上。針對上述高層次上定義的矩陣乘法運算元，透過利用其所攜帶的最佳化屬性，以及底層硬體的特點，設計了多條轉換的路徑（Pass），從而進一步把該運算元lower到MLIR框架提供的中間輔助層（此中選擇了Affine, Linalg，和Standard Dialect）。在這一層的轉換過程中，基本包含了所有的策略，如：Tile，定製化複製，unroll，vectorize等。然後再將中間的輔組層的Dialect，進一步lower到LLVM Dialect上，如下圖5所示。
最後，透過mlir提供的mlir-cpu-runner工具，可以執行最後生成的LLVM Dialect的結果。總體最佳化及執行測試的命令，如下圖6所示。其中，“-hopt”，“-hopt-vect”等，是從高層的運算元（hop.matmul）到中間輔組層的轉換路徑，每一條路徑都包含有相應的編譯策略，可以根據需要，靈活新增以及改變，“-convert-linalg-to-loops”, “-lower-affine”等時中間輔助層之間的轉換，最後轉換成LLVM Dialect。
總體上，一個GEMM運算透過在MLIR框架下實現，並重寫最佳化策略的路徑，可以得到如圖7所示的結果，其中箭頭1對應包含了所有重寫最佳化策略的MLIR實現，可以看到其能達到的計算速率為61.94GFLOPS，離理論上的峰值計算速率（75.2GFLOPS）比較接近，跟傳統的計算庫相比（如：BLIS，OpenBLAS，MKL等），也有著可以媲美的結果，其中的差距可能是傳統的計算庫有tunning的機制以及在編譯器後端生成彙編指令及機器碼有更成熟且高效的最佳化，因而可以得到更好的最佳化結果。總體而言，用MLIR重寫的GEMM最佳化演算法有著非常良好的表現。
另一方面，MLIR框架提供了非常完善的C++以及Python介面，因而可以很方便接入已有的計算庫，進行聯合最佳化。在[4]文中嘗試了用MLIR+BLIS的方法，將MLIR放在外側（提供手動最佳化功能），BLIS則作為micro-kernel放在內側（提供auto tunning功能），最終的結果如圖7中箭頭2所示。可以看出，對於DGEMM（雙精度），透過MLIR與BLIS的聯合最佳化，也可以達到接近峰值的效能，而其效能要比單獨的MLIR或者BLIS最佳化要差一點。但其實在SGEMM（單精度）的測試中，MLIR+BLIS的最佳化又要比單獨的MLIR或者BLIS最佳化要好一些，因而其中的效能在差異還需要進一步分析。總體而言，MLIR提供了非常完善的支援，可以融合已有的計算庫的最佳化演算法，去實現聯合的編譯最佳化。
透過上面對MLIR實現GEMM最佳化演算法的編譯的介紹，可以看出MLIR在其中有著非常突出的優勢。
首先，MLIR框架提供了非常完善的編譯器基礎設施，可以讓開發者不需要花費太多精力在編譯器周邊的實現，從而更加專注於編譯演算法的開發。同時，其基於多層中間表達的方式，可以讓編譯器更加模組化，可以讓編譯演算法利用不同層次的中間表達的抽象資訊，在不同的層次中逐步具體化，從而使得演算法實現更加層次化，更加易於實現及管理。
其次，MLIR框架提供了一直到最底層硬體的表示支援，由於其可以層次化在不同的中間表示層實現編譯演算法，可以在高層次的中間表示中實現不依賴於底層硬體的通用演算法，而在接近硬體底層中，開發針對性的路徑實現相應的編譯演算法，因而可以很方便地針對不同硬體目標開發統一的編譯環境。本人認為，這也是MLIR相對於一些現有的AI編譯器，如：TVM等，最有優勢的地方之一，由於其框架可以根據需要自行擴充套件Dialect，同時這些Dialect又在系統中遵循一套統一的正規化進行管理，因而對不同的編譯目標（硬體target）會有很強的擴充套件性，同時編譯器的工程管理又可以做到非常好的統一性。
另外，MLIR框架提供了完善的C++/Python介面，可以很方便地接入已有的最佳化演算法，快速實現演算法遷移。
本文主要介紹了矩陣乘法運算在MLIR編譯器框架實現的主要過程及內容，以及其在最佳化演算法的多層次實現，以及接入已有最佳化演算法的能力等方面的優勢。MLIR編譯框架，為編譯器的開發者，在多層中間表達轉換以及演算法最佳化等方面提供強大的基礎設施支援，降低開發編譯器的門檻。本人也希望透過本文，可以讓讀者對MLIR的工程實現過程更加清晰一些，從中得到一點啟發。
由於水平有限，文中存在不足的地方請各位讀者批評指正，也歡迎大家一起參與我們的討論。
[1] Chris Lattner, Mehdi Amini,Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle,Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infrastructure for the end of moore's law, 2020
[2] MLIR：https://mlir.llvm.org/
[3] Tze Meng Low, etc. Analytical Modeling Is Enough for High-Performance BLIS. 2016.
[4] UdayBondhugula, High Performance Code Generation in MLIR: An Early Case Study With GEMM. 2020.
壁仞科技研究院作為壁仞科技的前沿研究部門，旨在研究新型智慧計算系統的關鍵技術，重點關注新型架構，先進編譯技術和設計方法學，並將逐漸拓展研究方向，探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則，將積極投入各類產學研合作並參與開源社群的建設，為相關領域的技術進步做出自己的貢獻。
關鍵詞
框架
硬體
效能
編譯器
運算元
