一種運(yùn)算和數(shù)據(jù)協(xié)同優(yōu)化的深度學(xué)習(xí)編譯框架①

吳林陽 杜偉健 陳小兵 莊毅敏

(*中國科學(xué)院計(jì)算技術(shù)研究所計(jì)算機(jī)體系結(jié)構(gòu)國家重點(diǎn)實(shí)驗(yàn)室 北京 100190) (**中國科學(xué)院大學(xué) 北京 100049) (***上海寒武紀(jì)信息科技有限公司 上海 201308)

0 引 言

隨著大數(shù)據(jù)時代到來，深度學(xué)習(xí)算法已被廣泛應(yīng)用于安防、教育、醫(yī)療、貿(mào)易、交通、家居等各個領(lǐng)域。然而，深度學(xué)習(xí)應(yīng)用對硬件計(jì)算能力要求極高，通用處理器已經(jīng)無法滿足應(yīng)用的實(shí)時性需求。在此背景下，深度學(xué)習(xí)處理器蓬勃發(fā)展，逐漸成為處理深度學(xué)習(xí)應(yīng)用的主要載體。

深度學(xué)習(xí)算法和硬件的快速發(fā)展給編譯器設(shè)計(jì)帶來了巨大的挑戰(zhàn)，主要體現(xiàn)在2個方面：

(1) 通用性。如何設(shè)計(jì)統(tǒng)一的編譯框架，支持不同的硬件平臺后端(GPU、TPU、CPU、加速器)。

(2) 性能。針對特定的硬件平臺后端，如何把性能(延時、吞吐率)優(yōu)化到最好。

目前已有一些相關(guān)工作針對以上2個問題提出了解決方案。比如，Chen等人[1]提出了一種自動化、端到端的編譯優(yōu)化框架TVM。為了優(yōu)化程序的性能，TVM設(shè)計(jì)了面向算法的高級優(yōu)化和面向硬件的低級優(yōu)化兩層優(yōu)化架構(gòu)，并采用了操作融合、訪存延遲隱藏等關(guān)鍵技術(shù)。TVM通過中間表示來支持不同的硬件后端，同時在語言層面支持用戶自定義硬件優(yōu)化原語，從而解決通用性問題。又比如，Vasilache等人[2]提出了一種領(lǐng)域?qū)Ｓ谜Z言和編譯框架TC(tensor comprehensions)。為了優(yōu)化程序性能，TC在總體設(shè)計(jì)上采用了多面體編譯框架(優(yōu)化循環(huán)結(jié)構(gòu))，并在其中集成了如存儲層次提升、映射到寄存器、向量化、并行化等面向硬件結(jié)構(gòu)的優(yōu)化技術(shù)。相關(guān)工作還有DLVM[3]、Latte[4]等。

然而，上述相關(guān)工作在考慮性能優(yōu)化時，更側(cè)重于對程序的運(yùn)算部分進(jìn)行優(yōu)化，對數(shù)據(jù)的優(yōu)化非常有限，比如TVM只在高級優(yōu)化層面對數(shù)據(jù)布局進(jìn)行轉(zhuǎn)換。本文基于以下2點(diǎn)觀察：

(1) 深度學(xué)習(xí)處理器往往需要支持可變大小的向量乘矩陣、矩陣乘矩陣操作，為了配合指令集和硬件結(jié)構(gòu)進(jìn)行訪存、運(yùn)算優(yōu)化，數(shù)據(jù)往往需要根據(jù)不同的操作以及不同的運(yùn)算規(guī)模進(jìn)行不同的變換。

(2) 深度學(xué)習(xí)算法在推理過程中存在大量的常量數(shù)據(jù)(權(quán)值、偏置等)，可以采用部分求值的思想在編譯階段對運(yùn)算、數(shù)據(jù)進(jìn)行提前優(yōu)化。

提出一種運(yùn)算和數(shù)據(jù)協(xié)同優(yōu)化的深度學(xué)習(xí)編譯框架CDUCA(computation and data unified compile architecture)，以解決深度學(xué)習(xí)性能優(yōu)化問題。本文研究的內(nèi)容不涉及如何解決通用性問題。

