今天來閱讀一下最近 OSDI 放出的微軟的 Roller 這篇論文，題目為：《Roller: Fast and Efficient Tensor Compilation for Deep Learning》

• 論文鏈接：https://www.usenix.org/conference/osdi22/presentation/zhu
• 代碼鏈接：https://github.com/microsoft/nnfusion/

前段時間我分享了一下 OSDI 2021 PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections》 這篇論文的解讀。去年也分享了 OSDI 2020 《Ansor : Generating High-Performance Tensor Programs for Deep Learning》這篇論文的解讀。這兩篇論文的解讀可以在這個地址：https://github.com/BBuf/tvm_mlir_learn/tree/main/paper_reading 或者知乎主頁找到 。Ansor 的主要貢獻(xiàn)是做到了自動尋找高效的Schedule（循環(huán)展開、合并、分塊、緩存使用、改變并行度等等），不再需要開發(fā)者在TVM中基于Tensor Expression手寫Schedule模板，大大增強(qiáng)了算子編譯器（Tensor Compiler）的易用性并且對一些典型的算子和模型效果也很不錯，算是AutoTVM的升級版（因為AutoTVM還需要手動指定需要搜索的Schedule模板：https://zhuanlan.zhihu.com/p/508283737）。PET則不關(guān)心算子的Schedule，而是從部分等價變換的新角度出發(fā)去增加并行度或者改善緩存從而達(dá)到加速效果，和Roller這篇論文沒什么關(guān)系，其實讀不讀都沒關(guān)系。

最近不少人問一些tvm相關(guān)的問題，我也是業(yè)余看了下所以很多時候不能很好的解答，我建立一個討論TVM的微信群吧，有需要的讀者可以加一下互相問一問。請加微信 bbuf23333 入群，備注一下tvm吧。另外業(yè)余接觸編譯器這一年整理的這個知識倉庫已經(jīng)有500+ star了，謝謝大家。希望能得到更多關(guān)注。

https://github.com/BBuf/tvm_mlir_learn

無論是Ansor，AutoTVM還是PET（一部分代碼生成也是基于TVM AutoTVM/Ansor的）它們都面臨了同樣一個問題，那就是在對算子的Schedule進(jìn)行搜索時需要耗費大量的時間，在特定硬件上對一個常見的視覺模型進(jìn)行自動調(diào)優(yōu)和生成代碼kennel需要數(shù)小時。這嚴(yán)重阻礙了AI編譯器應(yīng)用于模型部署?；谶@個痛點，Roller橫空出世。

0x0. 標(biāo)題&作者&摘要

ROLLER：一個用于深度學(xué)習(xí)的快速高效的張量編譯器。作者來自微軟亞洲研究院以及多倫多大學(xué)等多所高校。

現(xiàn)代的張量編譯器雖然取得了很多的進(jìn)展，但通常這些編譯器都需要小時計的時間去搜索和生成高效的Kernel，這是因為現(xiàn)有張量編譯器通常指定的搜索空間很大。為了解決編譯時間長的問題，本文提出了Roller，它的核心是rTile，這是一種新的Tile抽象，它封裝了和底層加速器的關(guān)鍵特性一致的張量shape，從而通過限制shape的選擇實現(xiàn)高效的運行。Roller采用了基于rTile的遞歸構(gòu)建算法來生成目標(biāo)程序（rProgram）。最終，Roller可以在幾秒內(nèi)就生產(chǎn)高效的Kernel，性能可以媲美目前主流加速器上的其它張量編譯器，并且為IPU等新的加速器生產(chǎn)更好的Kernel。

還不能看出什么，繼續(xù)往下看吧。這里說的tile就是對輸入進(jìn)行分塊以適應(yīng)硬件的內(nèi)存結(jié)構(gòu)，我在之前的文章有詳細(xì)講到，不了解的同學(xué)可以先看一眼tile這部分的科普：https://zhuanlan.zhihu.com/p/508283737 。

0x1. 介紹

