MACO:基于訪存視角的卷積網(wǎng)絡(luò)自動代碼優(yōu)化①

張曉揚 肖俊敏 姚家樹 譚光明

(*中國科學(xué)院計算技術(shù)研究所高性能計算機(jī)研究中心 北京 100190)

(**中國科學(xué)院大學(xué) 北京 100049)

0 引言

推理優(yōu)化一直是人工智能(artificial intelligence,AI)與系統(tǒng)結(jié)構(gòu)領(lǐng)域交叉的研究重點。隨著智能物聯(lián)網(wǎng)產(chǎn)業(yè)的發(fā)展,越來越多的人工智能算法和智能硬件應(yīng)運而生,如何將海量算法和智能硬件有機(jī)結(jié)合成為系統(tǒng)性能優(yōu)化人員關(guān)注的問題。為實現(xiàn)應(yīng)用快速部署遷移的需求,大量研究開始關(guān)注自動優(yōu)化技術(shù)。同時,相比于硬件計算性能的高速發(fā)展,硬件的內(nèi)存性能提升較為緩慢,這也直接導(dǎo)致了現(xiàn)有的硬件內(nèi)存墻問題越來越嚴(yán)重,性能優(yōu)化目標(biāo)開始逐步從計算優(yōu)化過渡到訪存優(yōu)化。

卷積神經(jīng)網(wǎng)絡(luò)(convolutional neural network,CNN)的自動優(yōu)化技術(shù)通常包含2 個方面:全局計算圖優(yōu)化和算子內(nèi)核優(yōu)化。全局計算圖優(yōu)化是針對確定的計算任務(wù)以及算子間的數(shù)據(jù)依賴影響,確定最合適的算子搭配;算子內(nèi)核優(yōu)化是針對單個算子,確定特定硬件平臺上最合適的算子優(yōu)化參數(shù)。從內(nèi)存訪問的角度看,在全局圖優(yōu)化中存在的一個重要問題就是數(shù)據(jù)布局的確定。由于具體的算子性能與數(shù)據(jù)布局有相關(guān)性,且不同的算子間存在布局依賴性,所以需要從全局角度確定性能最優(yōu)的算子數(shù)據(jù)布局;而在算子內(nèi)核優(yōu)化方面,由于自動優(yōu)化技術(shù)多基于機(jī)器學(xué)習(xí)驅(qū)動的迭代式搜索方式實現(xiàn)[1],在沒有專家經(jīng)驗指導(dǎo)的搜索中會帶來很多無效搜索,因此可以從降低訪存提高數(shù)據(jù)重用的角度給予性能優(yōu)化搜索的方向性指導(dǎo)。

針對以上2 個訪存相關(guān)的自動優(yōu)化問題,目前也開展了一些研究工作。Li 等人[2]最先針對圖形處理器(graphics processing unit,GPU)上的不同數(shù)據(jù)布局的性能展開討論,但其工作主要關(guān)注如何快速進(jìn)行多維數(shù)據(jù)布局的轉(zhuǎn)換,而不關(guān)注應(yīng)當(dāng)何時進(jìn)行數(shù)據(jù)轉(zhuǎn)換。Zhang 等人[3]提出了針對數(shù)字信號處理器(digital signal processor,DSP)進(jìn)行數(shù)據(jù)布局轉(zhuǎn)換時機(jī)決策的模型,但是基于分類預(yù)測模型在一定程度上會造成全局性能的預(yù)測偏差。許浩博等人[4]提出面向多任務(wù)的神經(jīng)網(wǎng)絡(luò)特征挖掘來重用計算以達(dá)到性能提升的架構(gòu)設(shè)計,但這一設(shè)計存在一定局限性,如在單任務(wù)情況下不適用。而Chen 等人[5]提出的張量虛擬機(jī)(tensor virtual machine,TVM)框架,完全利用搜索算法和機(jī)器學(xué)習(xí)模型進(jìn)行全空間的內(nèi)核優(yōu)化搜索。Sabne[6]提出的加速線性代數(shù)(accelerated linear algebra,XLA)框架則充分利用了手工優(yōu)化的經(jīng)驗實現(xiàn)內(nèi)核級別的優(yōu)化。吳林陽等人[7]從數(shù)據(jù)布局、精度等方面提出運算與數(shù)據(jù)協(xié)同的深度學(xué)習(xí)編譯框架。這些框架的自動優(yōu)化技術(shù)能夠在一定程度上達(dá)到人工優(yōu)化的效果,提高模型高性能部署的效率,但專家經(jīng)驗知識提取的優(yōu)化模式限制了自動優(yōu)化技術(shù)的發(fā)揮空間,而基于搜索空間的極大優(yōu)化自由度又帶來了難以接受的優(yōu)化時間開銷,如何兼顧兩者的優(yōu)勢,成為內(nèi)核自動優(yōu)化過程中非常關(guān)鍵的問題。

1 相關(guān)背景

1.1 卷積數(shù)據(jù)布局

