梁佳利，華保健，呂雅帥，蘇振宇

1.中國科學技術大學軟件學院，合肥230000

2.中科寒武紀科技股份有限公司，北京100000

TVM（tensor virtual machine）[1]是一個開源深度學習編譯器[2-5]，通過提供圖級別的優(yōu)化[6]和算子級別的優(yōu)化，生成多硬件后端代碼。在算子實現(xiàn)時，不同的計算方法卻導致了生成代碼的局部性、并行性的差異[7-8]。TVM 的領域專用語言——張量表達式支持用簡潔的方式定義算子，并提供一系列調度原語對算子代碼做循環(huán)變換。用張量表達式描述的算子會被TVM翻譯為TVM IR（TVM intermediate representation）——一種樹狀的、易于描述循環(huán)計算的高級中間表示。TVM 在TVM IR 上實現(xiàn)了一系列分析和變換的優(yōu)化遍，來對算子的中間表示進行調整和優(yōu)化，最終生成目標硬件平臺上的高效代碼。

深度學習應用[9-12]中的重要算子，如卷積、矩陣乘和向量加等，都會被TVM 編譯生成TVM IR 上的多層嵌套循環(huán)[13]，經過循環(huán)變換后會產生與循環(huán)迭代變量相關的復雜表達式計算。例如，以下的代碼示例是向量加算子的TVM IR 表示：

1.for(i.outer:int32,0,4){

2.for(i.inner:int32,0,32){

3.C[(i.outer*32)+i.inner]=

4.A[(i.outer*32)+i.inner]+

5.B[(i.outer*32)+i.inner]}}

注意到在內層循環(huán)中存在被重復計算的表達式i.outer * 32，它應該被循環(huán)不變量外提優(yōu)化[14]提取到第一層和第二層循環(huán)之間，從而消除冗余計算。

然而，在深度學習編譯器中實現(xiàn)循環(huán)不變式外提優(yōu)化，存在以下研究挑戰(zhàn)：

首先，傳統(tǒng)編譯器中的循環(huán)不變量外提優(yōu)化算法，將計算結果保存在物理寄存器中，避免每次使用時都重復計算。但是TVM 等深度學習編譯器的高級中間表示無法感知物理寄存器的存在，盲目地外提所有循環(huán)不變量會帶來不必要的寄存器壓力，甚至降低算子的性能[15]；因此直接將傳統(tǒng)編譯器中的循環(huán)不變式外提算法應用到TVM 等深度學習編譯器中，并不能得到理想的優(yōu)化效果。

其次，表達式中操作數的順序會影響不變式的發(fā)現(xiàn)[16]；傳統(tǒng)的循環(huán)不變式外提算法，對深度學習算子中經常出現(xiàn)的、復雜嵌套條件表達式的檢測能力有限，無法檢測、合并、變換生成新的循環(huán)不變式。

最后，深度學習編譯器生成的算子，會被目標平臺編譯器進一步編譯優(yōu)化。目標平臺編譯器在一個更低級的中間表示上進行優(yōu)化[17]，而在兩個不同級別上的優(yōu)化常常存在沖突[18]，導致降低了循環(huán)不變式外提算法的有效性。Halide 開源編譯器[19]中實現(xiàn)了一個保守的循環(huán)不變式外提算法，但該算法只考慮循環(huán)內的表達式，而不對語句做處理；且在做循環(huán)不變式外提之前只對包含加減法的表達式做了重結合，沒有對條件表達式做優(yōu)化；因此，這些簡單的處理并不能完全解決上述研究挑戰(zhàn)。

本文提出了一種面向深度學習編譯器高級中間表示的、采用啟發(fā)式策略的循環(huán)不變式外提算法。首先，針對第一個研究挑戰(zhàn)，本文借鑒Halide 編譯器循環(huán)不變量外提優(yōu)化的代價模型，設計了衡量表達式計算代價的函數，來指導是否外提循環(huán)不變式，以一種簡單且易于實現(xiàn)的方式對消除冗余計算和寄存器分配之間的矛盾做了權衡。其次，針對第二個研究挑戰(zhàn)，本文算法在檢測循環(huán)不變量之前，先調整操作數結合順序和變換條件表達式的形式來對表達式做規(guī)范化，為循環(huán)不變量外提優(yōu)化以及后續(xù)的其他優(yōu)化提供更好的時機。最后，針對第三個研究挑戰(zhàn)，本文結合TVM IR 和目標平臺編譯器的具體特點，對分支和迭代數目比較小的內層循環(huán)做了特殊處理，解決了在不同級別的中間表示上的優(yōu)化沖突。