深度神經(jīng)網(wǎng)絡(luò)越來越重要，深度學(xué)習(xí)編譯器在硬件上生成高效的Kernel也越來越重要，并且取得了很多成功。但是當(dāng)代的編譯器在生成高效的Kernel時往往需要搜索數(shù)個小時甚至數(shù)天，因為它們都是把這些網(wǎng)絡(luò)中的算子實現(xiàn)成多重循環(huán)嵌套。張量編譯器通常需要對已實現(xiàn)的多重循環(huán)計算進(jìn)行循環(huán)展開、合并、分塊、緩存使用、改變并行度等調(diào)整以適應(yīng)硬件的內(nèi)存結(jié)構(gòu)（比如CPU的三級Cache和CUDA的global memory,l2 cache, l1 cache結(jié)構(gòu)）或者硬件特性（比如向量化，并行化）。這里涉及到非常大和復(fù)雜的搜索空間，所以搜索時間會很久。這篇文章提出的Roller解決了搜索時間長的問題，它有如下幾個特點。

• 首先，Roller不把DNN中的算子計算視為多層嵌套循環(huán)，而是視作數(shù)據(jù)處理管道，其中數(shù)據(jù)塊(tile) 在具有并行執(zhí)行單元（如GPU SM）和內(nèi)存層次結(jié)構(gòu)抽象的硬件上移動和處理。生成高效的Kernel的目標(biāo)變成了提高流水線吞吐量的目標(biāo)。

Roller將算子的計算過程建模為基于數(shù)據(jù)塊（tile）的流水線，即將不同大小的數(shù)據(jù)塊從多級內(nèi)存結(jié)構(gòu)中搬運到處理器如SM計算并逐級寫回。

• 然后，為了使得基于數(shù)據(jù)塊的流水線吞吐量最大化，要求每一級的數(shù)據(jù)塊（Tile）shape都必須匹配（論文中叫對齊）硬件的參數(shù)設(shè)置，比如memory bank, memory transaction length, 和 minimum schedulable unit (e.g., warp size in GPUs)這些和內(nèi)存帶寬以及并行度相關(guān)的設(shè)置。這個約束不僅可以使得張量程序在每一級內(nèi)存中都擁有很好的計算效率，這還大大降低了以前以多重循環(huán)為基礎(chǔ)的參數(shù)搜索空間，從而解決張量編譯器在編譯時因為搜索Schedule耗費的大量時間。
• 最后，對齊硬件的數(shù)據(jù)處理管道的性能是高度可預(yù)測的。因為內(nèi)存吞吐量可以從硬件規(guī)范或者Benchmark測試得出，這大大簡化了對不同硬件進(jìn)行對齊后做性能估計的難度，并不再需要基于硬件去構(gòu)建復(fù)雜的代價模型來估計性能。

基于這些想法，Roller提出了rTile，這是一種新的抽象，它封裝了和硬件加速器的關(guān)鍵特征和輸入張量shape一致的數(shù)據(jù)塊（Tile）shape（后面會詳細(xì)看）。然后將數(shù)據(jù)處理管道描述為基于rTile的程序（rProgram），由Load, Store, Compute 三個接口組成，作用于rTile。為了構(gòu)建高效的rProgram，Roller遵循了一個scale-up-then-scale-out的方法。它首先執(zhí)行Scale-up的過程，該過程采用基于rTile的遞歸構(gòu)造方法（Figure8）逐漸增加rTile shape大小，來構(gòu)造一個飽和加速器單個執(zhí)行單元（如SM）的rProgram。然后執(zhí)行Scale-out的過程，由于深度學(xué)習(xí)的計算模式和加速器的并行執(zhí)行單元的同質(zhì)性，它只是將生成的rProgram復(fù)制到其它并行執(zhí)行單元。這里的scale-up-then-scale-out可以叫做縱擴(kuò)和橫擴(kuò)。

Roller可以在沒有顯著開銷的情況下評估不同rTiles的性能。每種算子可以簡單的測試一下峰值和帶寬。由于對齊了硬件結(jié)構(gòu)，其它關(guān)鍵的性能因素比如rTile的內(nèi)存壓力可以從硬件規(guī)則分析得到。這樣就得到了一個高效的微評測模型，避免了其它編譯器所需的對每個配置進(jìn)行昂貴的在線分析，從而顯著加速了編譯過程。此外，由于嚴(yán)格的對齊要求，遞歸構(gòu)造過程可以快速生產(chǎn)一些想要的rTiles和rPrograms。綜合一下，Roller可以在幾秒內(nèi)生成高效的Kernel。