卷積在神經(jīng)網(wǎng)絡(luò)中是以張量形式存在的,張量是一個多維的數(shù)組結(jié)構(gòu)。由于內(nèi)存是個二維存儲結(jié)構(gòu),所以對于多維數(shù)組來說,張量會以特定的順序進(jìn)行數(shù)據(jù)排列,維度的排列先后順序會決定數(shù)據(jù)在內(nèi)存中的組織形式,即對應(yīng)不同的數(shù)據(jù)布局。卷積網(wǎng)絡(luò)中的張量通常包含4 個維度,即待處理特征圖數(shù)量N、特征圖通道數(shù)C、特征圖寬度W和高度H。這些維度的排列組合帶來了不同的數(shù)據(jù)布局,如NCHW、NHWC、CHWN 等。如圖1 所示,邏輯視圖展示了一個N=2(上下2 個獨立塊)、C=2(每個獨立塊疊放的2 層)、W=4、H=3 的四維張量邏輯結(jié)構(gòu),NCHW 和CHWN 分別展示了張量邏輯結(jié)構(gòu)在2種數(shù)據(jù)布局下的映射關(guān)系(紋理一一對應(yīng))。

圖1 張量邏輯布局與物理布局關(guān)系示意圖

1.2 內(nèi)核自動優(yōu)化技術(shù)

內(nèi)核自動代碼優(yōu)化技術(shù)是指通過對專家經(jīng)驗的歸納總結(jié)或者通過合理設(shè)計搜索空間,將優(yōu)化手段通過模式匹配或者搜索方法等方式自動應(yīng)用到代碼優(yōu)化中,以提高目標(biāo)代碼的性能,實現(xiàn)代碼與硬件的深度耦合。傳統(tǒng)的代碼優(yōu)化是件勞動量較大的工作,而自動優(yōu)化技術(shù)的出現(xiàn)能夠有效降低人力成本,使代碼的大規(guī)模部署成為可能。自動優(yōu)化技術(shù)是學(xué)術(shù)界和工業(yè)界都一直關(guān)注的問題。對于卷積算法而言,常見的解決方案是對特定卷積算法進(jìn)行模板設(shè)計,預(yù)留一些可優(yōu)化的參數(shù)軸,通過調(diào)優(yōu)的方式在預(yù)留的參數(shù)組合空間內(nèi)來確定合適的參數(shù)組合,如TVM[5],但模板設(shè)計需要一定的專業(yè)知識,這提高了開發(fā)者的優(yōu)化門檻,同時也將優(yōu)化的空間限制到了開發(fā)者的經(jīng)驗認(rèn)知內(nèi)。所以為了進(jìn)一步釋放自動優(yōu)化的空間,一些自動模板生成技術(shù)應(yīng)運而生,如Ansor[8]是通過規(guī)則匹配的方式,根據(jù)算法的簡單描述自動生成多種模板進(jìn)行內(nèi)核搜索,如采取不同的循環(huán)順序、選擇不同的并行深度等等。圖2 所示的是兩類基于搜索的自動優(yōu)化技術(shù)路線的差異。

圖2 基于模板的自動優(yōu)化與自動模板生成技術(shù)

1.3 紅藍(lán)卵石訪存通信模型

紅藍(lán)卵石游戲模型是由Hong 和Kung[9]于1981年提出的一個兩級內(nèi)存訪問模型。這一模型可以對算法訪存的下界進(jìn)行估計,而訪存下界可以用于指導(dǎo)算法的數(shù)據(jù)流設(shè)計,進(jìn)而為算法的內(nèi)核實現(xiàn)提供優(yōu)化方向。

紅藍(lán)卵石游戲模型所設(shè)計的游戲是基于一個描述算法操作過程的有向無環(huán)圖進(jìn)行的。令G(V,E)為有向無環(huán)圖的符號表示,其中V表示算法操作的頂點集合,E表示2 個操作的依賴關(guān)系的邊集。如果G(V,E) 滿足以下4 個屬性,那么滿足這些屬性的劃分則被稱為S-劃分。

屬性1:V被劃分為h個子集(V1,V2,……,Vh),滿足Vi之間互不相交,但它們的并集是V。

屬性2:對于每個Vi,都有一個最多包含S個頂點的支配集Di。Di是V的一個節(jié)點子集,使得從G的輸入到Vi的節(jié)點的任何路徑都包含Di中的一些節(jié)點。

屬性3:Vi的最小集被定義為Vi中沒有任何屬于Vi的頂點的頂點集。

屬性4:V1,V2,……,Vh之間沒有循環(huán)依賴。

下面的定理描述了基于S-劃分模型的通信下界Q:

2 問題與挑戰(zhàn)

2.1 卷積網(wǎng)絡(luò)數(shù)據(jù)布局的全局規(guī)劃問題

數(shù)據(jù)布局反映的是數(shù)據(jù)在內(nèi)存中的排布方式,其差異會影響到數(shù)據(jù)的局部性特征,進(jìn)而帶來不同的訪存讀寫效率,從而影響卷積的計算效率。而這種數(shù)據(jù)局部性特征,由于卷積張量的參數(shù)維度不同,在不同卷積層上的表現(xiàn)是不同的。圖3 所示的是AlexNet[10]對應(yīng)5 層卷積在CUDA(compute unified device architecture)平臺上2 種數(shù)據(jù)布局下的性能表現(xiàn),可以看到其中的性能差別[4]。