為驗證本文算法的有效性，在開源編譯器TVM的0.7 版本上實現(xiàn)了本文算法，并進行了實驗和結果分析。實驗選擇TVM 的TOPI測試集作為測試集合，對其包含的27 個算子對應的511 個測例進行了測試，這些算子都是TVM 框架頻繁用到的算子，其中包含了多種實現(xiàn)策略的卷積和激活等算子，測例都是典型的多層嵌套循環(huán)的數組計算，同時一些算子內的補零和邊界檢查操作在循環(huán)內產生了條件表達式和分支。這些算子測例描繪了常見深度學習領域算子特點，并且每一個算子都提供了大量在常見神經網絡中不同參數規(guī)模的測例，因此具有一定的代表性。實驗結果發(fā)現(xiàn)，在排除那些優(yōu)化前后代碼沒有發(fā)生變化的測例后，與未進行循環(huán)不變量外提優(yōu)化的測例相比，有47.6%的算子性能得到了顯著提升，最大加速比大于40.0%；所有測例的總時間加速比平均為22.7%。實驗結果證明，本算法對TVM 生成算子的性能產生了積極影響，加速了含有大量冗余計算的算子的執(zhí)行，完善了TVM 的算子優(yōu)化。

總結起來，本文的主要貢獻如下：

（1）提出了一個新的融合深度學習代價函數和啟發(fā)式策略的循環(huán)不變量外提算法；

（2）基于TVM 開源深度學習編譯器（0.7 版本），給出了原型系統(tǒng)設計與實現(xiàn)；

（3）選取深度學習領域典型的27 個算子及其511個測例，對原型系統(tǒng)進行了系統(tǒng)測試；實驗結果分析表明，該算法對提升TVM 的性能產生了積極作用。

1 背景與研究動機

本章介紹關于TVM 中間表示的相關背景，以給出本文工作的研究動機。

1.1 張量表達式和TVM IR

張量表達式是TVM 基于Halide 的算子描述和調度優(yōu)化分離的思想[7-8]，而設計實現(xiàn)的描述張量運算的領域專用語言。該語言定義張量間的輸入輸出關系，同時提供循環(huán)拆分、循環(huán)合并和循環(huán)分塊等一系列調度原語，這些調度原語的組合指定了計算輸出張量的具體策略[20]。

張量表達式首先被翻譯為TVM IR，TVM IR 是一種樹狀的中間表示，主要包含兩種類型的節(jié)點：（1）表達式節(jié)點；（2）語句節(jié)點。表達式節(jié)點主要包含：有名字的變量和常數、數組數據的加載、函數調用、算術邏輯表達式、條件表達式等；語句節(jié)點包含順序、循環(huán)、分支、賦值等。

下面實例定義了一個描述固定規(guī)模矩陣乘法算子的張量表達式：

變量A、B為規(guī)模為（1 024，1 024）的二維矩陣，k為求和歸約的那個循環(huán)“軸”。矩陣C由矩陣A和B計算得到。上述算子被翻譯為以下TVM IR：

為了高效實現(xiàn)上述代碼中的最內側循環(huán)，TVM使用了tile、split和reorder等操作，對矩陣乘法算子進行循環(huán)分塊、循環(huán)拆分和循環(huán)重排序：

1.2 研究動機

算子經過循環(huán)變換后有更好的并行性和數據局部性，且能夠充分利用底層硬件特性，加速計算[7,13,20-22]。但是算子經過循環(huán)變換后，數組下標表達式變?yōu)榱苏Z義等價，但更復雜的表達式，增加了額外標量運算。例如上述矩陣乘算子最后生成的TVM IR，對于最內層循環(huán)，其循環(huán)迭代變量是yi，下標表達式xo*32768+xi*1024+yo*32 在最內層循環(huán)迭代時，計算結果不變，在第一次迭代后的計算都屬于冗余計算。本算法的設計目標，是通過把循環(huán)中多次迭代結果保持不變的計算移動到循環(huán)外，從而減少冗余計算。另外，TVM 編譯生成的算子會經過目標平臺編譯器的編譯優(yōu)化，才能在常用深度學習硬件，如英偉達的GPU 和寒武紀的MLU 硬件上[23]，高效運行。這類硬件的典型特征是很多部件都為了運行計算密集的程序而設計。為了充分利用這些硬件特性，目標平臺編譯器會選擇展開內層循環(huán)，一些計算的延時通常被指令流水線所掩蓋，因此一些冗余計算并不影響運行時間，反而會降低物理寄存器壓力。而由于硬件更專注于計算，對于分支優(yōu)化處理的任務需要由編譯器優(yōu)化完成[24-25]。因此對布爾表達式、條件分支語句進行外提時需要謹慎處理，有時還需要變換條件表達式的形式，使得生成的代碼更利于底層編譯器的優(yōu)化。綜合以上討論，本文面向深度學習編譯器，提出了一種新的循環(huán)不變量外提算法，該算法考慮了TVM IR 和目標平臺編譯器的特點，通過外提循環(huán)不變量減少冗余計算，提高TVM 生成算子的性能。

2 相關工作

