面向權(quán)重可變的存算一體加速器的卷積算子調(diào)度

師緊想

（上海交通大學(xué)電子信息與電氣工程學(xué)院，上海200240）

0 引言

深度卷積神經(jīng)網(wǎng)絡(luò)迎來了巨大的發(fā)展，傳統(tǒng)CPU、GPU越來越難以滿足大規(guī)模神經(jīng)網(wǎng)絡(luò)對(duì)功耗和能效的需要。隨著Dennard定律縮放的失效[1]，存算一體加速器以低功耗、高能效的計(jì)算特點(diǎn)受到了廣泛關(guān)注。以憶阻存儲(chǔ)器（Resistive Random Access Memory，ReRAM）為基礎(chǔ)的存算一體加速器是存算一體一個(gè)重要研究方向[11-12]。

存算一體加速器具有存儲(chǔ)單元與計(jì)算單元一體的特點(diǎn)，因此給面向存算一體的編程和資源調(diào)度帶來了困難，一些存算一體加速器[2-3]仍采用手工映射的方式部署神經(jīng)網(wǎng)絡(luò)模型。本文面向憶阻器存算一體設(shè)計(jì)了指令生成工具，并設(shè)計(jì)了算子調(diào)度策略，解決面向存算一體神經(jīng)網(wǎng)絡(luò)加速器的編程、編譯問題。

1 憶阻器存算一體架構(gòu)

憶阻器存算一體加速器的基本單元是憶阻器交叉陣列（crossbar），當(dāng)前的存算一體加速器普遍采用多個(gè)憶阻器交叉陣列組成計(jì)算點(diǎn)陣，并將訓(xùn)練后的卷積神經(jīng)網(wǎng)絡(luò)的參數(shù)（權(quán)重）一次性映射在陣列上。例如ISAAC[2]采用了12k個(gè)crossbar組成一個(gè)從chip到tile到crossbar的層次化計(jì)算架構(gòu)，每個(gè)crossbar包含128行×128列個(gè)憶阻器單元，可以映射48Mb模型參數(shù)，足以映射如ResNet-18網(wǎng)絡(luò)的所有參數(shù)，而不需要考慮crossbar再次映射的情況（后文稱再次映射為權(quán)重更新）。

而實(shí)際上受限于工藝條件，crossbar很難做到128行×128列的規(guī)模[4,5,6]，并且從面積和利用率等實(shí)際情況考慮：①采用上萬個(gè)crossbar占用了很大的面積，增加了功耗；②在映射部分參數(shù)規(guī)模不大的網(wǎng)絡(luò)時(shí)，硬件資源利用率不高；③存在神經(jīng)網(wǎng)絡(luò)參數(shù)規(guī)?？蛇_(dá)到億級(jí)，對(duì)于這些極大規(guī)模的網(wǎng)絡(luò)，該架構(gòu)仍然無法一次性映射全部權(quán)重?？紤]這些情況，本文抽象出來存算一體架構(gòu)的一些共有特征，提出一個(gè)陣列規(guī)模小的權(quán)重可變的存算一體架構(gòu)，通過多次映射權(quán)重來完成整個(gè)網(wǎng)絡(luò)的計(jì)算，并且通過本文設(shè)計(jì)的軟件棧完成具有多次權(quán)重更新的編譯需求。

本文面向的架構(gòu)如圖1所示，其中左側(cè)是該架構(gòu)的tile級(jí)結(jié)構(gòu)，包含一個(gè)全局共享存儲(chǔ)器和4個(gè)core。通過高帶寬的H-tree互聯(lián)結(jié)構(gòu)，core可以訪問tile上的全局共享內(nèi)存，每個(gè)core包含9個(gè)邏輯crossbar，每個(gè)crossbar由64×64個(gè)2bits憶阻器單元構(gòu)成，即本文所基于的架構(gòu)可同時(shí)映射144Kb權(quán)重，最小矩陣向量乘執(zhí)行單元為64×64×8bits。由于架構(gòu)需要多次更新crossbar上的權(quán)重，考慮到crossbar更新權(quán)重的代價(jià)較高，本文將在后文設(shè)計(jì)相關(guān)的調(diào)度策略緩解這個(gè)問題。

