基於IREEHLO專案看MLIR編譯器實現的過程及優勢

摘要
MLIR(Multi-Level Intermediate Representation)[1]是谷歌團隊開發的一個非常熱門的開源編譯器框架,其提供了一套靈活的軟體基礎設施,來規範中間表示式(IR)及其相互之間的轉換,建立了一個非常友好的編譯器開發平臺,一些比較好的對MLIR框架解讀可以參考[2][3]。IREE專案也是谷歌推出的,基於MLIR進行統一的IR規範管理的一個從tensorflow模型到不同硬體加速器平臺的端對端專案,其前端的核心內容就是對XLA對應的HLO運算元進行MLIR實現。MLIR在編譯器實現方面,有一些非常優秀的特點,本文即是基於HLO專案對此做一個拋磚引玉的簡單討論。
IREE HLO專案介紹
MLIR HLO專案是IREE的重要的組成部分,IREE專案的IR結構如圖1所示。MLIR HLO的主體其實就是XLA(Accelerated Linear Algebra)編譯器高效能運算元組成的IR表示[4],其Operation基本上是和XLA中相對應的HLO是一一對應關係。MLIR HLO部分的輸入是由上端translator工具將TFgraph parse成TF IR並轉換成相應的HLO IR,並在其內進行一定的處理,包括Dialect之間的轉換,以及主要繼承於XLA相應的演算法處理,最後輸出或者lowering到後端的硬體target相對應的IR。
圖1 IREE專案MLIRIR示意圖
MLIR HLO專案的實現內容包含了許多工作,其主要包括:1.不同Dialect的定義,及其運算元(Operation)內容的定義;2.不同Dialect之間運算元轉換;3.Dialect內或之間的運算元鏈的演算法處理(比如:多面體處理,fusion操作)。本文接下來內容,即從這三方面進行分析討論。
定義Dialect運算元內容
一般編譯器IR定義的運算元,其內容主要包括:1.輸入變數;2.輸出的資料型別及形狀;3.其他附屬資訊(用來表徵一些最佳化資訊或者高層次資訊)。而如果不同層次的IR去進行獨立的開發,可能會使得IR的API混亂,並且部分的軟體基礎設施也會重複開發。而MLIR的Dialect可以很好解決這些問題,每一層的IR可以包含一種或者多種Dialect表示,而每一個Dialect包含各種自定義的運算元(Operation)。MLIR的Dialect有比較統一的寫法格式,並且由其提供的table_gen工具,最終可以為Dialect的各個運算元生成統一的C++介面。
圖2 HLO Dialect運算元的定義
一個典型Dialect運算元的table_gen的定義方式,如圖2所示,主要包含兩部分:1.特殊屬性部分(traits,圖2黃色部分);2.以及主體部分(圖2藍色部分)。特殊屬性部分定義了運算元的特殊操作類,主要以interface方式呈現,在特定運算元組合的最佳化方面有著重要的作用。主體部分即是MLIR定義的統一屬性部分,其包括:arguments,results等內容。其中arguments的內容,又可以根據需要,透過在build函數里面的addOperands及addAttribute函式(如圖3的虛擬碼),被分配到Operands(相當於op的input)或者Attributes屬性裡面。Operands,Attributes,Results這三個屬性,在MLIR軟體框架裡,又為其提供了不同的訪問函式(參見後面運算元鏈的演算法最佳化實現關於基於operand的get_using函式介紹),可以極大提高對運算元轉換及最佳化實現的效率。
圖3 build函式的虛擬碼
總體而言,MLIR在這一方面,為運算元的訪問,表示,及驗證等方面提供了非常完整的基礎設施,使得開發人員可以專注於運算元的內容定義,而不用花很多的精力去關心運算元的周邊介面表示及其他內容的實現。
實現Dialect間運算元轉換
Dialect之間的lowering過程,即是對其相應的運算元進行轉換過程。MLIR框架中的運算元,包含了豐富的介面函式,可以利用這些介面函式,可以快速得到需要轉換的運算元的資訊,包括:1. 函式getOperand得到operand內容;2.函式getAttribute得到所有的Attribute內容;3.函式getResult得到result內容;以及其他函式得到相對應其他的內容。再根據這些內容,計算出新的一個或者多個operations所需要的operands/attributes及其他的內容資訊,並組裝成新的operations。具體流程圖如圖4所示,包括:
圖4 不同Dialect的operation進行conversion的過程
1、對於舊的operation,透過相應的get函式得到attributes, results, operands等資訊;
2、步驟1中得到的資訊透過一定的轉換,得到新的attributes, operands等資訊;
3、建立新的operations,並將轉換後的資訊填充到這些operations內,同時移除舊的operation。
總體而言,MLIR在這一方面,也是提供了比較完備的API訪問介面,以及運算元建立及移除機制,可以讓開發人員更加專注於轉換的演算法實現,而省去轉換過程中周邊的一些適配工作。
實現運算元鏈的演算法最佳化
在對運算元從高層次到低層次convert過程中,對一些運算元根據計算特點進行演算法最佳化以達到最好的效能效果,也是編譯器重要的工作之一。MLIR HLO專案中,其演算法最佳化主要繼承於XLA HLO專案中已開發的一些最佳化內容。其中非常重要的一個內容是在HLO方言層面,對運算元進行fusion處理,即kLoop,KInput演算法處理。圖5即是一個kLoop的一個例子,其主要的思想是將所有由相同輸出形狀的連線運算元組成的運算元鏈(目前只支援element-wise性質的運算元),fuse到一起進行運算。
圖5 MLIR HLO fusion例子,上圖是輸入的mlir,下圖是fusion的結果
對於該例子,其整個演算法的流程示意圖,如圖6所示,主要的內容包括:
圖6 MLIR HLO fusion過程的示意圖(對應圖5中的例子)
1、基於MLIR形式的op_list資訊,去建立一個有向的連線圖;
2、透過深度優先演算法(DFS,Depth First Search),對有向連線圖進行查詢,透過一定的判斷條件,得到並融合可以fusion的pattern;
3、不斷迭代上一步的步驟,直到遍歷所有的節點,最後得到最終的fusion結果(圖5中的例子,最終是將整個的operation融合成一個大的fusion operation);
4、將fusion結果的有向連線圖和未參與融合的其他部分內容連線在一起,並輸出新的MLIR形式的op_list列表。
首先,在建立有向連線圖過程中,其主要內容是獲取每個節點的輸入及輸出節點(即:in和out),MLIR對此提供了兩個函式:
  • operand.getDefiningOp():可以獲取operand所對應的輸入operation節點;
  • result.getUses():可以獲取result所對應的輸出operation節點。