作為優(yōu)化循環(huán)內冗余計算的一種有效技術，循環(huán)不變式外提已經被廣泛研究和應用，很多編譯器都使用消除冗余計算的方法，優(yōu)化生成代碼性能。Aho 等[14]系統(tǒng)討論了循環(huán)不變量外提的概念。TVM[1]在圖級別的中間表示RELAY 上[6]，使用公共子表達式消除算法，在算子的層級上減少了相同輸入規(guī)模算子的重復計算；LLVM（low level virtual machine）[26]在LLVM IR 上使用了循環(huán)不變量外提技術，在基于基本塊的控制流圖上對循環(huán)內不變量指令外提；Steffen 等[27]提出了一種代碼移動和代數優(yōu)化的算法；Rosen 等[28]使用全局值編號來做冗余消除優(yōu)化。但是，這些研究都是做語義更高的計算圖級別中間表示，或是語義更低的底層中間表示上進行，缺失了在像TVM IR 這種源代碼級別的樹狀中間表示上的優(yōu)化算法。如RELAY 中間表示上沒有循環(huán)的概念，對于RELAY 的冗余消除優(yōu)化都是以整個算子為單位，優(yōu)化的粒度過大，并不能很好地對算子內的計算進行優(yōu)化。如在LLVM IR 上的循環(huán)不變量外提技術在一個更低級別的中間表示上進行，并沒有對于高語義的表達式形式如嵌套的條件表達式提供優(yōu)化方案。因此這些方法無法直接應用在深度學習編譯器的優(yōu)化場景中。不少的研究對于源代碼級別的優(yōu)化做了討論和實踐。Halide 編譯器[19]在把算子計算描述翻譯為Halide IR 后，針對循環(huán)內的表達式使用了公共子表達式刪除和循環(huán)不變量外提技術。Rawat等[15]提出了一個寄存器優(yōu)化框架，在源代碼級別對含有多層嵌套循環(huán)的高階模板計算進行了寄存器的優(yōu)化，提高指令并行性；akg 框架[29]研究了在華為mindspore后端的算子自動生成器akg 上，實現(xiàn)公共子表示消除算法，優(yōu)化了特定指令的重復計算；Vasilev等[30]提出了一種在遞歸函數上的循環(huán)不變量外提算法；Song 等[31]通過展開循環(huán)的前幾次迭代，來優(yōu)化循環(huán)不變量外提的效果；Bacon 等[32]全面總結了在高級語言上的高性能計算的優(yōu)化方法。但是，這些研究提供的消除冗余計算的方法沒有考慮到TVM 算子程序特點和深度學習加速硬件的特性。如akg 的公共子表達式消除只是針對特殊的張量指令進行，并沒有對于TVM 算子調度優(yōu)化產生的復雜下標表達式和條件表達式做優(yōu)化。深度學習領域的處理器雖然有很好的指令集并行架構，卻不擅長分支的處理。如Halide 編譯器實現(xiàn)的循環(huán)不變式外提算法，在分支的處理上忽略了對條件表達式的優(yōu)化。因此這些方法對于深度學習編譯器的優(yōu)化，還存在不足之處。

3 系統(tǒng)設計與實現(xiàn)

本章討論系統(tǒng)的設計與實現(xiàn)。在給出系統(tǒng)架構后，重點討論循環(huán)不變式外提算法，并給出規(guī)范化和代價函數。

3.1 系統(tǒng)架構

本文設計和實現(xiàn)的系統(tǒng)基于開源深度學習編譯器TVM 的0.7 版本。TVM 優(yōu)化[1]可以分為圖級別優(yōu)化和算子級別優(yōu)化兩個獨立的模塊，算子級別的優(yōu)化又可分為在張量表達式級別的基于機器學習的自動調度優(yōu)化和在TVM IR 級別的優(yōu)化。本算法針對算子生成的TVM IR 進行優(yōu)化，涉及對算子級別優(yōu)化模塊的修改。循環(huán)不變量外提算法在TVM 中的整體架構圖和關鍵編譯流程如圖1 所示。

圖1 TVM 中實現(xiàn)循環(huán)不變量外提算法的架構圖Fig.1 Architecture of loop invariant code motion algorithm in TVM

從整體架構上看，循環(huán)不變量外提算法由兩個關鍵部分組成：第一個是前處理部分規(guī)范化，該部分對表達式進行規(guī)范化處理，使得變化后的表達式更利于循環(huán)不變量的識別；第二個是不變量外提部分，對循環(huán)保持不變的表達式和語句進行識別和外提。TVM 通過優(yōu)化遍管理器組織一系列優(yōu)化遍在TVM IR 上進行優(yōu)化，向TVM 的優(yōu)化遍管理器注冊本算法的優(yōu)化遍，并安排新增的優(yōu)化遍到原生TVM 優(yōu)化遍序列的合適的位置（如圖1 中加粗邊框的遍所示）。同時使用了原生TVM 的循環(huán)判斷外提的優(yōu)化遍，對分支條件為循環(huán)不變量的分支語句做提升；修改原生TVM 的化簡優(yōu)化遍，禁止了原本對整數類型的賦值變量進行前向替換的操作。

大量研究表明，父母的信念是兒童信念發(fā)展的錨定起點(Ozorak, 1989; Boyatzis et al., 2006)。不同信仰的父母有著不同的觀念和行為方式，這些不同借助親子談話等方式，影響著兒童各種觀念的形成(Boyatzis et al., 2006)。 例如， Rosengren(2004)等人發(fā)現(xiàn)， 父母(天主教徒)被問及如何回答3～6歲的兒童的死亡問題時，大部分父母是用宗教相關的字眼來回應的，例如，天堂、靈魂、上帝等。有宗教信仰的父母的信念更加與眾不同，對兒童信念的影響也顯得尤為突出。