本文剩余章節(jié)內(nèi)容組織如下。第1節(jié)介紹CDUCA的總體設(shè)計(jì)。第2節(jié)介紹CDUCA各個模塊以及優(yōu)化技術(shù)。第3節(jié)對CDUCA生成模型的性能進(jìn)行了多個維度的實(shí)驗(yàn)評估。第4節(jié)對本文的工作進(jìn)行總結(jié)和展望。

1 總體設(shè)計(jì)

深度學(xué)習(xí)算法的主體是神經(jīng)網(wǎng)絡(luò)。神經(jīng)網(wǎng)絡(luò)由一系列的基本操作按一定的拓?fù)浣Y(jié)構(gòu)連接組成，每個操作包含一組或多組輸入、輸出神經(jīng)元，神經(jīng)元可以在操作之間共享。進(jìn)一步，神經(jīng)網(wǎng)絡(luò)可以用計(jì)算圖來表示，圖節(jié)點(diǎn)表示操作，邊表示神經(jīng)元。訓(xùn)練好的模型數(shù)據(jù)也是神經(jīng)網(wǎng)絡(luò)的重要組成部分，比如卷積操作的權(quán)值。

CDUCA編譯框架接收原始的神經(jīng)網(wǎng)絡(luò)計(jì)算圖和模型數(shù)據(jù)，輸出可部署到深度學(xué)習(xí)處理器上執(zhí)行的模型文件，其總體框架如圖1所示。

CDUCA包含計(jì)算圖引擎、代碼生成器、數(shù)據(jù)優(yōu)化器3個不同層次的組件/模塊。

(1) 計(jì)算圖引擎采用線性變換、常量折疊等優(yōu)化技術(shù)，對原始的計(jì)算圖和模型數(shù)據(jù)進(jìn)行面向算法的高級優(yōu)化，生成優(yōu)化后的計(jì)算圖和模型數(shù)據(jù)。

(2) 代碼生成器采用基于cost model的啟發(fā)式搜索策略，對計(jì)算圖引擎優(yōu)化后的計(jì)算圖和模型數(shù)據(jù)進(jìn)行運(yùn)算和數(shù)據(jù)協(xié)同編譯優(yōu)化，生成高效的目標(biāo)平臺代碼。

(3) 數(shù)據(jù)優(yōu)化器采用數(shù)據(jù)對齊、維度變換、精度調(diào)優(yōu)等優(yōu)化技術(shù)，對計(jì)算圖引擎優(yōu)化后的模型數(shù)據(jù)進(jìn)行面向硬件平臺的二次優(yōu)化，生成可部署模型。

圖1 CDUCA總體架構(gòu)圖

此外，得益于層次化的設(shè)計(jì)結(jié)構(gòu)，CDUCA編譯框架具有很強(qiáng)的擴(kuò)展性，它能夠方便地集成各種編譯優(yōu)化技術(shù)。比如，CDUCA可以在計(jì)算圖引擎模塊集成操作聚合技術(shù)，可以在代碼生成器模塊集成多面體循環(huán)優(yōu)化技術(shù)。如何擴(kuò)展這些優(yōu)化技術(shù)不作為本文討論的重點(diǎn)。

2 CDUCA模塊

本小節(jié)介紹CDUCA各個模塊以及優(yōu)化技術(shù)。

2.1 計(jì)算圖引擎

計(jì)算圖引擎采用線性變換和常量折疊技術(shù)來對運(yùn)算和數(shù)據(jù)進(jìn)行協(xié)同優(yōu)化。在介紹計(jì)算圖引擎的工作原理之前，本小節(jié)先介紹這2種優(yōu)化技術(shù)。

神經(jīng)網(wǎng)絡(luò)算法中常見的操作包含卷積、全連接、激活、池化、批規(guī)范化、縮放，本文將這些操作歸為2類：線性變換操作和非線性變換操作。所有的線性變換操作都能夠表示成向量或矩陣乘加的形式，如下：

Y=XW+B

(1)