作者團(tuán)隊在TVM和Rammer（Rammer可以看：https://www.msra.cn/zh-cn/news/features/osdi-2020-rammer）之上實現(xiàn)了Roller并開源了代碼。大量的實驗表明Roller可以在幾秒內(nèi)生產(chǎn)高度優(yōu)化的Kernel，特別是對于大型自定義的高成本算子。這在編譯時間上實現(xiàn)了3個數(shù)量級的改進(jìn)。Roller生成的Kernel可以和最先進(jìn)的張量編譯器乃至硬件廠商提供的加速庫相媲美，并且通常性更好（指接入新的硬件）。使用三個 rTile-based 的接口（Load, Store, Compute）描述一個程序，Roller可以輕松適應(yīng)不同的加速器如AMD GPU和Graphcore IPU。

0x2. 動機(jī)和關(guān)鍵觀察

• Excessive compilation time：張量編譯器編譯時間太長，影響生產(chǎn)。
• Observation and insights：我們觀察到對于深度學(xué)習(xí)算子的計算有不同的視角。以矩陣乘法為例子來說明我們的觀察。和將MatMul視為跨三重循環(huán)的現(xiàn)有編譯器不同，算子的計算過程也是一個數(shù)據(jù)處理管道。我們可以從A和B Load 2個子矩陣（tile），Compute 兩個子矩陣，Store 結(jié)果到C的內(nèi)存中。所以計算的性能取決于 Load-Compute-Store 管道移動一個 Tile 有多快。

影響流水線中所有步驟關(guān)鍵性能的因素是Tile shape和一維內(nèi)存空間中的布局。Figure1(a)說明C中一個元素的計算和內(nèi)存訪問的模式。假設(shè)所有矩陣存儲在行優(yōu)先的布局中，從B加載列會有1個跨步訪問。假設(shè)這里的事務(wù)內(nèi)存長度（the memory transaction length）是4，那么就有3/4的冗余數(shù)據(jù)讀取。所以數(shù)據(jù)塊的形狀應(yīng)該和內(nèi)存事務(wù)長度對齊，以實現(xiàn)高效的內(nèi)存訪問。在Figure1(b)中，當(dāng)以1x4 Tile的粒度計算B時不會有內(nèi)存帶寬浪費。除了內(nèi)存對齊之外，數(shù)據(jù)的Tile shape還應(yīng)該和硬件執(zhí)行單元如并行線程數(shù)對齊以避免浪費計算周期。此外，由于Cache的存在，Tile shape也會影響數(shù)據(jù)重用機(jī)會。例如Figure1(a)每次計算1x1 tile時需要讀取2mnk個數(shù)據(jù)。然而在Figure1(b)中只需要1.25mnk次讀取，因為來自A的一次數(shù)據(jù)讀取可以重復(fù)使用4次。如果沿M維度的tile 大小設(shè)置為4x4，總的reads可以減少到0.5mnk，總的數(shù)據(jù)讀取效率比Figure1(a)提高了10倍。

0x3. 系統(tǒng)設(shè)計

下面的Figure2描述了Roller的系統(tǒng)設(shè)計。Roller的輸入是使用TE表達(dá)式。該表達(dá)式由用戶生產(chǎn)或者從其它編譯器生成（這一步可能會發(fā)生一些融合操作）。Roller從TE中提取張量形狀并基于硬件規(guī)范來構(gòu)建rTiles，即對齊硬件的構(gòu)建塊?；趓Tiles，Roller提出了一種橫擴(kuò)縱擴(kuò)遞歸構(gòu)造算法，用于生成描述數(shù)據(jù)處理管道的高效張量化程序（rProgram）。在生成rProgram時，構(gòu)建算法通過微觀性能模型評估構(gòu)建的rProgram的性能，從而識別出良好的rTile配置。它建立在通過硬件抽象描述的設(shè)備上，僅公開和rTiles相關(guān)的接口：Load/Save/Compute。構(gòu)建的rProgram最終通過Codegen生成特定設(shè)備的最終Kernel。

0x3.1 Tensor Expression and rTile

Roller將TVM中引入的Tensor Expression引入作為編譯器的輸入，Tensor Experssion這里不講了，如果不了解可以看一下TVM里面chen tianqi寫的文檔。https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html