存算一體加速器具備軟件可編程性，如圖1右側(cè)所示，本文架構(gòu)的Core單元可以執(zhí)行從/向tile的全局存儲(chǔ)器上讀取/寫入數(shù)據(jù)的過程，并可以獨(dú)立譯碼和執(zhí)行計(jì)算過程。

圖1 存算一體層次化架構(gòu)

2 指令與代碼生成

2.1 流程介紹

當(dāng)前存算一體加速器上部署神經(jīng)網(wǎng)絡(luò)模型時(shí)大多采用手工映射的方法，為了減少開發(fā)時(shí)間，一些工作[7][8]設(shè)計(jì)了編譯軟件棧，但是缺乏算子優(yōu)化且沒有考慮權(quán)重更新的情況。Halide[9]是一個(gè)解耦計(jì)算與調(diào)度的代碼生成工具，本文的指令生成工具在算子優(yōu)化方面結(jié)合Halide的思想，提供了優(yōu)化算子的能力，并且考慮了crossbar的權(quán)重更新。

本文指令生成工具的代碼生成流程如圖2所示。

圖2 編譯流程圖示

首先利用Halide描述運(yùn)算，再將運(yùn)算描述轉(zhuǎn)換成Halide算子中間表示，并將循環(huán)嵌套表示的算子進(jìn)行任務(wù)劃分、分割成與硬件原語綁定的執(zhí)行單元，以方便后續(xù)代碼生成。最后將算子級(jí)中間轉(zhuǎn)換成LLVM IR，將LLVM IR翻譯成可執(zhí)行文件后，運(yùn)行可執(zhí)行文件，可以生成存算一體指令，最后借助驅(qū)動(dòng)工具將指令加載到Core單元并執(zhí)行。

2.2 算子任務(wù)劃分與并行

考慮到加速器執(zhí)行的最小運(yùn)算粒度是矩陣向量乘操作，因此需要將以循環(huán)嵌套描述的卷積等運(yùn)算拆分成硬件可執(zhí)行的矩陣向量乘操作單元。另外為了讓算子在多個(gè)Core上并行執(zhí)行，采取將以循環(huán)嵌套描述的算子的一些計(jì)算軸與計(jì)算核進(jìn)行綁定的方法?？紤]減少計(jì)算依賴等因此，本文將卷積的輸出通道與Core進(jìn)行綁定。

如圖3所示，本文首先在卷積算子的7層循環(huán)中將輸出通道C表示的循環(huán)軸分割成兩個(gè)子循環(huán)軸C_o與C_i，分割的因子為64，以便匹配硬件最小執(zhí)行單元，進(jìn)一步以因子4分割C_o通道，以綁定4個(gè)計(jì)算核心。最終形成從外到內(nèi)為C_o_outer、C_o_inner、C_i的循環(huán)順序，并在算子IR上的C_o_inner軸上利用語法標(biāo)記為并行，以便在代碼生成時(shí)將該變量與計(jì)算核心綁定。

3 算子優(yōu)化與調(diào)度空間探索

考慮到存算一體的輸入數(shù)據(jù)訪問以及crossbar權(quán)重更新代價(jià)，本文主要研究了兩個(gè)算子優(yōu)化的技術(shù)：①連續(xù)內(nèi)存訪問優(yōu)化：在片上輸入緩沖區(qū)資源足夠的情況下，讀取輸入數(shù)據(jù)時(shí)，應(yīng)盡量選擇讀取連續(xù)內(nèi)存排列的數(shù)據(jù)，以提高DMA訪問的效率；②權(quán)重更新次數(shù)優(yōu)化：在中間結(jié)果緩沖區(qū)足夠的情況下，盡量復(fù)用當(dāng)前權(quán)重，減少權(quán)重更新次數(shù)。