圖3 AlexNet 不同卷積層在不同數(shù)據(jù)布局下的性能表現(xiàn)

這些性能差距不僅來源于卷積核本身,也來源于卷積輸入數(shù)據(jù)和輸出數(shù)據(jù)的張量布局。圖4 展示了卷積神經(jīng)網(wǎng)絡(luò)中張量計算的關(guān)系。

在卷積神經(jīng)網(wǎng)絡(luò)中,張量以四維結(jié)構(gòu)呈現(xiàn),雖然計算過程包含了3 個張量共計12 個維度的計算關(guān)系,但從圖中可以發(fā)現(xiàn),并非所有維度都是獨立的變量。輸出通道取決于卷積核的數(shù)量,而權(quán)重張量的通道數(shù)一般等于輸入張量的通道數(shù),輸出張量與輸入張量有如下對應(yīng)關(guān)系:

其中跨步表示2 次卷積計算的跨度大小。由于張量是多維數(shù)據(jù)結(jié)構(gòu),而內(nèi)存是二維數(shù)據(jù)結(jié)構(gòu),所以張量需要從高維轉(zhuǎn)化為低維將數(shù)據(jù)存放到內(nèi)存中,而不同的存放順序決定了數(shù)據(jù)計算時的連續(xù)性。

由于不同層的數(shù)據(jù)局部性表現(xiàn)不同,如果相鄰層的最優(yōu)數(shù)據(jù)布局不一致,則需要損失某些層的性能,或進(jìn)行布局轉(zhuǎn)化,但布局轉(zhuǎn)化又會帶來額外的時間開銷。對于整個網(wǎng)絡(luò)而言,為了避免過多的數(shù)據(jù)布局轉(zhuǎn)換開銷,需要合理安排各個層的數(shù)據(jù)布局,以充分挖掘網(wǎng)絡(luò)的整體性能。目前已有不少研究工作針對GPU 的數(shù)據(jù)布局進(jìn)行研究,如Majeti 等人[11]探討了中央處理器(central processing unit,CPU) +GPU 混合架構(gòu)的布局自動調(diào)整問題。Weber 和Goesele[12]嘗試對CUDA 應(yīng)用程序的數(shù)據(jù)布局和內(nèi)核優(yōu)化配置進(jìn)行自適應(yīng)優(yōu)化。Zheng 等人[13]應(yīng)用數(shù)據(jù)布局優(yōu)化來提高GPU 上長短期記憶網(wǎng)絡(luò)-循環(huán)神經(jīng)網(wǎng)絡(luò)(long-short-term memory recurrent neural network,LSTM RNN)的性能。但如何從全局角度決策網(wǎng)絡(luò)層的數(shù)據(jù)布局,仍然是個有挑戰(zhàn)性的問題。因為轉(zhuǎn)換開銷的存在,全局的數(shù)據(jù)布局搭配不再是簡單的線性疊加問題。而為了能夠合理決策網(wǎng)絡(luò)層的數(shù)據(jù)布局,并且將這一數(shù)據(jù)布局對全局的影響加以刻畫,則需要對相應(yīng)層的性能進(jìn)行預(yù)測,但對于連續(xù)域的時間預(yù)測問題,精準(zhǔn)預(yù)測又較困難。

2.2 算子內(nèi)核的自動優(yōu)化問題

對于卷積神經(jīng)網(wǎng)絡(luò)而言,算子的內(nèi)核優(yōu)化主要針對卷積。卷積是卷積神經(jīng)網(wǎng)絡(luò)中計算量最大的部分,也是訪存行為最復(fù)雜的部分。卷積算子優(yōu)化是根據(jù)計算特點和硬件架構(gòu),對算子計算代碼進(jìn)行針對性改進(jìn)以提升算子運算效率的工作。目前的卷積優(yōu)化多基于人工優(yōu)化,許多大型硬件廠商為了提高相應(yīng)算法在自家平臺上的性能表現(xiàn),也提供了平臺專用的算子內(nèi)核集合,如目前主流的CPU、GPU 硬件廠商如Nvidia、Intel、ARM、AMD 等,都通過提供適合自身硬件的算子庫來支持AI 網(wǎng)絡(luò)模型算子內(nèi)核的加速,同時也有一些硬件平臺會有第三方開發(fā)者進(jìn)行支持,如尹寧[14]借助OpenBLAS 為龍芯平臺構(gòu)建優(yōu)化了深度學(xué)習(xí)算子庫。但由于這些優(yōu)化手段的針對性較強(qiáng),對于新型網(wǎng)絡(luò)或者不常用的卷積會存在一定的優(yōu)化短板,表現(xiàn)出長尾效應(yīng)。圖5 展示了CUDA 的深度神經(jīng)網(wǎng)絡(luò)(CUDA deep neural network,cuDNN)與TVM 調(diào)優(yōu)結(jié)果的性能對比,可以發(fā)現(xiàn)cuDNN[15]庫在部分卷積中性能表現(xiàn)較好,而在另外一些卷積下卻難以發(fā)揮出性能。這主要是因為很多優(yōu)化手段是針對于常見卷積,而這些優(yōu)化具有一定的滯后性,所以在算法和模型日新月異的發(fā)展趨勢下,現(xiàn)有優(yōu)化庫的穩(wěn)定性會導(dǎo)致軟件和硬件之間的差距難以快速被填補(bǔ)。因此,基于平臺和模型的內(nèi)核自動優(yōu)化技術(shù)成為解決這一問題的重要途徑。