3.2 循環(huán)不變代碼外提算法

深度學習編譯器中的循環(huán)不變量實例中，典型的是訪問數組元素的地址表達式的計算和條件表達式的計算，表達式中包含了循環(huán)迭代變量、程序參數和立即數。TVM IR 中的循環(huán)節(jié)點包含了循環(huán)嵌套關系以及對應的迭代變量，因此遞歸地判斷某個表達式是否是循環(huán)不變的，只需要構成表達式的值滿足以下判定條件：

（1）是常數（不可變的程序參數或立即數）；

（2）該變量的賦值在循環(huán)之外，或者；

（3）到達該變量的使用只有一個定值，且該定值的表達式也是循環(huán)不變的。

實際上，識別在一個循環(huán)中的不變表達式，只需要標記出，在循環(huán)的多次迭代中計算結果可能發(fā)生變化的表達式，那么剩下未標記的表達式都是該循環(huán)的不變式。因此，可按照以下計算步驟識別并提升循環(huán)不變式：

（1）標記循環(huán)迭代變量為“不可提升”；

（2）標記循環(huán)迭代變量到達的定值變量都是“不可提升”；

（4）將（3）中新的變量的賦值語句插入循環(huán)之前。

雖然單個常數是循環(huán)不變的，但是提升它沒有意義，因此算法不對單個常數做提升。

另外對于表達式中循環(huán)不變式的識別，表達式中操作數的順序非常關鍵；不同的操作數順序不但影響了某些循環(huán)不變的表達式能否被發(fā)現(xiàn)，還決定了目標編譯器中某些對于分支的優(yōu)化能否被觸發(fā)。因此，在發(fā)現(xiàn)循環(huán)不變量之前需要先對表達式進行規(guī)范化，利用運算符的結合律對表達式進行重結合等變換。

算法1循環(huán)不變量外提算法

算法1 給出了循環(huán)不變式外提算法LICM(T)的偽代碼。該算法接受原始程序T作為輸入參數，并返回優(yōu)化后的程序。算法首先調用規(guī)范化函數Normalize(T)對程序T中的語句和表達式做規(guī)范化。然后算法按照一定的順序，遍歷程序T中的語句S，并按語句S的可能語法形式進行分類討論：若S為賦值語句x=e，并且賦值右側表達式e被標記為“不可提升”，那么算法標記被該賦值語句定義的變量x為“不可提升”；注意到，算法中使用映射Promote[e]來記錄給定的表達式e是否可提升，特定的常量值⊥代表不可提升。若S為循環(huán)語句loop(x,e)，其中x是循環(huán)遍歷，e是循環(huán)體，則算法標記S的循環(huán)變量x為“不可提升”。對于語句S中那些沒有被標記為“不可提升”的表達式e，都可認為是循環(huán)不變量，因此，算法生成一個新變量z，將S中的表達式e用變量z改寫，并將識別的循環(huán)不變式的候選z←e添加到循環(huán)不變式集合R中供后續(xù)處理。

算法接著對循環(huán)不變式集合R中的每個循環(huán)不變式z←e進行分析。算法調用代價函數cost(e)計算外提表達式e的代價，當且僅當該代價大于等于給定的閾值K，并且表達式e沒有副作用時，算法才真正將不變式z←e外提；否則，算法放棄對不變式z←e的外提嘗試，并恢復被替換的表達式。代價常數K的取值與目標平臺硬件和編譯器密切相關，在通常情況下可取K為1。

LICM 算法的輸入T是一棵樹，設樹中的語句節(jié)點和表達式節(jié)點數量為n，Normalize 函數時間復雜度和空間復雜度為O(n)，具體分析在下節(jié)給出。LICM 算法在規(guī)范化之后，第一個循環(huán)遍歷T中的語句節(jié)點，該循環(huán)內的第一個內層循環(huán)遍歷語句節(jié)點內的表達式節(jié)點，并記錄對合法的表達式的賦值語句，第二個內存循環(huán)處理記錄的合法表達式的賦值，cost 函數的計算雖然是遞歸的，但是在計算過程中存儲每個表達式對應代價值，因此每個表達式只會被計算一次，所有cost的執(zhí)行時間和表達式節(jié)點數量線性相關。因此整個算法的時間復雜度和輸入T中語句節(jié)點和表達式節(jié)點數量線性相關，LICM 算法的時間復雜度為O(n)。整個LICM 算法除T外所涉及的存儲結構為R，Promote 和cost 記錄表達式的代價值所需的存儲。以上三個存儲結構的任意一個的大小最多為表達式節(jié)點數量，因此空間復雜度為O(n)。Halide 編譯器中的循環(huán)不變量外提算法也是在樹結構的中間表示上進行，時間復雜度為O(n)，空間復雜度為O(n)。本文算法的時間和空間復雜度都不差于Halide編譯器的循環(huán)不變量外提算法。