Roller引入rTile作為基本計算單元來組成張量計算。如Figure3所示，rTile封裝了沿給定張量表達(dá)式的expr的每個循環(huán)軸定義的多維tile shape。給定shape和expr，rTile可以靜態(tài)推斷所涉及的輸入和輸出數(shù)據(jù)塊。例如，沿軸i, j, k的tile shape表示上述Matmul表達(dá)式的rTile，其中每個rTile加載來自A的4x2個數(shù)據(jù)以及來自B的2x4個數(shù)據(jù)，進(jìn)行總共4x2x4 次 mul-add計算，并將4x4的數(shù)據(jù)tile寫回到C，如Figure4所示。

rTile的一個獨特屬性在于它必須和給定張量表達(dá)式中的底層硬件特征和輸入Tensor shape保持一致。所有這些對齊方式都由Figure3里rTile 的 shape 和 storage_padding 來控制，它們分別代表 rTile 的邏輯形式和物理布局。接下來，詳細(xì)闡述對齊的詳細(xì)要求：

• Alignment with the hardware execution unit 。首先，rTile的shape必須和它運行的執(zhí)行單元的并行度對齊。例如，在GPU上運行 rTile 的shape 大小必須是 wrap size的倍數(shù)比如 32 來達(dá)到最大的計算效率。當(dāng)在NVIDIA GPU中使用TensorCore時，rTile shape大小應(yīng)該是 16x16x16 的倍數(shù)。
• Alignment with memory transaction 。其次，數(shù)據(jù)塊（Tile）的 shape 應(yīng)該和內(nèi)存事務(wù)的長度保持一致，以實現(xiàn)最佳內(nèi)存訪問。具體來說，對于rTile的每個數(shù)據(jù)塊我們都應(yīng)該保證它的Leading dimension（如行優(yōu)先Tensor中的最內(nèi)層維度）是內(nèi)存事務(wù)長度的倍數(shù)。如Figure5(a)所示，在Roller中，張量內(nèi)存以緩存對齊的方式分配。因此，rTile可以避免浪費任何的內(nèi)存讀取，因為它的 shape 是和內(nèi)存事務(wù)長度對齊的。

最大程度的利用全局內(nèi)存帶寬，提高全局內(nèi)存加載效率是優(yōu)化Kernel的基本條件，非對齊的內(nèi)存會造成帶寬浪費，可參考：https://face2ai.com/CUDA-F-4-3-%E5%86%85%E5%AD%98%E8%AE%BF%E9%97%AE%E6%A8%A1%E5%BC%8F/

• Alignment with memory bank. 第三，數(shù)據(jù)塊的內(nèi)存布局應(yīng)該和Memory Bank對齊，以避免讀取沖突。例如，在Figure5(b)中數(shù)據(jù)塊a（shape為[3, 4] ）跨4個bank保存在內(nèi)存中，并由形狀為 [3, 1] 的塊讀取。將這個形狀為[3, 1]的小塊中的數(shù)據(jù)存儲在一個bank的naive方法將導(dǎo)致加載沖突。rTile通過padding來避免這種低效。給定一個Leading dimension為N的數(shù)據(jù)塊，由另外一個Leading dimension為n的塊讀取，我們延N維度做一個padding_size大小的padding。

其中B和L分別是bank數(shù)量和bank的寬度。每一個維度的padding大小被計算出來后存到Figure3中的storage_padding字段。對于Figur5(b)，通過padding_size為1的填充，所有的值 [3x1] 分布在不同的bank中，可以高效讀取。

GPU Shared Memory bank conflict: https://blog.csdn.net/Bruce_0712/article/details/65447608

• Alignment with tensor shape 最后，rTile的shape應(yīng)該和輸入張量表達(dá)式的張量shape對齊。否則，計算不能被rTile均勻且分，浪費計算資源或者產(chǎn)生大量的邊界檢查開銷。一個簡單的解決方案是沿著Tensor的維度（大小為 ）進(jìn)行padding，padding的大小為，使得 時rTile shape在維度i大小的倍數(shù)。但是較大的padding kennel會浪費計算，所以Roller將張量padding限制在內(nèi)，并且需要滿足以下公式：。這確保了計算的浪費百分比以 ε 為上界。有了這個限制，我們可以枚舉所有滿足這個條件的有效 rTile 形狀。

• Deriving allrTiles. 鑒于上述對齊要求，對于特定的張量表達(dá)式和硬件設(shè)備，Roller 使用以下接口增量導(dǎo)出所有符合條件的 rTiles：

vector<int> GetNextAlignedAxisSize(rTile T, Dev d),