圖5 cuDNN 與調(diào)優(yōu)性能的對比

自動優(yōu)化技術(shù)雖然能夠加速填補(bǔ)算法與硬件之間的性能溝壑,擺脫繁重的人力工作,但自動優(yōu)化技術(shù)的應(yīng)用仍舊存在很大的問題,其中之一就是巨大的參數(shù)搜索空間所帶來的巨大優(yōu)化開銷。圖6 展示了AlexNet 在TVM 框架內(nèi)的部分卷積參數(shù)搜索空間,可以看到搜索空間可達(dá)千萬級別。而基于搜索的參數(shù)優(yōu)化方法對運行時對數(shù)據(jù)反饋依賴很大,因此在如此海量的空間內(nèi)獲取一個性能較好的內(nèi)核實現(xiàn)是非常耗時的。由于這一優(yōu)化過程包含搜索、編譯、測試等環(huán)節(jié),在不剪枝搜索空間的情況下,很難通過簡單的設(shè)備疊加來實現(xiàn)搜索效率的提高。而對于如Ansor 這類自動模板生成的優(yōu)化方法而言,TVM 的搜索空間僅僅是其中的一個子集,所以雖然這些方法給了更大的優(yōu)化空間和優(yōu)化機(jī)會,但同時也引入了更大的搜索開銷,為自動優(yōu)化技術(shù)的推廣帶來了障礙。

圖6 AlexNet 卷積在TVM 中的參數(shù)搜索空間

為了既能保留優(yōu)化機(jī)會,同時又可以降低搜索開銷,本文基于卷積核復(fù)雜的內(nèi)存讀寫行為進(jìn)行建模,并將建模結(jié)果應(yīng)用到自動優(yōu)化的過程中,以基于訪存角度給予搜索空間合理的設(shè)計方向,能夠有助于在縮減搜索空間的同時,提高搜索效率,并保證搜索的內(nèi)核性能。

紅藍(lán)卵石游戲模型是進(jìn)行算法訪存下界估計的經(jīng)典模型,利用這一模型能夠確定更合適的搜索空間設(shè)計方向,有效降低不必要的搜索過程。然而,這一模型進(jìn)行多階段算法分析時難度很大。因為一般的紅藍(lán)卵石模型只適用于一個單一的計算過程,而對于卷積這類包含多階段的問題,如果直接應(yīng)用原始方法,會在不同階段轉(zhuǎn)換的過程中帶來冗余的數(shù)據(jù)寫回和讀取操作。因此,在這種情況下,通過簡單的數(shù)學(xué)證明來獲取訪存的通信下界是難以完成的,所以需要針對卷積算法分析發(fā)展更合適的推導(dǎo)方法。

2.3 主要解決的問題

根據(jù)上述分析,本文從全局圖優(yōu)化和算子優(yōu)化角度出發(fā),提出了布局規(guī)劃和內(nèi)核優(yōu)化2 個方面的研究,具體優(yōu)化流程如圖7 所示。

圖7 卷積網(wǎng)絡(luò)的優(yōu)化流程

在內(nèi)核優(yōu)化方面,本文力求從訪存角度減少自動優(yōu)化的搜索開銷,高效地獲得最優(yōu)化的代碼;而在布局規(guī)劃方面,本文旨在通過算法設(shè)計和預(yù)測模型來獲取合理的全局內(nèi)存數(shù)據(jù)布局方案,以將自動優(yōu)化后的內(nèi)核進(jìn)行有效的搭配,實現(xiàn)全局性能上的最優(yōu)。本文的具體研究貢獻(xiàn)如下。

(1)在算子內(nèi)核優(yōu)化層面,從訪存行為特征出發(fā),重新發(fā)展了傳統(tǒng)的紅藍(lán)卵石模型,以適應(yīng)多階段算法的通信下界建模。借助提出的建模方法,對卷積的多階段算法I/O 下界進(jìn)行分析,為卷積內(nèi)核的優(yōu)化方向提供參照標(biāo)準(zhǔn)。在這一結(jié)論的指導(dǎo)下,進(jìn)行合理的數(shù)據(jù)流設(shè)計,進(jìn)而從數(shù)據(jù)讀寫的角度提高卷積算法計算的數(shù)據(jù)局部性,提高數(shù)據(jù)重用率。并將這一數(shù)據(jù)流設(shè)計的條件,應(yīng)用到自動模板生成技術(shù)中,以降低海量搜索空間所帶來的巨大開銷,為自動優(yōu)化技術(shù)的應(yīng)用掃清障礙。

