# FlexSA: Flexible Systolic Array Architecture for Efficient Pruned DNN Model Training <body> <div style="text-align: right;"> <p>2024.03.20</p> <a href="https://arxiv.org/abs/2004.13027">Paper link</a> </div> --- [TOC] --- ## Background ### Systolic array ##### &emsp; &emsp;&emsp;![image](https://hackmd.io/_uploads/ryEo6WTj6.png) ##### &emsp; &emsp;&emsp;上圖為SA的經典示意圖，因為往往access memory都比實際計算時間還要長，所以下面那個就是一個1D SA的範例，一次將同一個data隨時間傳遞給一整排的PE ##### &emsp; &emsp;&emsp;SA的優點是結構簡單且成本低，但缺點是靈活性很差，而且資料需要適當安排，下面以一個2D 3*3 SA為例 ##### &emsp; &emsp;&emsp;![image](https://hackmd.io/_uploads/r1RDgf6i6.png) ## Why FlexSA? 眾所皆知的是SA的彈性非常差，要是所需要的矩陣沒有很剛好的被分割成它的大小，就會造成許多浪費，以下圖為例 ![Screenshot from 2024-03-17 22-13-18](https://hackmd.io/_uploads/r1rW2d40p.png) 在（b)情況，由於pruning導致現在所需要的形狀無法剛好與SA匹配，因此有不少的浪費。這種情況在SA越大的情況會更加嚴重。 因此就會想到如果將一個大塊的SA切成若干塊的小SA，這樣對於每一塊的SA確實就不容易有大量浪費，就如同（c)情況。 但這樣又會衍生另外一個問題，越小的SA就越無法善用SA最大的優點：data reuse，因此會需要很大的Global Buffer，而且面積也會提高。 因此本篇提出FlexSA，顧名思義就是 有彈性的SA。 設計好幾種模式，利用compiler去規劃現在該使用哪種模式才能讓PE達到最高的利用率。 ## Evaluation SA size and Global Buffer Layout 本篇事前主要做了2種分析，分別是 1. 在128 * 128的PE大小上，該拆分成幾核? ![image](https://hackmd.io/_uploads/ryloP7rCT.png) 藍虛線是指高pruning強度下的PE利用率 藍實線是指低pruning強度下的PE利用率 藍虛線是指高pruning強度下的從GBUF到LBUF的data傳輸量 藍實線是指低pruning強度下的從GBUF到LBUF的data傳輸量 很明顯的，拆分成越多核，肯定會有更高的PE使用率，但data傳輸量會大幅增加 因此看起來4核的64 * 64SA是不錯的方案 2. 不同Global Buffer的設計的優缺點，還有面積差異 ![image](https://hackmd.io/_uploads/HyLcPmBA6.png) 上圖是講2種不同的Global Buffer的設計 4.a是多核共享一個Global Buffer 4.b是每個核有自己的Global Buffer ![image](https://hackmd.io/_uploads/HJhjwXrA6.png) 上圖紅線指增加的data傳輸通路的面積 上圖藍線指拆分Global Buffer與Local Buffer增加的額外邏輯的面積 因此結合兩張圖可以看出，若是使用4.a的方案，也就是共享Global Buffer的情況不僅僅會增加data傳輸，增加的面積也過多。但若是使用4.b的方案，也就是每個核都有自己的Global Buffer的情況，則會有額外的邏輯面積，尤其是核數過多的情況。 因此看起來4核的64 * 64SA是不錯的方案，但Buffer的設計仍須被優化。 ## FlexSA ![image](https://hackmd.io/_uploads/ByrpvXHRT.png) 主要想法就是透過compiler去規劃，在不浪費的前提下，使用的SA越大越好，但若是會有過多的浪費，就根據情況來切割SA，去盡可能地減少浪費。 如上圖所示，FlexSA基於四個共享Global Buffer的SA核心，但每個核心獨立運行，透過不同的operation mode去決定現在各個核心要獨立運作還是協作，紅線即為額外的data path和path switch。 ### 4 modes ![image](https://hackmd.io/_uploads/HJnRvQrAT.png) 優先順序是 FW > VSW = HSW > ISW #### 1. Full Wave (FW) 4個核心同時協作，一開始由上而下每一個核頂部的Local Buffer各自傳入輸入，由於高度64相較於原本的128還小，所以可以節省一半的input shift。 接著Core0和Core2左邊的Local Buffer向右傳輸入給Core1和Core3。 最後由上而下累加，Core0和Core1的輸出也會傳給Core2和Core3，然後由Core2和Core3傳出到Output Buffer。 由於Data Reuse與原本單核的128 * 128相同，所以與單獨4核的64 * 64相比，flow減少一半。 #### 2. Vertical sub-wave (VSW) 顧名思義就是垂直的Core一組，也就是Core0與Core2協作，Core1與Core3協作，目的主要是要處理瘦長的GEMM碎片。 一樣先由4核各自的Local Buffer由上到下傳入，接著Buffer0負責Core0和Core2的傳輸，Buffer1負責Core1和Core3的傳輸。 最後一樣由上而下累加，Core0和Core1的輸出也會傳給Core2和Core3，然後由Core2和Core3傳出到Output Buffer。 #### 3. Horizontal sub-wave (HSW) 顧名思義就是水平的Core一組，也就是Core0與Core1協作，Core2與Core3協作，目的主要是要處理胖短的GEMM碎片。 一樣先由4核各自的Local Buffer由上到下傳入。 接著Core0和Core2左邊的Local Buffer向右傳輸入給Core1和Core3。 最後不同的是Core0和Core1的輸出直接傳到Core2和Core3下方的Output Buffer。 #### 4. Independent sub-wave (ISW) 顧名思義就是4個Core各自獨立。 一樣先由4核各自的Local Buffer由上到下傳入。 接著結合VSW與HSW的傳輸路徑去進行計算。 ## Tiling ![image](https://hackmd.io/_uploads/rJpkOmSAp.png) 想法也很簡單，由於FW mode優先級最高，所以優先切成128 * 128，剩下的就是胖矮的給HSW，瘦高的給VSW，太小的給ISW。 註:這裡的blkN和blkK為FlexSA核心的寬與高，本篇為128。 blkM的大小是非定點輸入，是圖7左側的Local Buffer/128 ## ISA 1. FlexSA mode：類型、每個波的blkM、blkN和blkK的大小。為了映射到wave所需要的micro-controller， e.g.data path選擇、input shift、output shift、accumulation。 2. Vector Load and Store：為了實現Global Buffer到Local Buffer之間的傳輸。 ## Compilation for GEMM Processing ![image](https://hackmd.io/_uploads/SyRldmrRa.png) 2-3行去以128為單位切割矩陣的寬 4-6行代表著Buffer所能存的高度，IsWideWave紀錄是不是寬波 7-9行代表以128為單位切割矩陣的高，IsTallWave紀錄是不是高波 10-11行由Wide和Tall來判斷mode 12-13行LdLBUF_V指令做Vertical方向的輸入 14-15行ShiftV指令的值與wave的高度有關 16-17行LdLBUF_H指令做Horizon方向的輸入 18-19行ExecGEMM指令收到mode和每個波的blkM、blkN和blkK 21-22行在迭代完k維度的迴圈之後，使用StLBUF指令將psum存儲到Global Buffer中 ## Evaluation Method ![image](https://hackmd.io/_uploads/SypCdorRa.png) ![未命名绘图 (3)](https://hackmd.io/_uploads/ByyZxTHR6.jpg) 開發了一個Instruction-level simulator，用於評估FlexSA的性能和能源效率。 | Model | ResNet50、Inception v4、MobileNet v2 | | -------- | -------- | | Pruning Strength | ResNet50在訓練期間使用PruneTrain作為剪枝機制進行剪枝，有兩種不同的模型剪枝強度、Inception v4通過應用ResNet50的相同剪枝統計信息來進行人工剪枝 | | Batch size | 32(ResNet50、Inception v4) 、 128(MobileNet v2) | | Method | 使用mixed-precision MAC來執行conv和FC Layer的GEMMs | | SA core | 700MHz (23TFLPOS in every layout) | | Global Buffer | 10MB | | HBM2 DRAM | 270 GB/s | | I/O | 支援double buffer | 小技巧：可以將共享的輸入保存起來，避免data傳輸 ## Evaluation Result ![image](https://hackmd.io/_uploads/Bk9MdmHAa.png) 圖 10a 說明使用不同核心配置在訓練 ResNet50、Inception v4 和 MobileNet v2 模型在訓練時的平均 PE 利用率。實驗使用無限（理想）的 DRAM bandwidth，以將 PE 利用率的影響僅限於wave和core。 圖 10b 才加了memory bandwidth的限制。 ![image](https://hackmd.io/_uploads/rytXd7SAT.png) 圖 11 將不同model的conv和 FC Layer的 Global Buffer 到 Local Buffer 的芯片內流量進行比較 ![image](https://hackmd.io/_uploads/B1oE_QHA6.png) 圖 12a 評估了使用不同core layout的每個訓練迭代消耗的dynamic power。每個顏色表示不同加速器operation消耗的能量分佈 註：COMP 表示 FlexSA core的Mix-precision MAC operation消耗的能量 OverCore 表示over-core數據傳輸能量 ![image](https://hackmd.io/_uploads/SJHrdmrAa.png) 上圖顯示不同的model在各種不同operation mode的佔比 註：文章說上面的估計都偏保守，用SIMD可能會更快，但這個我還不懂QQ ## Conclusion FlexSA compiler支持使用mode選擇和ISA來實現FlexSA PE的最佳利用。總的來說，與使用單個大核相比，FlexSA 將pruning同時訓練的 CNN 模型的 PE 利用率提高了 37%，並將芯片上的reuse增加了 1.7 倍，節省了 28% 的dynamic power。 除此之外，多個FlexSA的組合，看起來也是相當有潛力的。 --- ## Reference [Systolic Array](https://zhuanlan.zhihu.com/p/26522315) [【機器學習2021】神經網路壓縮(Network Compression) (一) ](https://www.google.com/search?q=pruning%E6%A9%9F%E5%99%A8%E5%AD%B8%E7%BF%92&sca_esv=f115bf2ae4b2ba76&sca_upv=1&sxsrf=ACQVn08I8o64jbdkKE5XtE_ZrsHfCs-ADQ%3A1710667950603&ei=rrj2ZZK8JMa3vr0Pqoif0Ac&oq=pruning&gs_lp=Egxnd3Mtd2l6LXNlcnAiB3BydW5pbmcqAggCMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADMgoQABhHGNYEGLADSKgJUABYAHABeAGQAQCYAQCgAQCqAQC4AQHIAQCYAgGgAgSYAwCIBgGQBgqSBwExoAcA&sclient=gws-wiz-serp#fpstate=ive&ip=1&vld=cid:d0ca49ee,vid:utk3EnAUh-g,st:0) </body>