所有的非線性變換操作都無法表示成式(1)的形式。其中，X、Y是變量，W、B是常量(提前訓(xùn)練好的模型數(shù)據(jù))。按照上面的分類，卷積、全連接、批規(guī)范化、縮放操作是線性變換操作，最大池化、激活是非線性變換操作。以全連接操作為例，式(1)中的X、Y分別表示全連接的輸入、輸出神經(jīng)元矩陣，W表示權(quán)值矩陣，B表示偏置矩陣。上面列舉的其他線性變換操作也能轉(zhuǎn)換成式(1)，具體轉(zhuǎn)換的方法本文不做詳細(xì)描述。

假定有2個連續(xù)的線性變換操作，第2個操作使用第1個操作的輸出作為輸入：

Y1=X1W1+B1

(2)

Y2=X1W2+B2

(3)

由于線性變換滿足分配律和結(jié)合律，又由于W矩陣和B矩陣是常量，因此式(2)、(3)可以做以下等價線性變換：

Y2=(X1W1+B1)W2+B2

(4)

Y2=X1W1W2+B1W2+B2

(5)

W′=W1W2,B′=B1W2+B2

(6)

Y2=X1W′+B′

(7)

因此，原始的2步計(jì)算式(2)和式(3)通過線性變換為式(4)和式(5)和常量折疊式(6)優(yōu)化，簡化成了一步計(jì)算式(7)。這2種優(yōu)化技術(shù)在減少運(yùn)算量的同時也對常量模型數(shù)據(jù)進(jìn)行了壓縮，一方面減少了模型數(shù)據(jù)的存儲開銷，另一方面減少了運(yùn)行時的訪存量。

舉一個具體的應(yīng)用來說明這2種優(yōu)化技術(shù)是實(shí)用的。ResNet[5]是經(jīng)典的圖像分類網(wǎng)絡(luò)，它包含大量的卷積+批規(guī)范化+縮放子結(jié)構(gòu)，如圖2所示。

圖2 ResNet子結(jié)構(gòu)

上面3個操作都是線性變換操作，因此可以進(jìn)行線性變換和常量折疊優(yōu)化。具體的優(yōu)化步驟參考式(2)～(7)。

計(jì)算圖引擎的具體工作原理如下(圖3)。首先掃描原始的計(jì)算圖，識別其中連續(xù)的線性變換操作；然后對連續(xù)的線性變換操作執(zhí)行線性變換和常量折疊優(yōu)化，把多個操作節(jié)點(diǎn)合并成一個操作節(jié)點(diǎn)，把多份模型數(shù)據(jù)合并成一份模型數(shù)據(jù)；最后輸出簡化后的計(jì)算圖和模型數(shù)據(jù)。

圖3 計(jì)算圖引擎工作原理圖

2.2 代碼生成器

代碼生成器在生成目標(biāo)代碼時充分考慮了對運(yùn)算和數(shù)據(jù)進(jìn)行協(xié)同優(yōu)化。深度學(xué)習(xí)處理器為了優(yōu)化訪存性能,往往會在靠近計(jì)算單元的地方設(shè)計(jì)片上緩存(on-chip memory)，訪問片上緩存的速度遠(yuǎn)比從片外DDR快。為了簡化結(jié)構(gòu)設(shè)計(jì)，深度學(xué)習(xí)處理器并沒有將片上緩存設(shè)計(jì)成cache這種軟件透明的結(jié)構(gòu)，因此軟件編程人員需要顯示地管理這些片上緩存[6,7]。合理利用片上緩存能夠極大地提升程序的性能，單從優(yōu)化運(yùn)算的角度出發(fā)無法做到這一點(diǎn)。為了在支持可變大小的向量乘矩陣、矩陣乘矩陣操作應(yīng)用場景下充分利用片上緩存，代碼生成器必須同時考慮運(yùn)算優(yōu)化和數(shù)據(jù)優(yōu)化。另一方面，由于深度學(xué)習(xí)處理器支持向量化操作以及按固定維度進(jìn)行循環(huán)展開[8]，從提高訪存效率的角度出發(fā)，片外數(shù)據(jù)往往需要做特殊的對齊以及維度變換。