(2)在全局圖優(yōu)化層面,針對全局?jǐn)?shù)據(jù)布局決策需要性能預(yù)測為指導(dǎo)的問題,設(shè)計能夠進(jìn)行性能預(yù)測的數(shù)據(jù)布局分析模型,減少實際測試帶來的優(yōu)化開銷,提出全局?jǐn)?shù)據(jù)布局的規(guī)劃算法,采用混合數(shù)據(jù)布局的方式來充分發(fā)揮卷積網(wǎng)絡(luò)各層的性能。

3 卷積內(nèi)核生成方法優(yōu)化

對于卷積內(nèi)核的自動優(yōu)化,本文旨在進(jìn)行合理的優(yōu)化空間設(shè)計,并基于最先進(jìn)的自動模板生成技術(shù),從訪存角度改進(jìn)模板生成時的優(yōu)化搜索空間。首先需要解決的問題是如何獲得一個科學(xué)的訪存最優(yōu)下界;然后基于訪存最優(yōu)下界進(jìn)行數(shù)據(jù)流的設(shè)計;最后根據(jù)數(shù)據(jù)流合理設(shè)計優(yōu)化搜索空間。

紅藍(lán)卵石模型是建立最優(yōu)訪存下界的經(jīng)典模型,但傳統(tǒng)的紅藍(lán)卵石模型更適合于單階段算法。對于卷積這類多階段復(fù)合算法而言,直接推斷P(S) 是十分困難的,因此需要對這一方法加以改進(jìn)。本文采取對P(S) 的下界進(jìn)行估計的方法,對于一個有向無環(huán)圖G(V,E),這里使用符號|·|來表示任意集合的頂點數(shù)目,例如|V|表示集合V中的頂點數(shù)量。令PS為包含有向無環(huán)圖G(V,E) S-劃分所有可能性的集合,并且每個PS中的元素表示一些G(V,E) 的S-劃分。令

H(S) 表示G(V,E) 任一S-劃分集合數(shù)量的下界。根據(jù)P(S) 的定義,結(jié)合式(1)和(2),Q可表示為

因此,只需估計H(2S) 而非P(2S)。由于H(S)依賴于的值,所以對Vi的細(xì)粒度分析是關(guān)鍵。

如果能找出Vi與G(V,E) 所有子計算的關(guān)系,那么估計Vi的頂點數(shù)量成為可能。在推斷Vi上界之前,首先記Gj(Uj,Ej) 為G(V,E) 一個多階段劃分中任一子圖對應(yīng)的子計算,那么該子計算的輸入頂點必須是Gj-1(Uj-1,Ej-1) 輸出頂點,且Uj內(nèi)部頂點集互不相交。

現(xiàn)在的目標(biāo)是估計|Vi| 的最大值。為了方便表示,記Uj的輸出集合為,Vi的支配集記作Di。此處定義最大頂點生成函數(shù)。對于給定的任意整數(shù)k和第j個子計算,定義的最大頂點生成函數(shù)如下:

其中,| Θ(D)∩U∩Uj| 和|Θ(D)∩U∩分別表示由支配集D在2 個頂點集合U∩Uj和U∩中生成的頂點數(shù)量。顯然,φj和ψj提供了Uj和中的頂點數(shù)量上界估計,這些頂點是由滿足|D∩Uj| +| Θ(D)∩|≤k條件的D生成的。

此外,通過最大生成函數(shù)φj和ψj可獲得強(qiáng)有力的分析方法來估計Vi∩Uj和Vi∩中Di生成的頂點數(shù)?；诖?對任意S-劃分G(V,E) 的|Vi| 有一個上界:

假設(shè)有向無環(huán)圖G(V,E) 描述了一個n階段算法。所有子計算都對應(yīng)于一個有向無環(huán)圖的多階段劃分。給定大小為S的快速內(nèi)存,為了實現(xiàn)整個算法的計算,在快速內(nèi)存和慢速內(nèi)存之間的I/O 操作數(shù)Q滿足:

雖然式(7)和(1)是相似的,但是對多階段算法而言估計T要比得到P更容易。

基于以上的多階段算法分析方法,下面以直接卷積為例建立I/O 下界。

圖8 展示了直接卷積的有向無環(huán)圖G(V,E)。直接卷積包含2 個階段:通過輸入圖片和卷積核的輸入生成大量的乘積項;將乘積項相加形成一個基于求和樹的最終輸出。求和樹是一個樹形結(jié)構(gòu)的有向無環(huán)圖子圖,除了輸入頂點外,該樹的其他頂點的入度最多為2,并且樹的所有輸入都會被累加到一起并且只有一個輸出。在求和過程之后就完成了直接卷積。因此,G(V,E) 的多階段劃分可以被記作G(V,E)=G1(U1,E1) ∪G1(U1,E1),對于直接卷積的有向無環(huán)圖,其頂點總數(shù)為

圖8 直接卷積的有向無環(huán)圖

對于不同的滑動窗口,在一個輸入圖像中每個輸入(元素)的最大重用數(shù)記作R,它的值是R=是跨步的大小。對于直接卷積這一兩階段的算法,可以證明ψ1=φ1,φ1(k1)并且φ2(k2) ≤k2-1 對于任何整數(shù)k1和k2都是可用的,這可以用于估計T(S)。

對于直接卷積,T(S),建立I/O 直接卷積的下界為