傳統(tǒng)編譯器中的循環(huán)不變量外提算法，如LLVM中的外提算法，都是在基于基本塊的流圖上設計和實現(xiàn)的。這種實現(xiàn)方式主要的缺點是程序在下降的過程中丟失了如循環(huán)、作用域以及條件表達式等高層程序信息，因此在執(zhí)行循環(huán)不變量外提算法前這些必要信息都要通過額外的操作來恢復。另外，在編譯器低層次的中間表示中，表達式或語句的表示形式發(fā)生了改變，這導致原本在高級表示上顯而易見的優(yōu)化時機在后面的優(yōu)化中很難被發(fā)現(xiàn)。而本算法充分利用了深度學習編譯器的主要特點，面向其高級中間表示實現(xiàn)循環(huán)不變式外提，彌補了傳統(tǒng)循環(huán)不變量外提算法的不足。另外，該算法和已有的類似循環(huán)不變式算法也很不相同，如本算法使用的新的前置處理方式對表達式做變換，而不僅僅是像Halide 那樣只對加法減法表達式做重結合；這使得本算法對冗余消除得更為徹底，避免了需要對外提的不變量做公共子表達式消除的額外工作，讓目標平臺編譯器更加簡化。

3.3 規(guī)范化

規(guī)范化指的是通過對語句和表達式進行保語義的變換，使得變換后的程序更利于后續(xù)循環(huán)不變量外提，或目標平臺編譯器的優(yōu)化。規(guī)范化最主要的操作包括：調整表達式中操作數的順序和對條件表達式和分支的合并操作。調整操作數順序是指利用特定的代數性質，如結合律、交換律和分配律等，對表達式進行重結合優(yōu)化，將其劃分成常數部分、循環(huán)不變部分和變量部分。盡管這本身并不能改進性能，但由于某個表達式的子表達式是否為循環(huán)不變表達式，很大程度上取決于操作數的運算順序和結合關系[16]，因此，調整表達式操作數順序能夠給循環(huán)不變式外提優(yōu)化帶來明顯的提升效果。例如，假設下面給定的兩個表達式中，只有i是循環(huán)迭代變量，j為外層循環(huán)的迭代變量：

表達式（1）對應的抽象語法樹如圖2 所示，算法會計算得到兩個候選的循環(huán)不變式：(3 ＜j)和(j＜56)。而表達式（1）在規(guī)范化過程中，把(3 ＜j)和(j＜56)放到了一起，得到表達式（2），其對應的抽象語法樹如圖3 所示，表達式順序調整后的作用如下：變換之前的兩個循環(huán)不變式(3 ＜j)和(j＜56)在調整順序后形成了新的循環(huán)不變式((3 ＜j)&&(j＜56))，比原先的多外提了一個取并的運算操作。根據以上的觀察，調整表達式的操作數結合順序的思路為：對于某個滿足結合律的運算，把操作數按照所包含循環(huán)迭代變量對應循環(huán)的嵌套深度，以遞增或遞減進行排序后結合到一起，這樣的操作能夠將分散的循環(huán)不變式結合到一起，利于循環(huán)不變量表達式的發(fā)現(xiàn)?；谝陨咸攸c，可按如下步驟把循環(huán)不變的表達式結合到一起：

圖2 表達式（1）的抽象語法樹Fig.2 Abstract syntax tree of expression(1)

圖3 表達式（2）的抽象語法樹Fig.3 Abstract syntax tree of expression（2）

（1）為每一個表達式e綁定一個整數rank：

（2）若表達式e為常數，則為其綁定rank為0；

（3）若表達式e為循環(huán)迭代變量，則為其綁定rank為嵌套循環(huán)的深度，否則：

①綁定定義該變量的操作數的最大rank值。

②對于滿足結合律的運算，按照rank值遞增為每個表達式排序后，重新結合為新的表達式。

條件表達式對應判斷條件的成立和不成立，都有對應的返回值，而在嵌套的條件表達式中，多個判斷條件成立與否可能對應于同一個返回值，對此，可以合并這些具有相同返回值的判斷條件，使得變換后的表達式中的循環(huán)不變量更完整。例如，假設下面的表達式中只有i是循環(huán)迭代變量：

表達式（3）對應的語法樹如圖4 所示，算法確定的候選循環(huán)不變式為3 ＜j和j＜56。表達式（3）經過化簡后得到條件表達式：

圖4 表達式（3）的抽象語法樹Fig.4 Abstract syntax tree of expression（3）

其對應的語法樹如圖5 所示，算法發(fā)現(xiàn)了更好的候選循環(huán)不變表達式((3 ＜j)&&(j＜56))，并且((3 ＜i)&&(i＜56))和((3 ＜j)&&(j＜56))這兩個表達式的形式更利于實現(xiàn)與布爾表達式相關的窺孔優(yōu)化。最后，對于相鄰的具有相同跳轉條件的分支，算法可以合并這些分支語句。算法2 給出了規(guī)范化算法Normalize(T)的偽代碼實現(xiàn)。

圖5 表達式（4）的抽象語法樹Fig.5 Abstract syntax tree of expression（4）

算法2表達式規(guī)范化算法