在給定設(shè)備指定參數(shù)后，它返回rTile shape里每個維度的下一個對齊大小。這是通過在每個維度逐漸增加尺寸大小直到滿足所有對齊要求來計算的。rTile抽象允許Roller被擴(kuò)展以支持新的對齊要求，這是通過GetNextAlignedAxisSize接口來實現(xiàn)的。

• Calculating data reuse score rTile一個有趣的特性是我們可以通過調(diào)整它的shape來隱式的控制內(nèi)存流量。增加rTile 大小通常會以占用更多內(nèi)存為代價為程序帶來更多的數(shù)據(jù)重用機(jī)會。給定一個rTile T和在每一個軸上的下一個對齊大小，我們可以通過

計算出軸的數(shù)據(jù)重用分?jǐn)?shù) ，其中 是通過用GetNextAlignedAxisSize得到的下一個對齊大小替換軸處維度大小獲得的一個更大的rTile。函數(shù)Q(T)和F(T)計算以T的粒度執(zhí)行計算時的內(nèi)存流量和內(nèi)存占用，這可以根據(jù)給定張量表達(dá)式和硬件內(nèi)存規(guī)范直接推斷（0x3.3節(jié)內(nèi)容）。更大的意味著在使用相同的內(nèi)存時可以節(jié)省更多的內(nèi)存流量。內(nèi)存重用分?jǐn)?shù)在構(gòu)建高效的 rProgram（使用 rTiles）中起著至關(guān)重要的作用。

0x3.2 Tensor Program Construction

• rTile program。給定 rTile 和現(xiàn)代加速器的內(nèi)存分層結(jié)構(gòu)，張量計算可以自然地被看成數(shù)據(jù)流處理管道。計算從最低的內(nèi)存級別加載數(shù)據(jù)塊（在rTile中指定），在加速器的執(zhí)行單元上對rTile進(jìn)行計算，并將結(jié)果數(shù)據(jù)塊寫回最低的內(nèi)存級別。對于每個內(nèi)存級別，定義了一個特定的rTile和該內(nèi)存級別的特性保持一致。 因此，Roller將張量計算描述為具有分層 rTile 配置的數(shù)據(jù)處理管道，成為rProgram。

Figure6展示了具有三個存儲層（L0，L1，L2）的設(shè)備上的rProgram。rProgram由每個 內(nèi)存層次 的 rTile 和 rTile 指令(Load, Store, Compute) 來描述。

Figure7(a)展示了Figure7(b)對應(yīng)的MatMul程序。Figure7(c)說明了rProgram如何映射到設(shè)備的每個內(nèi)存層次。具體來說，每次它從內(nèi)存L2中加載一個A的4x4小塊和B的4x8小塊到L1中。然后從L1中加載一個A的2x1和B的1x2小塊到L0(寄存器)中。每次計算完成后，結(jié)果的2x2小塊會直接從L0寫回到L2。

給定一個數(shù)據(jù)處理流水線，對應(yīng)的rProgram的優(yōu)化目標(biāo)就是最大化流水線的吞吐量。這個目標(biāo)可以轉(zhuǎn)化為滿足三個條件：1）計算和內(nèi)存移動應(yīng)該充分利用硬件的特性。2）吞吐量應(yīng)該達(dá)到瓶頸階段（接近峰值）。3）需要有足夠的并行度來利用所有的并行執(zhí)行單元。因此，Roller提出以下rProgram的構(gòu)建策略：首先通過構(gòu)建單核 rProgram在一個內(nèi)核上縱向擴(kuò)展，使得Kernel的硬件利用率飽和。然后通過復(fù)制構(gòu)建的單Kernel橫擴(kuò)以利用硬件的并行度。

• Scaling up an rProgram 。由于rTile的對齊屬性確保了硬件的效率，Roller可以只專注于通過構(gòu)建正確的rTile shape來最大化每個內(nèi)存層次的吞吐量。通過利用0x3.1節(jié)中定義的數(shù)據(jù)重用分?jǐn)?shù)，單核rProgram構(gòu)建算法從初始化rTile開始，并逐漸將其擴(kuò)大到rTile中收益最大的軸（也即具有最大重用分?jǐn)?shù)的）。注意，構(gòu)造算法不需要精確的數(shù)據(jù)重用分?jǐn)?shù)，它只是選擇最大的一個來最大化吞吐量。在此過程中，內(nèi)存的性能會提高直到達(dá)到計算峰值或者最大的內(nèi)存容量。上述過程從上到下對每個內(nèi)存層次進(jìn)行重復(fù)，直到構(gòu)建出所需的rProgram。請注意，如果某些張量表達(dá)式的數(shù)據(jù)重用分?jǐn)?shù)保持不變，比如elemetwise算子，Roller將只為頂層構(gòu)建rTiles并從底層內(nèi)存內(nèi)存加載它們。