借助等式(8)中的通信下界可以建立近似I/O最優(yōu)策略方法。在提出的一般I/O 下界理論中,根據(jù)式(6)可知,I/O 下界結(jié)果的最高階項必定由某個φj決定。以直接卷積為例,φ2決定了I/O 下界的最高階項,也就是φ2指向了涉及I/O 操作最多的過程。這說明需要針對直接卷積的第2 階段進(jìn)行優(yōu)化。φ2指明應(yīng)該用求和樹上最少的輸入得到最多的輸出,即充分利用求和樹的輸入。而求和樹上的輸入來源于第1 階段的乘積頂點,也就是數(shù)據(jù)流設(shè)計應(yīng)該將有效片上內(nèi)存盡量分配給這些乘積頂點。由于直接卷積沿通道方向的數(shù)據(jù)只計算一次,所以不需要在通道方向上分配更多額外空間,只需要讓輸入圖像的寬、高和卷積核數(shù)量維度上能夠有更高的相互復(fù)用率,對應(yīng)于最終輸出就是其寬、高、通道維度。

圖9 展示了一個大小為x×y×z的輸出圖像的子塊(輸出張量深色部分)。為了達(dá)到最小的片外內(nèi)存訪問,本文傾向于令xyz的結(jié)果逼近S/Np,其中Np是處理核心的數(shù)量。為了計算出子塊x×y×z的結(jié)果,需要對應(yīng)于x′ ×y′位置的所有通道維度的輸入以及與部分輸出通道相關(guān)聯(lián)的z個卷積核,如圖9 所示。由于片上內(nèi)存受限而且需要用于存儲盡可能多的輸出結(jié)果,所以在加載圖片輸入和卷積核的時候,需要連續(xù)加載而不是一次性加載到快速內(nèi)存中。在每個階段,輸入圖片x′ ×y′ ×α的一部分和相應(yīng)的z個卷積核權(quán)重Hker×Wker×α被加載到片上內(nèi)存中。由于輸入的第i通道僅能夠被權(quán)重的第i層通道重用,所以為了保證輸出子塊在有限的片上內(nèi)存中能夠盡可能大,令α=1,也就是說數(shù)據(jù)流設(shè)計要求以固定的通道數(shù)加載x′×y′貼片并沿著通道維度的方向滑動。將x′ ×y′的輸入貼片和相應(yīng)的z卷積核權(quán)重加載到片上內(nèi)存后,可以對輸出子塊執(zhí)行部分和的計算。

圖9 直接卷積數(shù)據(jù)流

為了更新整個輸出子塊,需沿通道方向連續(xù)滑動x′×y′輸入塊,并加載相應(yīng)的輸入和權(quán)重,執(zhí)行部分更新。因此,更新每個輸出子塊只需要從片外內(nèi)存加載所需的輸入和權(quán)重到片上內(nèi)存一次。同時,不同子塊由Np處理器并行地處理。在數(shù)據(jù)流設(shè)計中,共有(WoutHoutCout)/(xyz) 個輸出子塊。為更新每個子塊,需從輸入圖像中輸入x′y′Cin個元素,從z個卷積核中輸入WkerHkerCinz個元素。由于R=WkerHker/μ2,x′≈μx,y′≈μy,數(shù)據(jù)讀入I/O 為

其中,當(dāng)且僅當(dāng)xy=Rz時等式成立?；赗=WkerHker/μ2,x′≈μx,y′≈μy,xy=Rz條件的要求使得x′y′=zWkerHker。它決定了每個x′ ×y′貼片的最優(yōu)尺寸。此外,存儲輸出I/O 是HoutWoutCout。當(dāng)xyz≈S/Np且xy=Rz時,I/O 的總量為

如果Np=1 且,則數(shù)據(jù)流能夠達(dá)到I/O 下界。因為S在卷積網(wǎng)絡(luò)的應(yīng)用中一般不大于kB 級別,所以很容易滿足條件。

基于以上數(shù)據(jù)流設(shè)計,本文用科學(xué)的指導(dǎo)方法進(jìn)行搜索空間合理剪枝。根據(jù)數(shù)據(jù)流給出的最優(yōu)性條件,將其應(yīng)用于自動模板生成技術(shù)中。通過將最優(yōu)性條件應(yīng)用到模板調(diào)度設(shè)計的過程中,模板的自動生成被限制到科學(xué)范圍內(nèi),從而降低后續(xù)搜索帶來的時間開銷和硬件開銷,而且也保證了在此空間內(nèi)性能參數(shù)搭配的合理性。具體過程如圖10 所示。

4 全局卷積數(shù)據(jù)布局優(yōu)化

為了挖掘網(wǎng)絡(luò)的全局性能,發(fā)揮出自動調(diào)優(yōu)后的卷積性能,需要確定每層卷積最合理的數(shù)據(jù)布局,因?qū)又埔?所以本文采取混合張量布局策略,允許不同卷積層選擇不同的張量布局。但由于不同的張量布局之間存在轉(zhuǎn)換開銷,因此需要從全局層面來決策整個張量布局的搭配。由于卷積網(wǎng)絡(luò)通常較深,所以其對應(yīng)的搭配空間較大,在自動進(jìn)行布局規(guī)劃的過程中,需要盡可能全面且迅速地判斷規(guī)劃方案的性能,進(jìn)而做出最優(yōu)決策。張量布局的規(guī)劃算法過程如算法1 所示。