該算法接受原始程序T作為輸入，并輸出優(yōu)化后的程序。算法遍歷程序T中的表達式e，并根據表達式e的語法形式對其進行處理：如果表達式e為條件表達式if(e1,e2,e3)，則算法調用函數condCollapse(e)，對可以進行判斷條件合并的表達式e進行優(yōu)化，并返回優(yōu)化后的結果。接下來，算法按照前面討論的步驟，計算表達式e中的所有變量和子表達式的rank值。對于完成rank值計算的表達式e，算法遍歷其所有滿足結合律的操作符⊕，調用deconstruct(e,⊕)將由操作符⊕連接的子表達式轉換為有序子表達式序列L；然后，算法調用函數stableSort(L)使用穩(wěn)定排序算法將L按照rank值從小到大排序并返回排序后的序列L；最后，算法調用函數groupExpressions(L,⊕)將L中相同rank值的子表達式用操作符⊕連接為新的表達式并返回變換后的序列L，并調用函數construct(L,⊕)用操作符⊕依序將L中的所有子表達式連接，并返回變換后的最終結果e。算法對程序T中的所有表達式遍歷結束后，調用函數ifconcat(T)合并T中相鄰的具有相同分支條件的分支語句，并返回變換后的程序T。函數ifconcat(T)本質上完成控制流優(yōu)化，通過對分支語句進行合并優(yōu)化，可有效減少語句和表達式的數量，減少優(yōu)化算法的執(zhí)行時間。Normalize 算法的輸入T是一棵樹，設樹中的語句節(jié)點和表達式的數量為n，外層循環(huán)遍歷T中的表達式節(jié)點，condCollapse 所需時間和判斷條件數量線性相關；loopDepth 函數遞歸求每個變量的嵌套深度，在遞歸過程中保存遇到的變量的嵌套深度，因此每個變量的深度只會被求一次，所需時間和表達式數量線性相關。最后一個內層循環(huán)對表達式重結合，deconstruct、groupExpressions 和construct 都是和表達式數量相關的線性時間；stableSort 排序算法根據rank值進行排序，因為每個表達式的rank值都是0 到嵌套深度之間的整數值，可以使用計數排序，該排序算法為線性時間復雜度，ifconcat 和T的語句數線性相關，因此Normalize 算法時間復雜度為O(n)?？臻g復雜度的分析和LICM 的思路相同，也為O(n)。

傳統(tǒng)的標量編譯器在進行循環(huán)不變式外提前，也都會執(zhí)行代數化簡和表達式重結合等優(yōu)化遍，來完成前置的處理和準備工作。而在深度學習編譯器中，和內存讀寫相關操作的優(yōu)化應該在更加低級的中間表示上進行，代數化簡和常量傳播等優(yōu)化由TVM 原生的simplify 優(yōu)化遍完成，因此，規(guī)范化算法可以更專注于對下標和條件表達式進行特定變換，以有利于循環(huán)不變量外提優(yōu)化和控制流優(yōu)化的實現(xiàn)。

3.4 代價函數

代價函數cost(e)接受表達式e作為輸入，計算并返回e的外提代價值。當且僅當該函數的返回值大于等于給定的閾值K時，編譯器才認為有足夠的收益，并將該表達式e進行外提。需要計算代價函數cost(e)的原因是TVM 生成的高層算子代碼，還要經過目標平臺編譯器的編譯，才能生成真正目標硬件上的可執(zhí)行程序。由于目標平臺編譯器會在更低層次的中間表示上對程序做一系列優(yōu)化，兩個不同層次的優(yōu)化之間可能產生矛盾，比如外提循環(huán)不變量和后端寄存器分配優(yōu)化存在矛盾，當循環(huán)內的一條指令的執(zhí)行需要一個不變表達式的執(zhí)行結果時，有兩種方法可以獲取這個結果：（1）每次指令執(zhí)行前重新計算；（2）將計算結果保存在變量中，每次指令執(zhí)行前使用該變量。方法（1）中重復計算相同的表達式增加了冗余計算，但是減少了變量聲明周期，降低了寄存器壓力，利于有序寄存器分配優(yōu)化。方法（2）消除了多次計算不變表達式的冗余計算，但是增加了變量生命周期，增加寄存器壓力，可能導致后續(xù)寄存器分配優(yōu)化效果下降。為了權衡消除冗余計算和寄存器分配優(yōu)化之間的矛盾，避免兩個不同層次的優(yōu)化相互干擾，設計了一個代價函數cost(e)來計算表達式的代價，指導算法外提循環(huán)不變式。表達式代價值代表了計算表達式所需開銷，也就是外提表達式的收益，因此僅當代價值大于某一閾值K時，算法才對該表達式做外提。對表達式e的代價函數cost(e)定義為：

當e是常數或變量時，其代價定義為0；當e是復合表達式時，其代價遞歸定義在其子表達式a上，另外加上其運算符⊕的代價。下面以一個具體的表達式x/1024+y為例，詳細說明代價值的計算方式。根據定義有：