圖3 卷積運(yùn)算在多核之間的任務(wù)劃分

3.1 優(yōu)化連續(xù)內(nèi)存訪問

考慮圖4所示的兩種讀取輸入的方式，當(dāng)讀取的輸入張量是張量A表示的張量窗口時(shí)，由于第一行數(shù)據(jù)和后行的數(shù)據(jù)在內(nèi)存的存放存在間隔，因此需要使用至少三次內(nèi)存訪問操作（或者啟動(dòng)三次DMA操作）才能將完整輸入讀入到數(shù)據(jù)緩沖區(qū)。

圖4 從內(nèi)存中讀取輸入張量數(shù)據(jù)的兩種訪問形式

當(dāng)輸入緩沖區(qū)還存在余量時(shí)，考慮張量B表示的讀取方式，由于第一行數(shù)據(jù)和第二行數(shù)據(jù)在內(nèi)存中是連續(xù)存放的，因此只需要一次內(nèi)存訪問指令。顯然第二種讀取方式減少讀取次數(shù)，提高了訪問效率。本文通過在編譯時(shí)檢測(cè)讀取的目標(biāo)張量數(shù)據(jù)是否連續(xù)存放，如公式1所示。

并利用存儲(chǔ)感知輸入緩沖區(qū)的余量決定讀取張量的大小，最后決定每次讀取張量塊的大小。

3.2 優(yōu)化權(quán)重更新次數(shù)與調(diào)度空間探索

前文敘述了本文架構(gòu)在映射神經(jīng)網(wǎng)絡(luò)模型時(shí)權(quán)重?zé)o法一次性映射到crossbar上，因此在計(jì)算中需要多次更新權(quán)重。因此在編譯階段本文需要解決兩個(gè)問題，一是在算子循環(huán)嵌套的何處進(jìn)行權(quán)重更新，二是如何調(diào)度算子使得權(quán)重更新次數(shù)減少。對(duì)于第一個(gè)問題，結(jié)合上文中的算子任務(wù)劃分策略，對(duì)于輸入layout為[N,K,H,W]（batch_size、輸入通道數(shù)、長(zhǎng)、寬），卷積核layout為[C,K,m,m]（輸出通道數(shù)、輸入通道數(shù)、卷積核大?。┑木矸e運(yùn)算，可以得到，在卷積運(yùn)算循環(huán)嵌套的輸出通道C_o與輸入通道K_o所表示的循環(huán)變量改變時(shí)，權(quán)重需要進(jìn)行更新。因此本文在C_o和K_o軸設(shè)置了一個(gè)語法標(biāo)記，如圖5所示，這個(gè)標(biāo)記通過調(diào)用一個(gè)接口函數(shù)將Invalid標(biāo)記傳送到代碼生成模塊，代碼生成模塊根據(jù)Invalid標(biāo)記來判斷是否將接下來的LoadWeight指令插入到指令流中。對(duì)于第二個(gè)問題，結(jié)合圖5所示描述，通過Halide提供的循環(huán)重排方法，盡量將h_o、w_o軸排到c_o_outer與k_outer之后，減少了權(quán)重更新的次數(shù)。

本文在算子優(yōu)化上借鑒了Halide解耦計(jì)算描述與計(jì)算調(diào)度的思想。首先對(duì)算子如卷積運(yùn)算設(shè)置一個(gè)默認(rèn)的調(diào)度模板，然后在此模板的基礎(chǔ)上利用reorder、split、tiling等方法對(duì)默認(rèn)的循環(huán)嵌套進(jìn)行循環(huán)重排、分塊等操作，以探索在指定硬件資源上的最佳計(jì)算方案。