Figure8展示了詳細(xì)的構(gòu)建算法。給定一個張量表達(dá)式expr和目標(biāo)設(shè)備dev，該算法在頂層內(nèi)存構(gòu)造一個初始化的rTile T并遞歸的放大T（對應(yīng)第4行的EnlargeTile）。每一步，它都會枚舉下一個更大的rTile T‘，最大程度的提高數(shù)據(jù)重用得分（對應(yīng)第10行的GetNextRTileShapes）。如果T'達(dá)到內(nèi)存容量（第13行）或者數(shù)據(jù)塊加載的吞吐量MemRef(T')超過了峰值計算吞吐量 MaxComputePer f(T')（第17行），算法記錄當(dāng)前的rTile并在下一個內(nèi)存級別繼續(xù)EnlargeTile。否則，它會在當(dāng)前內(nèi)存層級繼續(xù)擴(kuò)大T'（第20行）。構(gòu)建在最低的內(nèi)存層級完成（第6行），產(chǎn)生一個結(jié)果并重復(fù)運行直到產(chǎn)生K個rPrograms（來容忍編譯器的隱藏因素影響），注意，這里的MemPer f(T′)和MaxComputePer f(T′)是基于dev和0x3.3節(jié)的微性能模型推導(dǎo)出來的。

• Scaling out an rProgram。鑒于大多數(shù)DNN算子的計算模式和加速器中的并行執(zhí)行單元的同質(zhì)性，Roller 通過將計算統(tǒng)一劃分為大小等于最低內(nèi)存層級 rTile 的 rTiles，簡單地將在一個執(zhí)行單元上構(gòu)建的 rProgram 復(fù)制到其他單元。我們通過將所有rTiles平均分配到所有執(zhí)行單元來實現(xiàn)這一點。注意，Roller更喜歡將reduce軸分配到同一執(zhí)行單元上，因為它們可以在更高的內(nèi)存層級中共享recue的結(jié)果。請注意，Roller并不假設(shè)會獨占所有的計算單元，系統(tǒng)可以在橫向擴(kuò)展時顯示地控制rProgram的并行度。
• Small operator and irregular tensor shape 。橫向擴(kuò)展算法天然有利于有足夠并行度的算子。例如，分區(qū)數(shù)明顯大于執(zhí)行單元數(shù)的。對于小算子，算法的整體性能kennel會受到并行執(zhí)行單元利用率低的影響。這里可以通過Rammer編譯器的同時調(diào)度一些小Kernel來解決。然后另外一種方法是對于每個rProgram，Roller嘗試沿著具有最小數(shù)據(jù)重用分?jǐn)?shù)的軸收縮rTiles，來實現(xiàn)足夠的并行度。請注意，和其它對齊規(guī)則一樣，此枚舉過程每次都會返回下一個對齊的Tile大小，這是一個高效的過程，和整個構(gòu)建過程相比產(chǎn)生的成本可以忽略。

另外大算子可能包含不規(guī)則的尺寸較小的張量維度，而Roller由于對齊要求kennel無法生成足夠數(shù)量的rProgram。為了解決這個問題，Roller通過軸融合pass將張量表達(dá)式轉(zhuǎn)換為規(guī)范的形式。具體來說，對于所有設(shè)計的張量，如果在一個張量中存在兩個相鄰的軸，這些軸在所有的其它張量中既存在又相鄰，或者都缺失，Roller就可以安全的合并這兩個軸。如，一個輸入和輸出張量形狀都是[17, 11, 3]的張量，Roller會把這三個維度fuse起來變成。除了軸融合外，Roller還嘗試在張量填充機(jī)制中貪心的增加參數(shù)，直到kProgram構(gòu)建完成。

0x3.3 Efficient Evaluation of an rProgram

在構(gòu)建算法中，Roller需要評估rProgram的性能。Roller無需評估真實硬件設(shè)備中端到端的rProgram，只需要評估 rTile 的性能，如Figure8中的MemPerf和MaxComputePerf。