不同的運算⊕根據其需要的時鐘周期的不同定義了不同的代價，其中加法指令需要的時鐘周期較少，設置cost(+)=1，除法指令需要的時鐘周期較多，設置cost(/)=3。因此有cost(x/1 024+y)=1+3+0+0+0=4。這個基礎的代價函數cost(e)可以修改為適合于不同硬件平臺的代價函數，例如，對于GPU 平臺，其平臺上的編譯器NVCC 通常選擇展開迭代范圍很小的循環(huán)，并在此基礎上執(zhí)行其他的優(yōu)化。因此，可以設置循環(huán)迭代范圍很小的循環(huán)中的表達式代價值為0，避免兩個不同層次的優(yōu)化互相干擾。

4 實驗結果與分析

4.1 研究問題

實驗主要回答以下兩個研究問題：

（1）性能。本文提出的新的循環(huán)不變式外提算法是否能夠提升實際TVM 生成算子的性能？以及對哪些算子會造成一些特殊情況性能下降？如果有下降，原因是什么？

（2）正確性。本文算法改變了操作數的順序，對于優(yōu)化后生成的算子的正確性（即計算精度）是否會造成影響？

4.2 實驗平臺

實驗使用英偉達的NVCC 編譯器的10.0 版本作為GPU 平臺的編譯器。所有實驗數據都在表1 所示配置的服務器上收集得到。

表1 實驗環(huán)境Table 1 Setup of experiment

4.3 數據集

實驗選擇了TVM 的TOPI 算子庫中的部分算子在不同輸入規(guī)模下的實現(xiàn)，作為本次實驗的基準測試集，本次測試涉及27 個算子和511 個測例。表2 列出了本次實現(xiàn)選擇的所有算子及其測例數量。這些算子屬于深度學習框架和TVM 常用的算子，并且實驗對每一個算子都提供了大量在常見神經網絡中不同輸入規(guī)模的測例，在這些測例上的實驗結果具有一定的代表性。

表2 算子測試集Table 2 Benchmark of operators

4.4 實驗結果與分析

在研究算法對性能的影響時，為了減少實驗誤差，認為那些優(yōu)化前后生成代碼相同和運行時間小于50 μs 的測例為無效測例，且不考慮運行時間波動低于1%的測例。表2 給出了每個算子的有效測例數量，稱有效測例數量大于0 的算子為有效算子。定義算子運行加速比S為：

其中，T0為原生TVM 生成算子的運行時間，T1為本次實驗的原型系統(tǒng)生成算子的運行時間（即經循環(huán)不變式外提優(yōu)化后算子運行時間）。定義算子平均加速比A為所有測例的加速比的算術平均值：

圖6（a）給出了有效算子的平均加速比，圖6（b）給出了有效算子運行時間加速比。圖中的綠色虛線代表加速比為101%，橙色實線代表加速比為99%。從算子平均加速比的角度來看，相比原生TVM 編譯得到的算子，本文算法編譯得到的算子中有10 個算子性能得到了提升，占有效算子的47.6%，性能提升最明顯的算子為conv3d_winograd，其加速比達到141.46%；只有1 個算子性能下降，占有效算子的4.7%，該性能下降的算子為conv2d_hwcn，其加速比為97.60%；其余算子性能不變。從算子總時間加速比的角度來看，相比原生TVM 編譯得到的算子，本文算法編譯得到的算子中有10 個算子性能得到了提升，占有效算子的47.6%，性能提升最明顯的算子為conv3d_transpose_ncdhw，加速比為142.91%；只有1個算子性能下降，占有效算子的4.7%，該性能下降的算子為conv2d_hwcn，加速比為97.10%；其余算子性能不變。所有算子的總時間加速比為122.7%。對性能得到提升的測例，進一步分析了其優(yōu)化前后生成的代碼，發(fā)現(xiàn)表達式中的循環(huán)不變式被提出了循環(huán)外，分支和條件表達式也得到了相應的優(yōu)化。首先，循環(huán)內的指令數目減少了，一部分循環(huán)不變的指令被移到了循環(huán)外；其次，計算謂詞的指令得到了優(yōu)化，一些有公共操作數的比較指令被減法或其他指令替換；再次，分支指令也用謂詞執(zhí)行來替換，減少了分支指令的數量；最后，因為循環(huán)內指令減少或者簡化，nvcc 編譯器對內層循環(huán)進行了更大范圍的循環(huán)展開，同時其他的優(yōu)化遍對展開后的循環(huán)體進行了更好的優(yōu)化。對優(yōu)化后性能下降的一個測例，進一步分析發(fā)現(xiàn)只有少量的簡單表達式被提出了循環(huán)外，循環(huán)內的控制流結構和指令數基本沒有發(fā)生變化，這些簡單指令的延遲原本就會被流水線隱藏；另外，一些指令和同步指令的相對位置發(fā)生了變化，一些原本在同步指令后的計算被移動到了同步指令之前，可能改變后指令調度時的優(yōu)化效果。需要特別注意的是，盡管這個特定算子的性能出現(xiàn)下降，但其2.9%的小幅下降仍然在可接受的范圍內。

圖6 算子加速比Fig.6 Operator speedups