基於上述兩個功能函式,可以非常方便得獲取有向連線圖節點的輸入及輸出(in和out),建立運算元的有向連線圖。
其次,MLIR的operation定義有特殊屬性部分內容(即traits,主要以interface形式呈現),這些內容可以在某些特定的實現方面提供支援。HLO專案定義了InferFusibilityOpInterface這樣的一個特殊屬性介面,該屬性介面包含了多個功能函式,都是用來判斷兩個連線的operations是否可以fuse,而需要參與fusion演算法的operations則需要新增該特殊屬性。HLO fusion演算法判斷的過程如圖7所示。判斷的條件包括:operation本身是否是可以融合的,operation與其輸入節點(withOperand)或者輸出節點(withConsumer)是否可以融合,兩個operations的輸出形狀是否相等。這樣,透過這些特殊屬性,可以很方便建立operation的判斷過程。
圖7 fusion的判斷過程示意圖
最後,將融合完成後的有向圖,恢復到MLIR的op_list形式,MLIR也提供了相應的一些功能函式,極大地方便了這一過程的實現。
總體而言,MLIR在對新演算法的開發過程,其提供的基礎設施框架可以減少周邊適配的工作,讓開發者更多專注於演算法本身的實現,而對已有演算法的移植也提供了非常友好的相容性。
結論與思考
本文透過結合MLIR HLO專案,分析了基於MLIR框架編譯器的一些關鍵點的程式碼實現以及優勢。總體而言,基於MLIR的框架,為編譯器多IR的實現提供了豐富且統一的軟體API介面,以及其他非常高效實用的基礎設施軟體框架支援,可以讓常規的編譯器開發更為便捷和高效。當然,也有一些觀點認為MLIR並沒有解決一些編譯器的根本問題,比如auto-tiling及auto-schedule等,而且過多的Dialect表示層級劃分反而可能加劇IR層級的碎片化問題,同時高層次的IR表示存在運算元表述的不完備性等問題。但本人認為,這些也都是傳統編譯器的老大難問題,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、Chris Lattner, "Thoughts on Tensor Code Generation in MLIR", 2020-01-24, Talks and Related Publications – MLIR (llvm.org)
4、https://www.tensorflow.org/mlir/xla_gpu_codegen
關於壁仞科技研究院
壁仞科技研究院作為壁仞科技的前沿研究部門,旨在研究新型智慧計算系統的關鍵技術,重點關注新型架構,先進編譯技術和設計方法學,並將逐漸拓展研究方向,探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則,將積極投入各類產學研合作並參與開源社群的建設,為相關領域的技術進步做出自己的貢獻。
掃碼關注我們

相關文章