為此，Roller針對硬件抽象層（HAL）中描述的設(shè)備構(gòu)建了一個微觀模型。HAL將加速器建模為具有分層內(nèi)存的多個并行執(zhí)行單元，HAL公開了三個基于rTile的接口：Load，Save，Compute。執(zhí)行單元被抽象為rTile Execution Unit（TEU），它通過Compute接口對數(shù)據(jù)塊進(jìn)行計算?？梢詫⒍鄠€TEUs組織為一個組，它們可以協(xié)同加載和存儲Tiles。HAL將不同的內(nèi)存層（如寄存器，共享內(nèi)存，DRAM）視為一種統(tǒng)一類型，暴露了影響Tile性能的硬件規(guī)范。硬件規(guī)范包括內(nèi)存容量，事務(wù)長度，緩存行大小，Memory Banks數(shù)量，可以通過Figure9的getDeviceSpec獲取。

• Micro performance model 。借助硬件抽象層，Roller可以輕松推導(dǎo)出rTile（和rProgram）的性能。首先，給定一個rTile，可以從rTile的張量表達(dá)式expr和shape（Figure9中的MemFootprint 和 MemTraffic 接口）靜態(tài)推斷出產(chǎn)生的內(nèi)存占用（包括padding）和跨不同層的內(nèi)存流量。計算數(shù)據(jù)重用分?jǐn)?shù)并檢查rTile是否已經(jīng)超出內(nèi)存容量。其次，為了計算rTile的MaxComputePerf，Roller通過積極擴(kuò)大Tiles shape使得TEU飽和，進(jìn)行一次性分析以測量峰值計算吞吐量。此性能數(shù)據(jù)緩存在Roller中，供將來在構(gòu)造算法中查詢。最后，對于給定的rTile，Roller還估計MemPerf，即從內(nèi)存低層加載到更高層的性能。給定rTile中對齊的內(nèi)存訪問，加載常規(guī)Tile的延遲可以簡單地通過將總流量處于內(nèi)存帶寬來建模。對于所有TEU共享的內(nèi)存層，我們平均分配帶寬。對于較小的訪問內(nèi)存，Roller對每種設(shè)備類型進(jìn)行一次離線分析并緩存結(jié)果。值得注意的是，微觀性能模型只需要在Tile shape完全對齊的情況下準(zhǔn)備，這是Roller的關(guān)鍵要求。

4. 實現(xiàn)細(xì)節(jié)

• 代碼生成：給定固定的代碼結(jié)構(gòu)（如Figure6中的一個rProgram），Roller基于預(yù)定義的模板生成代碼（TVM 內(nèi)置調(diào)度原語）。在每個內(nèi)存層級加載和存儲數(shù)據(jù)塊由 TVM 的 cache_read 和 cache_write 原語實現(xiàn)。rTile 上的分區(qū)是通過 split 和 fuse 完成的。一些rTile計算的原語是通過TVM內(nèi)置API完成的?；谀０?，給定的rProgram可以直接生成cuda代碼。
• Tensor Padding：Roller依靠張量padding將rTiles和張量shape對齊。在實踐中，最底層內(nèi)存（例如 DRAM）中的大多數(shù)張量是由外部程序（例如 DNN 框架）分配的，因此我們只需在上層內(nèi)存（例如共享內(nèi)存）中應(yīng)用padding。Roller的張量padding目前需要輸入張量表達(dá)式來指定它是否允許填充，以及默認(rèn)的填充值（例如，0 表示 MatMul 運算符）。對于 Memory Bank 對齊的storage padding，我們利用 TVM 的 storage_align 原語添加padding。
• Performance profiling。Roller實現(xiàn)了兩個性能分析器。一個微觀性能分析器和一個內(nèi)核分析器。前者通過micro-benchmark生成內(nèi)存帶寬，計算吞吐量等硬件指標(biāo)。這是針對每種設(shè)備類型和張量表達(dá)式的一次離線分析。后者描述了top K個kPrograms中最快的kernel，如果k大于1則用于每一個編譯結(jié)果。在實際應(yīng)用中，特定內(nèi)核代碼的性能也會受到設(shè)備編譯器和硬件相關(guān)隱藏因素的輕微影響，Roller 幾乎無法控制。這些因素包括不同指令類型的指令密度、寄存器分配行為、設(shè)備編譯器優(yōu)化、warp 調(diào)度開銷等。特別是在 NVIDIA GPU 上，Roller 依靠 nvcc 將生成的 CUDA 代碼編譯成機(jī)器代碼。但是，nvcc 的專有優(yōu)化可能會對程序執(zhí)行行為產(chǎn)生不利影響。因此，Roller 利用內(nèi)核分析器快速評估性能最佳的 rProgram 并選擇最佳的。較大的 K 通?？梢蕴岣遦ernel質(zhì)量。在評估前 10、20 和 50 個結(jié)果后，我們的經(jīng)驗表明，前 10 名可以獲得大多數(shù)情況下的最佳結(jié)果。請注意，Roller 的內(nèi)核分析器不同于以前編譯器中由機(jī)器學(xué)習(xí)算法驅(qū)動的評估過程 ?；?ML 的方法通常需要數(shù)百甚至數(shù)千個順序評估步驟，而 ROLLER 僅并行分析數(shù)十個候選者。未來，我們計劃實現(xiàn)匯編級代碼生成，以緩解高級設(shè)備編譯器中的隱藏問題。