另外，根據本算法對于不同類型的算子性能提升幅度不同，具體分析了本文循環(huán)不變式外提算法的使用場景。對于conv3d_winograd、conv3d_transpose_ncdhw 這類實現(xiàn)復雜的算子，經過前端調度優(yōu)化后得到的中間表示包含大量的嵌套循環(huán)和嵌套條件表達式，本文算法對這類使用場景能夠減少循環(huán)內的計算和分支，利于后端編譯器優(yōu)化的進行。在這類使用場景下本文算法能夠起到顯著的性能提升效果。對于像conv2d_hwcn、depthwise_conv2d_back_input 等這類算子代碼中循環(huán)和條件表達式嵌套數較少，表達式形式較為簡單，部分簡單形式的冗余計算在底層編譯器優(yōu)化中也能被消除。且由于外提計算導致其中同步指令的相對位置發(fā)生變化，可能改變后續(xù)指令調度優(yōu)化的效果。在這類使用場景下本文算法對于算子性能的提升可能得不到令人滿意的效果。

為了進一步證明算法改進的有效性，選擇Halide編譯器的循環(huán)不變量外提優(yōu)化與本文算法的優(yōu)化效果進行比較，實驗將Halide 編譯器的循環(huán)不變量外提優(yōu)化遍經過適配移植到了TVM 上。在原生優(yōu)化遍的基礎上，添加了對于條件表達式和其他符合運算結合律的運算的外提支持。實驗結果如圖7 所示，本次實驗中最小的加速比為90%，因此將加速比減去89%后以2 為底取對數，以上操作是為了直觀地對比兩個算法的優(yōu)劣，每個縱坐標計算方式為lb((s-89%)×100)，其中，s表示算子加速比。紅色的柱形表示Halide 編譯器的循環(huán)不變量外提算法的優(yōu)化效果。本次實驗的有效算子測例中，忽略性能波動小于1%的測例。同時本文認為兩個算法優(yōu)化算子性能差異小于1%時為實驗誤差。實驗結果顯示，本文算法相比于另一比較算法，在5 個算子的性能提升效果上明顯更優(yōu)，占有效測試算子的23.81%；僅在1 個算子的性能提升效果上較差。其余71.42%的算子測例的性能提升效果差別不大。從優(yōu)化效果明顯提升的例子，如conv3d_transpose_ncdhw、depthwise_conv2d_back_weight 和conv3d_winograd 等復雜算子測例的生成代碼中觀察到，相比比較算法，本文算法在對于嵌套條件表達式、不變式的重結合方面表現(xiàn)都更好，外提了一部分沒有被比較算法發(fā)現(xiàn)的不變式，減少了跳轉指令數量。而觀察性能下降的例子conv2d_hwcn，經過Halide 循環(huán)不變量外提算法優(yōu)化后算子性能沒有發(fā)生變化，而本算法的變化導致了性能的輕微下降，大約為2.9%。生成代碼中控制流結構和指令數基本不變，但有一部分計算和同步指令的相對位置發(fā)生了變化，可能導致底層編譯優(yōu)化效果的下降。從總體效果來看，本文算法對算子性能提升效果明顯優(yōu)于比較算法，進一步證明了本文對于循環(huán)不變量外提算法改進的有效性。

圖7 本文算法與Halide編譯器的LICM 算法對比Fig.7 Comparison between proposed algorithm and Halide’s LICM

本次實驗研究的第二個問題為算子的準確性，算子的準確性主要由算子計算結果的精度來衡量。影響算子計算準確性的主要因素有硬件計算單元本身精度和數據計算順序。硬件計算單元精度和實驗硬件平臺緊密相關，本次實驗選擇NVIDIA Tesla P4圖形處理器。數據計算順序主要依賴于前端的調度，另外優(yōu)化遍對于中間表示的變換也會改變計算順序，本文算法實現(xiàn)的優(yōu)化遍只對于滿足運算結合律的運算改變計算順序，不會改變運算結果。為研究本文算法的正確性，對所有測例在GPU 上的運行結果與Python 實現(xiàn)的等價測例在CPU 上的運行結果進行了比較，由于浮點數的相等測試的約束，本實驗設置誤差容忍度為1E-5。實驗結果表明100%的測例實驗結果誤差小于1E-5，這證實了本文算法的正確性。

4.5 結論

本文以TVM（0.7 版本）為基礎，實現(xiàn)了添加循環(huán)不變式外提優(yōu)化算法的原型系統(tǒng)，并在TVM 的TOPI算子測試集上針對27 個算子的511 個測例進行測試，并收集實驗結果。對實驗結果的分析表明，本文提出的循環(huán)不變式外提算法，彌補了閉源編譯器NVCC在對復雜表達式的冗余刪除和控制流優(yōu)化存在的不足，完善了TVM 生成算子的優(yōu)化過程，對算子性能提升產生了積極的作用。

5 結語

本文針對深度學習編譯器，提出了一種新的循環(huán)不變式外提優(yōu)化算法，并基于開源TVM 的0.7 版本實現(xiàn)了原型系統(tǒng)。對TVM 中的TOPI 算子庫中的27個算子和511 個測例進行了實驗，實驗結果表明該新優(yōu)化算法在保證正確性的情況下有效提升了算子性能。本文工作為TVM 或者其他深度學習編譯器移植傳統(tǒng)循環(huán)不變式外提算法提供了有益參考。