可以發(fā)現(xiàn),算法1 可以全面地考慮不同的情況,并及早對不合適的規(guī)劃設(shè)計采取終止搜索行為。而效率方面,在于當(dāng)確定相應(yīng)規(guī)劃策略時所獲得的性能表現(xiàn),也就是解決各層的性能預(yù)測問題。獲得張量布局的性能最自然的方法是對設(shè)備相關(guān)的所有內(nèi)存訪問操作進(jìn)行建模,并與不同張量布局的量化時間進(jìn)行比較。但這種內(nèi)存訪問模型的建模難度大,同時由于計算行為還受到運行時其他因素的影響,所以準(zhǔn)確性難以保證。因此,在實際情況中對所有操作進(jìn)行建模是難以實現(xiàn)的。

為了更加準(zhǔn)確地對不同層的合理數(shù)據(jù)進(jìn)行標(biāo)記,需要采取更加具備泛化能力的方法來獲取張量布局的確定標(biāo)準(zhǔn)。同時,由于是全局優(yōu)化問題,在很多設(shè)備上需要考慮張量布局的轉(zhuǎn)化開銷,因此無法簡單地選擇各層最好的張量布局,這就要求標(biāo)準(zhǔn)的獲得以數(shù)值預(yù)測為好。

采用機(jī)器學(xué)習(xí)模型進(jìn)行建模不失為一種泛化能力較強(qiáng)的方法。但機(jī)器學(xué)習(xí)的方法眾多,且機(jī)器學(xué)習(xí)方法難以在數(shù)值預(yù)測上達(dá)到良好的性能。因此,本文選擇采用神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索的方法進(jìn)行模型的構(gòu)建,以獲得高精度的預(yù)測模型。

本文模型是基于自動優(yōu)化后的內(nèi)核數(shù)據(jù)訓(xùn)練而來的,具體的輸入?yún)?shù)包括數(shù)據(jù)布局的編號、輸入張量維度、權(quán)重張量維度、輸出張量維度、計算跨步等等,以R2 得分為評估目標(biāo),進(jìn)行回歸擬合。在使用時,當(dāng)獲取到相關(guān)的卷積層參數(shù)后可在模型的幫助下,獲得不同數(shù)據(jù)布局編號下的性能預(yù)測結(jié)果。

本文的策略包含預(yù)處理和網(wǎng)絡(luò)優(yōu)化2 個階段。預(yù)處理階段提取卷積網(wǎng)絡(luò)中各層的配置信息,并基于算法1 進(jìn)行布局方案搜索,搜索過程中配合性能預(yù)測方法進(jìn)行決策,將獲得的結(jié)果返回給卷積網(wǎng)絡(luò)進(jìn)行合理的張量布局標(biāo)記。部署卷積網(wǎng)絡(luò)的設(shè)備根據(jù)標(biāo)記將權(quán)重張量提前進(jìn)行轉(zhuǎn)換,并將流式數(shù)據(jù)按照相應(yīng)層提前做好布局轉(zhuǎn)換輸出,如圖11 所示。

圖11 數(shù)據(jù)布局優(yōu)化流程

5 實驗

5.1 數(shù)據(jù)布局模型測試

實驗發(fā)現(xiàn),相比于采用默認(rèn)布局,采取模型決策的布局能夠顯著減少推理用時,ResNet-18、AlexNet、VGG-11 的性能在采取模型決策后分別提升了1.28 倍、1.32 倍、1.29 倍,如圖12 所示。

圖12 默認(rèn)布局與模型決策布局的性能對比

通過對不同算法的測試發(fā)現(xiàn),在相同的訓(xùn)練數(shù)據(jù)集下,支持向量機(jī)(support vector machine,SVM)回歸模型[16]、線性回歸模型[17]、梯度提升回歸樹(gradient boosted regression trees,GBRT) 回歸模型[18]對不同數(shù)據(jù)布局下卷積進(jìn)行性能預(yù)測的R2 得分分別為0.162 584、0.312 771、0.897 487,而借助autoML[19]技術(shù)構(gòu)建的模型R2 得分為0.998 412,如圖13 所示。

眾所周知,R2 得分的取值范圍為[-1,1],取值越接近1 表示模型的預(yù)測表現(xiàn)越接近最優(yōu)[20]。而在實際統(tǒng)計測試中也可以發(fā)現(xiàn),通過神經(jīng)網(wǎng)絡(luò)模型預(yù)測的時間結(jié)果與實際的性能差距并不大。如圖14所示,在與表現(xiàn)最好的GBRT 算法進(jìn)行對比中發(fā)現(xiàn),預(yù)測的閾值(即發(fā)生預(yù)測錯誤時,選定布局與最佳布局的性能差距)較大(即大于0.01 ms)時,本文的方法分布比例僅占2.2%,且最大性能預(yù)測差距不超過0.05 ms,而GBRT 在預(yù)測閾值較大時的分布比例占5.6%,且約2.8%的性能預(yù)測差距超過0.05 ms。這說明GBRT方法的波動性大。由于布局轉(zhuǎn)換開銷的存在,當(dāng)性能預(yù)測的差距較小時對全局的性能影響不大,而在性能預(yù)測差距較大時產(chǎn)生的影響會顯著提高。因為在全局?jǐn)?shù)據(jù)布局預(yù)測時,模型預(yù)測的穩(wěn)定性比較重要,所以應(yīng)當(dāng)盡量避免較大性能預(yù)測差距。