CDUCA代碼生成器在設(shè)計(jì)時充分考慮了上面幾個問題，并提出了以下解決方案。

采用基于cost model的啟發(fā)式搜索策略來協(xié)同優(yōu)化運(yùn)算和數(shù)據(jù)。cost model給出程序的運(yùn)行時間估計(jì)，它考慮的核心因素是訪存時間、運(yùn)算時間以及兩者的覆蓋率。增大覆蓋率是為了隱藏計(jì)算、訪存延時，在此基礎(chǔ)下，盡量減小訪存和運(yùn)算時間就能提升程序的性能。在生成目標(biāo)平臺指令時存在多種選擇，比如，分幾次加載一塊完整的數(shù)據(jù)，數(shù)據(jù)存放的位置(片上緩存/片外DDR),單次運(yùn)算數(shù)據(jù)量的大小等。使用暴力搜索的開銷過大，因此，本文提出了一種啟發(fā)式的指令生成策略，一個重要的啟發(fā)式策略就是盡量把片上緩存用滿。

代碼生成器的完整工作流程如下(圖4)。首先對計(jì)算圖引擎優(yōu)化后的計(jì)算圖進(jìn)行cost model建模，然后采用啟發(fā)式搜索策略來生成高度優(yōu)化的目標(biāo)平臺指令。

圖4 代碼生成器工作原理圖

2.3 數(shù)據(jù)優(yōu)化器

數(shù)據(jù)優(yōu)化器根據(jù)目標(biāo)平臺代碼所需要的最優(yōu)數(shù)據(jù)布局，對模型數(shù)據(jù)進(jìn)行面向硬件平臺的二次優(yōu)化。主要的優(yōu)化包括：

(1) 按照運(yùn)算的要求對數(shù)據(jù)進(jìn)行對齊，加快訪存速度。深度學(xué)習(xí)處理器向量化運(yùn)算對數(shù)據(jù)對齊有一定的要求，對齊訪問能夠加快訪存速度。

(2) 按照運(yùn)算的要求對模型數(shù)據(jù)進(jìn)行維度變換，提高訪存局部性。卷積等算法需要將數(shù)據(jù)解釋成多維數(shù)組，多維數(shù)組的排布順序會影響訪存的跳轉(zhuǎn)次數(shù)。為了盡量減少訪存跳轉(zhuǎn)，需要對數(shù)據(jù)進(jìn)行維度變換。

(3) 精度調(diào)優(yōu)，提高運(yùn)算速度。深度學(xué)習(xí)處理器往往支持低精度運(yùn)算[8，9]，比如16比特量化、8比特量化，不同的精度運(yùn)算性能不同。因此，在編譯時期，可以對模型數(shù)據(jù)進(jìn)行精度選優(yōu)，從而提高性能。

數(shù)據(jù)優(yōu)化器的工作原理如下(圖5)。先按照目標(biāo)平臺代碼對數(shù)據(jù)布局的要求，對模型數(shù)據(jù)進(jìn)行上述3種優(yōu)化；然后將目標(biāo)代碼和優(yōu)化后的模型數(shù)據(jù)打包，生成可部署模型。

圖5 數(shù)據(jù)優(yōu)化器工作原理

3 實(shí)驗(yàn)和評估

本文在DLPlib[10]的基礎(chǔ)上進(jìn)行了編程模型擴(kuò)展，使其支持了網(wǎng)絡(luò)級別的計(jì)算，并在內(nèi)部實(shí)現(xiàn)了CDUCA編譯優(yōu)化框架。本文對CDUCA編譯生成的模型進(jìn)行了多方面的性能評估。接下來先介紹本文使用的硬件平臺和測試集。

(1) 硬件平臺