為了探索最佳調(diào)度方案，本文首先基于調(diào)度模板和輸入負(fù)載的大小生成調(diào)度空間，在調(diào)度空間中通過遍歷并評(píng)估各個(gè)點(diǎn)以找到最佳調(diào)度方案，并將最佳調(diào)度生成到記錄文件，在遇到相同計(jì)算負(fù)載時(shí)可以直接從記錄文件中提取相應(yīng)的調(diào)度方案，減少不必要的調(diào)度空間探索時(shí)間。

如圖5所示，本文首先將卷積的循環(huán)軸如C軸、K軸進(jìn)行第一步拆分，如公式（2）、（3）所示，其中C_i與K_i用以匹配硬件最小執(zhí)行單元，再進(jìn)一步對(duì)C_o、K_o以及H,W,N軸采取類似的分割策略，可以生成一個(gè)由C_o_outer,C_o_inner,H_o,H_i,W_o,W_i,N_o,N_i組成的一個(gè)與變量順序和幅度有關(guān)的調(diào)度空間。

圖5 卷積循環(huán)分塊與調(diào)度（分塊因子由調(diào)度空間中所選的調(diào)度方案決定）

由于不同的排列組合下的申請(qǐng)buffer空間、中間結(jié)果所占的空間、內(nèi)存訪問次數(shù)以及權(quán)重更新次數(shù)都有所不同，因此可以尋優(yōu)一個(gè)最佳調(diào)度。

4 實(shí)驗(yàn)結(jié)果與分析

4.1 實(shí)驗(yàn)方法

為了驗(yàn)證本文指令生成工具編譯流程的正確性和調(diào)度策略的有效性，本文基于PUMA[10]仿真器進(jìn)行修改以支持本文所提出的存算一體指令，并對(duì)生成的指令塊進(jìn)行仿真驗(yàn)證，結(jié)合生成的各級(jí)中間表示進(jìn)行分析以驗(yàn)證編譯流程的正確性。最后通過與默認(rèn)調(diào)度方案對(duì)比，驗(yàn)證本文優(yōu)化方案的有效性。

4.2 實(shí)驗(yàn)結(jié)果

以batch_size=1,height=28,width=28,in_filter=128,out_filter=256,wkernel=3,padding=[1,1],stride=[1,1]的卷積運(yùn)算為例，本文指令生成工具生成的含有硬件信息的算子級(jí)中間表示如圖6所示，本文在基于PUMA[10]的仿真器上對(duì)生成的指令流進(jìn)行了仿真，驗(yàn)證了生成結(jié)果的正確性。圖7顯示在本文的減少權(quán)重更新策略調(diào)度下，部分神經(jīng)網(wǎng)絡(luò)的權(quán)重更新次數(shù)相比與默認(rèn)模板減少了90%以上。

圖6 生成的Conv2d-ReLU算子描述

圖7 算子優(yōu)化前后的映射權(quán)重的次數(shù)對(duì)比

4.3 結(jié)果分析

圖6中間表示表明了本文指令生成工具生成代碼的正確性，另外之所以本文的調(diào)度方案相比與沒有調(diào)度的方案有很大提升，是因?yàn)槟J(rèn)計(jì)算方案趨向于產(chǎn)生完成計(jì)算結(jié)果，因此產(chǎn)生大量權(quán)重更新，在調(diào)度之后，趨向于在中間結(jié)果緩沖區(qū)有余量的情況下優(yōu)先產(chǎn)生中間結(jié)果，因此產(chǎn)生大量權(quán)重復(fù)用的情景，從而減少了權(quán)重更新次數(shù)。

5 結(jié)語

針對(duì)憶阻器存算一體加速器的代碼生成與優(yōu)化問題，本文設(shè)計(jì)了一套自動(dòng)指令生成的流程，在算子級(jí)中間表示上利用Halide解耦計(jì)算與調(diào)度，在生成的調(diào)度空間上探索映射卷積運(yùn)算的最佳調(diào)度方案，最后相比于默認(rèn)計(jì)算規(guī)則，本文減少了85%~96%的權(quán)重更新次數(shù)，大大減少了權(quán)重更新的開銷。