還有一些NIVIDIA GPU/AMD ROCm/Grphcore IPUs具體硬件上的一些實現(xiàn)細(xì)節(jié)，感興趣的可以自己看下論文。

5. 評測

這里主要看一下在cuda上的結(jié)果。

Figure 10 繪制了我們基準(zhǔn)測試中 119 個算子的平均kernel性能，按算子類型和 ID 排序。我們將大型算子（例如，kernel時間大于 5ms）繪制在 y 軸為對數(shù)尺度的頂部子圖中，而底部 4 個子圖是其它中小型算子。首先，與 CUDA 庫 (CudaLib) 相比，Roller 可以為 81.5% 占比的算子獲得可比的性能（即在 10% 以內(nèi)），并且對于 59.7% 的算子來說甚至更快。我們觀察到，Roller 表現(xiàn)較差的大多數(shù)算子是具有 3×3 或更大濾波器的卷積算子，它們通常在 cuDNN 中使用更有效的數(shù)值算法（例如，Winograd [23]）來實現(xiàn)，并且難以用張量表示表達(dá)。這就是在這些情況下 Ansor 和 TVM 也比 CudaLib 慢的原因。其次，與 TVM 和 Ansor 相比，Roller 也可以分別為 72.3% 和 80.7% 占比的算子獲得可比的性能。其余的 27.7% 和 19.3% 主要是小算子或張量形狀不規(guī)則，難以與硬件對齊。然而，這些算子的kernel執(zhí)行時間通常相對較短，例如平均只有 1.65 毫秒和 1.16 毫秒。在所有算子的 54.6% 和 65.5% 占比中，Roller 甚至可以分別比 TVM 和 Ansor 生成更快的kernel。我們觀察到這些算子中的大多數(shù)都是大型且耗時的。正如上面的子圖所示，當(dāng)算子大于 5 毫秒（最高 343 毫秒）時，Roller 可以為這些算子中的大多數(shù)實現(xiàn)更好的性能，例如，與 TVM 和 Ansor 相比，平均速度提高了 1.85 倍和 1.27 倍。

下面的Figure11還比較了算子編譯的平均時間：

可以看到相比于TVM和Ansor，Roller的算子編譯時間在數(shù)秒內(nèi)，比TVM和Ansor的搜索時間快了2個數(shù)量級。

這里的Table展示了幾個經(jīng)典的神經(jīng)網(wǎng)絡(luò)的性能和編譯時間，可以發(fā)現(xiàn)Rooller相比于TVM和Ansor可以獲得相當(dāng)?shù)男阅埽梢詫⒕幾g時間從幾十個小時縮短到幾百秒鐘，可以大大提高模型的實際生產(chǎn)周期。

6. 結(jié)論&評價

為了解決編譯時間長的問題，這篇論文提出了Roller，它的核心是rTile，這是一種新的tile抽象，它封裝了和底層加速器的關(guān)鍵特性一致的張量shape，從而通過限制shape的選擇實現(xiàn)高效的運行。Roller采用了基于rTile的遞歸構(gòu)建算法來生成目標(biāo)程序（rProgram）。最終，Roller可以在幾秒內(nèi)就生產(chǎn)高效的Kernel，性能可以超越目前主流加速器上的其它張量編譯器，并且為IPU等新的加速器生產(chǎn)更好的Kernel。