圖14 相應(yīng)預(yù)測閾值下的分布比例

5.2 內(nèi)核性能測試

圖15 展示了自動優(yōu)化卷積相較于cuDNN 的性能提升,可以看出,無論是NCHW 還是NHWC 布局的卷積,其性能較原始相比都有顯著提升,整體加速比為2.24 倍。其中,當(dāng)卷積寬度和高度較小時,在單批次的情況下性能提升更加明顯,隨著卷積寬度和高度的增大,加速效果開始下降。同樣地,當(dāng)卷積的輸出通道數(shù)增加時,性能的提升效果開始下降。這說明隨著卷積的增大,設(shè)備更容易發(fā)揮出計算性能,原始的優(yōu)化難度降低。

圖15 自動優(yōu)化卷積與cuDNN 卷積在NCHW 和NHWC 布局下相比的性能提升

同時,由于搜索空間在經(jīng)過數(shù)據(jù)流重新設(shè)計后獲得了剪枝,避免了大量無效的搜索過程,因此本文的方法相比于原始的搜索空間,搜索效率得到了提升。尤其在搜索的初始階段,性能提升較快,說明本文的方法搜索起點更好。隨著搜索過程的進(jìn)行,本文的方法能夠更快地找到一個性能更好的內(nèi)核優(yōu)化方案,如圖16 所示。

圖16 搜索空間剪枝前后的搜索效果對比

另外,圖17 展示了其中一些卷積的L1 Cache命中率測試結(jié)果。觀察發(fā)現(xiàn),經(jīng)過調(diào)優(yōu)后的內(nèi)核L1 Cache 命中率要顯著高于cuDNN 卷積內(nèi)核,這從側(cè)面印證了調(diào)優(yōu)后的內(nèi)核在數(shù)據(jù)重用方面的表現(xiàn)要優(yōu)于cuDNN。

圖17 調(diào)優(yōu)內(nèi)核與cuDNN 內(nèi)核的L1 Cache 命中比

5.3 端到端性能測試

圖18 展示了端到端的性能表現(xiàn)。當(dāng)采取模型決策的混合數(shù)據(jù)布局后,獲得的網(wǎng)絡(luò)模型性能顯著高于基準(zhǔn),ResNet-18、AlexNet、VGG-11 的性能分別提升了1.07 倍、1.62 倍、1.34 倍,說明這些自動優(yōu)化方法在端到端的性能提升上發(fā)揮了作用。

圖18 端到端性能測試

6 結(jié)論

本文改進(jìn)了經(jīng)典的紅藍(lán)卵石訪存模型的建模方法,通過提出新的訪存下界估計方法,降低了多階段復(fù)合算法的建模難度,使卷積I/O 下界的推導(dǎo)成為可能。獲得卷積I/O 下界后,通過對卷積數(shù)據(jù)流的合理設(shè)計,搜索空間得到了優(yōu)化剪枝,并將其應(yīng)用于最新的自動模板生成方法的搜索空間設(shè)計中,通過避免大量無效搜索來降低搜索時間的成本開銷。而對于全局的數(shù)據(jù)布局確定問題,通過神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索技術(shù)設(shè)計的新型神經(jīng)網(wǎng)絡(luò),配合數(shù)據(jù)布局回溯算法,能夠?qū)崿F(xiàn)對卷積數(shù)據(jù)布局的預(yù)測和搭配。實驗證明,本文所使用的神經(jīng)網(wǎng)絡(luò)能夠有效預(yù)測出不同卷積的性能差異,并提供一個較為合理的數(shù)據(jù)布局搭配方法,較默認(rèn)的數(shù)據(jù)布局在多個網(wǎng)絡(luò)上性能均有1.3 倍左右的提升。同時,經(jīng)過對紅藍(lán)卵石模型的改進(jìn),本文的方法成功獲得了多階段算法的訪存下界,并且經(jīng)過訪存分析后重新設(shè)計的數(shù)據(jù)流,能夠有效降低自動模板生成技術(shù)下的內(nèi)核搜索空間,降低搜索成本。在此基礎(chǔ)上,本文獲得的內(nèi)核性能得到了保證,其較cuDNN 有2.24 倍的加速,同時搜索效率獲得有效提升。

未來的工作主要包括兩方面:一方面是全局方面引入更多的優(yōu)化維度,除卷積布局外,還可在圖結(jié)構(gòu)、算法等方面進(jìn)行性能預(yù)測,以帶來更大的全局優(yōu)化空間;另一方面是局部優(yōu)化上引入對算子融合等情況下的空間設(shè)計以及更加簡單的約束域推導(dǎo)方式。