本文采用與DaDianNao[11]類似的設(shè)計(jì)，用verilog實(shí)現(xiàn)了大部分硬件邏輯，并在SynopsysFPGA上進(jìn)行綜合。在100 MHz的現(xiàn)場可編程門陣列(field programmable gate array, FPGA)主頻下，單個處理器核心的峰值運(yùn)算性能可以達(dá)到0.2 GFOP/s。DaDianNao包含多個處理核心，本文在測試時只使用了單個運(yùn)算核心，但這并不表示CDUCA框架不支持多核系統(tǒng)。

(2) 測試集

本文選取了典型的圖像分類網(wǎng)絡(luò)如AlexNet[12]、ResNet[5]和典型的目標(biāo)檢測網(wǎng)絡(luò)如Yolov2[13]、Faster-RCNN[14]作為測試集。

本文將CDUCA生成的模型與手工優(yōu)化模型進(jìn)行了性能對比。手工優(yōu)化模型是指由專業(yè)人士使用匯編指令手工實(shí)現(xiàn)算法模型，這種方式能夠把特定的算法模型性能優(yōu)化到極致，其缺點(diǎn)是開發(fā)成本太高。

圖6給出了CDUCA模型與手工優(yōu)化模型性能的對比。其中CDUCA-NOPT系列表示的是關(guān)閉了計(jì)算圖引擎的CDUCA。對比CDUCA-NOT和CDUCA可以看出，計(jì)算圖引擎優(yōu)化對ResNet50、Yolov2性能提升很大(10%以上)，而對AlexNet、Faster-RCNN提升不大，因?yàn)锳lexNet中不包含連續(xù)的線性變換操作，而Faster-RCNN主要開銷并不是在線性操作部分。表1給出了CDUCA模型和手工優(yōu)化模型性能的比值，在平均值上，CDUCA生成的模型性能能達(dá)到手工優(yōu)化模型性能的86.5%。

圖6 CDUCA模型與手工優(yōu)化模型性能對比

表1 CDUCA模型與手工優(yōu)化模型性能對比

AlexNetResNet50Yolov2Faster-RCNNMean89.1%90.7%83.0%83.2%86.5%

4 結(jié) 論

本文從深度學(xué)習(xí)算法和深度學(xué)習(xí)硬件平臺的特點(diǎn)出發(fā)，提出了一種運(yùn)算和數(shù)據(jù)協(xié)同優(yōu)化的深度學(xué)習(xí)編譯框架CDUCA。在計(jì)算圖引擎模塊，采用線性變換和常量折疊技術(shù)對原始計(jì)算圖和模型數(shù)據(jù)進(jìn)行面向算法的協(xié)同優(yōu)化；在代碼生成器模塊，采用基于cost model的啟發(fā)式搜索技術(shù)，對運(yùn)算和數(shù)據(jù)進(jìn)行協(xié)同優(yōu)化，生成目標(biāo)代碼；在數(shù)據(jù)優(yōu)化器模塊，采用數(shù)據(jù)對齊、維度變換、精度調(diào)優(yōu)技術(shù)，對模型數(shù)據(jù)進(jìn)行面向硬件平臺的二次優(yōu)化。CDUCA在整個框架設(shè)計(jì)中貫穿運(yùn)算和數(shù)據(jù)協(xié)同優(yōu)化的思想。實(shí)驗(yàn)結(jié)果表明，CDUCA生成的模型性能能達(dá)到手工優(yōu)化模型性能的86.5%。

此外CDUCA框架層次化的設(shè)計(jì)結(jié)構(gòu)具有很強(qiáng)的擴(kuò)展性，能夠在各個模塊集成各種編譯優(yōu)化技術(shù)。比如，CDUCA可以在計(jì)算圖引擎模塊集成操作聚合技術(shù)，可以在代碼生成器模塊集成多面體編譯優(yōu)化技術(shù)，集成這些優(yōu)化技術(shù)可以作為本文的后續(xù)研究工作。

總結(jié)來說，在深度學(xué)習(xí)編譯優(yōu)化領(lǐng)域，對運(yùn)算和數(shù)據(jù)進(jìn)行協(xié)同優(yōu)化非常重要，這是后續(xù)通用深度學(xué)習(xí)編譯框架設(shè)計(jì)中應(yīng)該重點(diǎn)考慮的。