樊 哲 南子淵 郝一帆 杜子?xùn)|③ 陳云霽

(?中國科學(xué)院計算技術(shù)研究所計算機(jī)體系結(jié)構(gòu)國家重點實驗室 北京 100190)

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

0 引言

深度學(xué)習(xí)作為一個熱門研究方向,其呈現(xiàn)出多樣化和輕量化的發(fā)展趨勢。多樣化趨勢主要體現(xiàn)在模型和硬件的不斷更迭。模型發(fā)展如AlexNet[1]、ResNet[2]、CapsNet[3]和Transformer[4]等;硬件發(fā)展如中央處理器(central processing unit,CPU)、圖形處理器(graphics processing unit,GPU)、DianNao 系列[5-6]、張量處理器(tensor processing unit,TPU)[7]以及各種針對特定算法的專用加速器[8-10]。輕量化趨勢主要體現(xiàn)在深度學(xué)習(xí)應(yīng)用從云端到邊端的遷移,即越來越多的深度學(xué)習(xí)應(yīng)用被部署到算力受限的邊緣端設(shè)備[11-13]。這對深度學(xué)習(xí)應(yīng)用的算法優(yōu)化提出了越來越高的要求。上述發(fā)展趨勢帶來了深度學(xué)習(xí)應(yīng)用的編程和性能之間的矛盾。

深度學(xué)習(xí)張量程序(張量算子的低層級實現(xiàn)[14])的自動生成框架致力于解決編程與性能之間的矛盾。其被廣泛用于深度學(xué)習(xí)編譯器(如張量虛擬機(jī)(tensor virtual machine,TVM)[15])以自動生成高性能的張量程序。其通過搜索的方式找尋在指定硬件后端上性能最優(yōu)的張量程序,之后交由對應(yīng)的硬件編譯器生成二進(jìn)制指令,在減輕用戶編程負(fù)擔(dān)的同時盡量保證性能。

為了生成高效的深度學(xué)習(xí)張量程序,必須考慮靜態(tài)數(shù)據(jù)的布局問題。原因主要有3 個方面:(1)對靜態(tài)數(shù)據(jù)的布局變換可發(fā)生在編譯時期,而不占用程序的運行時間;(2)在深度卷積神經(jīng)網(wǎng)絡(luò)的前向推理(inference)過程中,靜態(tài)數(shù)據(jù)(如權(quán)值、偏置等)參與的運算高達(dá)總運算量的90%[16-17];(3)數(shù)據(jù)布局可以提升程序訪存局部性和硬件資源利用率[18-20]。因此,靜態(tài)數(shù)據(jù)的布局對深度學(xué)習(xí)張量程序的性能影響十分顯著;而如何確定合適的靜態(tài)數(shù)據(jù)布局,使得程序執(zhí)行效率最高,是深度學(xué)習(xí)張量程序自動生成框架面臨的重大挑戰(zhàn)。

作為目前應(yīng)用最廣泛、最具前景的深度學(xué)習(xí)張量程序自動生成框架,Ansor[14]應(yīng)對上述挑戰(zhàn)的方案是:根據(jù)預(yù)先指定的單一數(shù)據(jù)布局策略,訓(xùn)練性能預(yù)測模型,依據(jù)該模型搜索最佳性能的張量程序。但其存在單一策略非最優(yōu)和性能預(yù)測模型不準(zhǔn)確的問題。具體地,單一的數(shù)據(jù)布局策略無法使Ansor在所有任務(wù)上都生成性能最優(yōu)的程序,而對于一個任務(wù),其最優(yōu)策略無法預(yù)先確定,因此預(yù)先指定單一策略的Ansor 生成的程序仍有性能提升空間(即單一策略非最優(yōu)的問題)。另外,指定一種策略后,Ansor 會依據(jù)性能預(yù)測模型C從任務(wù)的程序空間(即搜索空間)中采樣一個程序A,用對A進(jìn)行布局變換得到的程序X去訓(xùn)練C,以指導(dǎo)下一次采樣,直到搜索過程結(jié)束。最后Ansor 輸出搜索過程中遇到的性能最優(yōu)程序X?。從優(yōu)化角度看,X所在程序空間即為優(yōu)化空間。C并非直接在優(yōu)化空間中尋優(yōu),而是通過在搜索空間中采樣,間接在優(yōu)化空間中尋優(yōu)。優(yōu)化空間的程序因經(jīng)過了數(shù)據(jù)布局,訪存局部性要比搜索空間的程序好,因而優(yōu)化空間和搜索空間在訪存特性上存在巨大差異。這種差異導(dǎo)致Ansor 得到的程序X?仍有性能提升空間(即性能預(yù)測模型不準(zhǔn)確的問題)。

針對上述問題,本文提出基于自適應(yīng)靜態(tài)數(shù)據(jù)布局(adaptive layout,AL)策略的深度學(xué)習(xí)張量程序自動生成框架AL-Ansor。其通過自適應(yīng)地選取多種靜態(tài)數(shù)據(jù)布局策略來解決單一策略非最優(yōu)的問題;通過多種策略共同訓(xùn)練性能預(yù)測模型C來解決模型不準(zhǔn)確的問題。具體地,AL-Asnor 可先以直接尋優(yōu)的布局策略預(yù)訓(xùn)練C,將搜索空間限制在靜態(tài)數(shù)據(jù)訪存局部性較好的程序子空間。然后,再以間接尋優(yōu)的布局策略(此時搜索空間和優(yōu)化空間的不一致已被緩解)微調(diào)C,使之更容易在搜索空間中采樣出程序A?,其經(jīng)布局變換后的程序X?的性能最優(yōu)。本文在32 核Intel Xeon CPU 上對多個卷積層進(jìn)行實驗驗證,結(jié)果表明,在同樣的搜索次數(shù)下,相較于3 種基于單一策略的Ansor,AL-Ansor 生成的張量程序分別有13.81%、12.41%和16.59%的平均性能提升。

綜上,本文的主要貢獻(xiàn)有以下3 點:

(1) 分析了Ansor 的工作原理,發(fā)現(xiàn)了其單一策略非最優(yōu)和性能預(yù)測模型不準(zhǔn)確的問題。

(2) 提出了自適應(yīng)靜態(tài)數(shù)據(jù)布局策略AL 以及基于此策略的張量程序自動生成框架AL-Ansor。

(3) 在典型深度神經(jīng)網(wǎng)絡(luò)的多個卷積層上進(jìn)行實驗驗證,證明了AL 策略的有效性。

1 背景

1.1 張量程序自動生成框架

由于深度學(xué)習(xí)模型和硬件的多樣化發(fā)展,基于搜索的張量程序自動生成框架成為一個重要研究方向。這些框架以描述深度學(xué)習(xí)模型/算子的計算圖(computational graph)和目標(biāo)硬件信息作為主要輸入,以優(yōu)化硬件執(zhí)行時間為主要目標(biāo)執(zhí)行搜索過程,來自動生成對不同硬件后端具有通用表達(dá)能力的張量程序。

張量程序自動生成框架按照其搜索空間的定義方式可以分為2 類:基于模板搜索(如文獻(xiàn)[21])和無模板搜索(如文獻(xiàn)[14,22]),如圖1 所示?；谀０逅阉鞯乃阉骺臻g由用戶自定義模板來指定;而無模板搜索的搜索空間無需用戶參與指定。AutoTVM 是基于模板搜索的一種典型框架。在該框架下,用戶需要針對不同硬件自定義不同算子的搜索模板。搜索模板的內(nèi)容包括算子的實現(xiàn)方式,輸入數(shù)據(jù)應(yīng)該如何變換布局,以及輸入數(shù)據(jù)以什么方式(循環(huán)結(jié)構(gòu)、循環(huán)順序、向量化、并行化等)進(jìn)行計算等。算子搜索模板內(nèi)的不確定因素(循環(huán)的拆分因子、是否向量化等)定義了輸入計算圖對應(yīng)的張量程序搜索空間。但這樣以自定義模板方式定義的空間是非常受限的,也不易擴(kuò)展[14]。Ansor 是無模板搜索的一種典型框架。其搜索空間由內(nèi)置的與算子解綁的通用派生規(guī)則(derivation rules)定義:這些規(guī)則可作用于任意輸入計算圖來實現(xiàn)對計算圖的改寫,每個改寫的計算圖對應(yīng)一種張量程序?qū)崿F(xiàn),這些構(gòu)成輸入計算圖的張量程序空間,即搜索空間。綜上,Ansor 無需為每個計算圖專門定義搜索空間,也無需用戶參與,因此,其更容易應(yīng)用于新的算子,也緩解了用戶的編程負(fù)擔(dān),是目前使用最廣泛、最有前景的深度學(xué)習(xí)張量程序自動生成框架。

圖1 張量程序自動生成框架示意圖

1.2 數(shù)據(jù)布局優(yōu)化

數(shù)據(jù)布局指的是數(shù)據(jù)在內(nèi)存中的擺放順序。改變一個程序所需數(shù)據(jù)的布局,會影響程序的訪存局部性和硬件的計算資源利用率。數(shù)據(jù)布局優(yōu)化,即通過改變數(shù)據(jù)布局,改善程序的訪存局部性,使整個程序的性能更好。

圖2 展示了矩陣乘程序的一個數(shù)據(jù)布局優(yōu)化過程。圖2(a)展示了一種實現(xiàn)(M,K)×(K,N) 矩陣乘的程序A。假設(shè)2 個輸入矩陣U、V以及輸出矩陣Z都按行(row-major)存儲。對A的輸入數(shù)據(jù)V進(jìn)行數(shù)據(jù)布局優(yōu)化后,會得到2 個程序:圖2(b)所示的布局變換程序L和圖2(c)所示的布局后程序P。布局變換程序L可根據(jù)程序A對數(shù)據(jù)V的下標(biāo)訪問規(guī)則自動生成(Ansor 中的數(shù)據(jù)布局變換也是這樣實現(xiàn)的),使得布局后的數(shù)據(jù)V’在同樣的循環(huán)結(jié)構(gòu),即程序P中擁有更好的訪存局部性。布局后程序P和程序A的差別如圖2(c)中加粗部分所示,僅在于數(shù)據(jù)V’。數(shù)據(jù)布局優(yōu)化除了使程序P擁有更好的訪存局部性外,還可以提升硬件計算資源的利用率。對于一個擁有NI個KI維向量點積運算單元的硬件平臺,程序P的最內(nèi)兩重循環(huán)(即圖2(c)中虛線框部分)可以直接由這些向量點積運算單元處理,而程序A由于對數(shù)據(jù)V的訪存不連續(xù),無法直接利用這些運算單元。如果數(shù)據(jù)V是編譯時可確定的靜態(tài)數(shù)據(jù),則上述數(shù)據(jù)布局優(yōu)化稱為靜態(tài)數(shù)據(jù)布局優(yōu)化,此時布局變換程序L可在編譯過程中執(zhí)行,不會帶來運行時開銷。

圖2 數(shù)據(jù)布局優(yōu)化示例(虛線框部分可進(jìn)行并行加速,從而提高硬件計算資源利用率;加粗部分展示了P 和A 的差異)

對于深度學(xué)習(xí)應(yīng)用而言,數(shù)據(jù)布局優(yōu)化是提升程序性能的一個重要優(yōu)化策略。Caffe[23]在實現(xiàn)卷積操作時,先對輸入數(shù)據(jù)進(jìn)行im2col 布局變換,再調(diào)用高效的通用矩陣乘(general matrix-matrix multiplication,GEMM)接口實現(xiàn)卷積功能。文獻(xiàn)[24]為在單指令數(shù)據(jù)流(single instruction multiple data,SIMD)架構(gòu)上實現(xiàn)高效的卷積運算,需要將輸入布局設(shè)置為NCHWc,權(quán)值布局設(shè)置為KCRSck。LDLNDK(neural network development kit with labled data layout)[25]使用標(biāo)簽標(biāo)記靜態(tài)數(shù)據(jù)布局信息,在編譯時進(jìn)行布局變換,在深度學(xué)習(xí)加速器上達(dá)到數(shù)倍性能提升。CDUCA(computation and data unified compile architecture)[26]的數(shù)據(jù)優(yōu)化器模塊在部署前對數(shù)據(jù)進(jìn)行面向硬件平臺的數(shù)據(jù)布局等優(yōu)化,提升硬件平臺的訪存和運算效率。NeoCPU[27]針對神經(jīng)網(wǎng)絡(luò)在CPU 上的前向推理過程,在圖級別對各算子的神經(jīng)元數(shù)據(jù)布局進(jìn)行規(guī)劃,再對靜態(tài)權(quán)值數(shù)據(jù)進(jìn)行相應(yīng)的布局變換,達(dá)到高性能。

2 問題定義

為了更加清晰地進(jìn)行后續(xù)論述,本文在此定義并引入一些符號,并對張量程序自動生成框架的靜態(tài)數(shù)據(jù)布局問題進(jìn)行定義。

2.1 任務(wù)及程序空間

任務(wù)T由輸入數(shù)據(jù)、輸出數(shù)據(jù)以及輸入數(shù)據(jù)到輸出數(shù)據(jù)的映射規(guī)則共同定義。輸入數(shù)據(jù)根據(jù)其內(nèi)容是否為編譯時確定又可以分為靜態(tài)輸入數(shù)據(jù)和動態(tài)輸入數(shù)據(jù)。例如,圖2 所示為一個將M×K維輸入數(shù)據(jù)和K×N維輸入數(shù)據(jù)進(jìn)行矩陣乘得到M×N維輸出數(shù)據(jù)的任務(wù)MatMul。

程序解決給定任務(wù)(實現(xiàn)輸入數(shù)據(jù)到輸出數(shù)據(jù)映射規(guī)則)的一個有限可終止的步驟序列(包括訪存步驟和計算步驟)。對于圖2 所示程序A和P,雖然其計算步驟是相同的,但是其輸入(分別為V和V’)和訪存步驟不同,因此是2 個不同的程序。

任務(wù)T的程序空間AT解決給定任務(wù)T的所有程序構(gòu)成的空間。圖2 中程序A和(L+P) 均為任務(wù)MatMul 的程序空間AMatMul中一點。對于一個具體的搜索過程(如Ansor 的搜索過程),任務(wù)T的程序空間受限于搜索空間的生成方式,但為了描述的方便,會將該搜索過程所能探索到的程序空間記為AT。

任務(wù)T的布局分離程序空間ST是對空間AT中每個程序進(jìn)行數(shù)據(jù)布局優(yōu)化所得到的程序構(gòu)成的空間。顯然,該空間是AT的一個子空間。在圖2 所示的MatMul 任務(wù)中,對程序A∈AMatmul進(jìn)行數(shù)據(jù)布局優(yōu)化后得到的組合程序(L+P)為SMatMul空間中一點。

任務(wù)T的布局后程序空間PT是對空間AT中每個程序進(jìn)行數(shù)據(jù)布局優(yōu)化所得到的布局后程序構(gòu)成的空間。因此,該空間的程序訪存局部性都很好,與空間AT在程序訪存特征上存在較大差異。根據(jù)定義,AT和ST上的點可以在PT上找到點與之對應(yīng),反之,PT上的點也都可以在AT和ST上找到點與之對應(yīng)。但PT并不是空間AT或ST的子空間,因為PT中程序解決的任務(wù)并不是T本身,而是T經(jīng)過布局變換程序后的子任務(wù)(該子任務(wù)的輸入由于數(shù)據(jù)布局變換,已不同于任務(wù)T中對應(yīng)的輸入,如圖2 中的V和V’)。圖2(c)所示程序P即為PMatMul空間中一點。

圖3 展示了上述術(shù)語的關(guān)系。對于一個任務(wù)T,從其程序空間AT中取一程序A,其通過數(shù)據(jù)布局優(yōu)化會產(chǎn)生布局變換程序L和布局后程序P。其中,程序(L+P) 屬于任務(wù)T的布局分離程序空間ST,程序P屬于任務(wù)T的布局后程序空間PT。虛線展示了程序A到P的數(shù)據(jù)布局優(yōu)化路徑。反過來,PT中的每一個程序也都可以在ST和AT中找到對應(yīng)的程序。即對于一個任務(wù)T,這3 個空間中存在一一對應(yīng)的路徑。

圖3 3 種程序空間的關(guān)系以及其中程序的對應(yīng)關(guān)系

2.2 張量程序自動生成框架的靜態(tài)數(shù)據(jù)布局

對于張量程序自動生成框架,給定任務(wù)T和目標(biāo)硬件H,其靜態(tài)數(shù)據(jù)布局問題定義如下。

定義1在程序空間XT中搜索滿足優(yōu)化目標(biāo)G的張量程序X?,使得在H上解決任務(wù)T所需的運行時間最短。

根據(jù)張量程序自動生成框架所采取的搜索策略(靜態(tài)數(shù)據(jù)布局策略)的不同,定義1 中的程序空間XT(優(yōu)化空間)和輸出程序X?會有所不同。

對于本文關(guān)注的Ansor 而言,其共有3 種搜索策略:(1)不進(jìn)行布局(no layout,NL)策略;(2)運行時布局(layout during running,LR)策略;和(3)編譯時布局(layout during compiling,LC)策略。其中第3 種策略只能用于靜態(tài)數(shù)據(jù),而前2 種對于動態(tài)數(shù)據(jù)和靜態(tài)數(shù)據(jù)均適用。在Ansor 中,不同策略使用不同的優(yōu)化空間XT(AT或ST或PT),但卻使用同一個搜索空間AT,因為定義于輸入計算圖上的通用派生規(guī)則無法直接生成位于ST和PT中的程序。因此,Ansor 只能從AT中(準(zhǔn)確地來說應(yīng)該是AT-ST)采樣程序進(jìn)行搜索,然后經(jīng)數(shù)據(jù)布局變換映射到相應(yīng)優(yōu)化空間(AT或ST或PT)上進(jìn)行測量。這一點體現(xiàn)在優(yōu)化目標(biāo)G上。

對于策略NL,定義1 中的XT是任務(wù)T的程序空間AT,輸出程序X?是空間AT中的程序A?。優(yōu)化目標(biāo)G如式(1)所示,其中,timeof(A,H) 表示程序A在硬件H上的執(zhí)行時間。

對于策略LR,定義1 中的XT是任務(wù)T的布局分離程序空間ST,輸出程序X?是空間ST中的程序(L+P)?。優(yōu)化目標(biāo)G如式(2)所示,其中,LT0(A)表示對程序A進(jìn)行數(shù)據(jù)布局優(yōu)化得到的布局變換程序,LT1(A) 表示對程序A進(jìn)行數(shù)據(jù)布局優(yōu)化得到的布局后程序。

對于策略LC,定義1 中的XT是任務(wù)T的布局后程序空間PT,優(yōu)化目標(biāo)G如式(3)所示,搜索的輸出程序是空間PT中的程序P?。

對于策略LR 和策略LC,由于優(yōu)化空間XT和搜索空間AT并非同一空間,因此其優(yōu)化做法——依據(jù)在優(yōu)化空間XT上定義的優(yōu)化目標(biāo)G去采樣位于搜索空間AT中的程序A,以求A經(jīng)過變換后得到的程序X能有最小的運行時間timeof(X,H)——是無法達(dá)到目標(biāo)的。尤其是對策略LC,因為其優(yōu)化空間PT與搜索空間AT中程序的訪存特征差異過大。

3 采用自適應(yīng)靜態(tài)數(shù)據(jù)布局策略的AL-Ansor

3.1 Ansor 的工作原理及靜態(tài)數(shù)據(jù)布局策略

Ansor 是一種深度學(xué)習(xí)張量程序自動生成框架,其工作原理如圖4 所示。

圖4 Ansor 工作原理示意圖

Ansor 以待優(yōu)化算子(記為任務(wù)T)的計算圖、目標(biāo)硬件信息、搜索次數(shù)以及指定的某一靜態(tài)數(shù)據(jù)布局策略(NL、LR 或LC)為輸入,經(jīng)過指定次數(shù)的搜索后,輸出在目標(biāo)硬件上執(zhí)行時間最短的張量程序X?。在每一輪搜索過程中,Ansor 主要分為4 步進(jìn)行:(1)程序采樣。依據(jù)性能預(yù)測模型C,從由輸入計算圖和圖上的通用派生規(guī)則共同定義的張量程序空間AT中采樣若干程序。(2)靜態(tài)數(shù)據(jù)布局。依據(jù)輸入的靜態(tài)數(shù)據(jù)布局策略,對上一步采樣出的每一個程序A進(jìn)行相應(yīng)的變換,得到程序X。(3)編譯并執(zhí)行。將程序X編譯至目標(biāo)硬件平臺,并將相應(yīng)二進(jìn)制指令部署至目標(biāo)硬件平臺上運行,這個過程中會得到程序X的指令、運行時間等信息。(4)訓(xùn)練性能預(yù)測模型C。從上一步的編譯及運行信息中提取特征訓(xùn)練性能預(yù)測模型C,用來指導(dǎo)下一輪搜索過程中第一步的程序采樣。

在Ansor 的一次完整搜索流程中,需要指定固定的一種靜態(tài)數(shù)據(jù)布局策略。不同的靜態(tài)數(shù)據(jù)布局策略會影響性能預(yù)測模型C的訓(xùn)練。

對于策略NL,程序X就是采樣程序A。該策略下,性能預(yù)測模型C以程序A的信息來訓(xùn)練。此時Ansor 的搜索空間和優(yōu)化空間均為AT,優(yōu)化目標(biāo)為2.2 節(jié)中的式(1)。

對于策略LR,程序X是采樣程序A經(jīng)過靜態(tài)數(shù)據(jù)布局變換而產(chǎn)生的布局變換程序L和布局后程序P的組合。該策略下,性能預(yù)測模型C以(L+P)∈ST的特征信息訓(xùn)練,來進(jìn)行程序A∈AT的選擇。此時Ansor 的搜索空間為AT,而優(yōu)化空間為AT的子空間ST,優(yōu)化目標(biāo)為2.2 節(jié)中的式(2)。

對于策略LC,程序X是采樣程序A經(jīng)過靜態(tài)數(shù)據(jù)布局變換產(chǎn)生的布局后程序P,因為布局變換程序L被指定在編譯時執(zhí)行。該策略下,性能預(yù)測模型C以P∈PT的特征信息訓(xùn)練,來進(jìn)行程序A∈AT的選擇。此時Ansor 的搜索空間為AT,而優(yōu)化空間為訪存局部性更好的布局后程序空間PT,優(yōu)化目標(biāo)為2.2 節(jié)中的式(3)。優(yōu)化空間和搜索空間不在同一個空間,且其中的程序訪存特征具有較大的差異。

表1 列出了以32 核Intel Xeon CPU 為目標(biāo)硬件,以ResNet-50 的4 個卷積塊(包括輸入補(bǔ)pad、2D卷積、加偏置以及做ReLU 激活)為搜索任務(wù),Ansor分別在3 種靜態(tài)布局策略下搜索1024 次得到的最優(yōu)程序X?的運行時間?？梢钥闯?單一的策略NL、LR 和LC 并不總在所有搜索任務(wù)上表現(xiàn)得最好。

表1 Ansor 的3 種靜態(tài)數(shù)據(jù)布局策略搜索結(jié)果

3.2 自適應(yīng)靜態(tài)數(shù)據(jù)布局策略

從Ansor 的工作原理可知,其存在以下2 個問題:

(1)靜態(tài)數(shù)據(jù)布局策略需要預(yù)先指定,而針對不同的任務(wù),最優(yōu)的策略不同,因此預(yù)先指定的靜態(tài)數(shù)據(jù)布局策略往往不是最優(yōu)的(單一策略非最優(yōu))。

(2)對于編譯時布局策略LC,搜索空間AT與優(yōu)化空間PT存在較大差異,主要體現(xiàn)在程序?qū)o態(tài)數(shù)據(jù)的訪存特性上,PT上的程序訪存特性更好。由此訓(xùn)練得到的性能預(yù)測模型難以找到性能最優(yōu)的程序(性能預(yù)測模型不準(zhǔn)確)。

針對這些問題,本文提出自適應(yīng)靜態(tài)數(shù)據(jù)布局策略(AL)。其主要特點是,在搜索過程中不是只使用一種固定的靜態(tài)數(shù)據(jù)布局策略,而是根據(jù)一個自適應(yīng)策略選擇函數(shù)Select_Strategy 選擇并應(yīng)用多種靜態(tài)數(shù)據(jù)布局策略(NL/LR/LC),來共同訓(xùn)練性能預(yù)測模型C。一種可能的自適應(yīng)策略是:在進(jìn)行策略LC 前,先使用策略NL 和策略LR 來訓(xùn)練模型C。這樣不僅有機(jī)會探索到不同策略下的最優(yōu)程序,還可以縮小待搜索空間。之后再采取策略LC 微調(diào)模型C,此時優(yōu)化空間和搜索空間的差異不那么顯著,使得模型C更容易在搜索空間AT中找到程序A?,其對應(yīng)的布局后程序P?在目標(biāo)硬件上執(zhí)行時間最短。

自適應(yīng)靜態(tài)數(shù)據(jù)布局策略AL 的偽代碼描述如算法1 所示。

算法1 的輸入是任務(wù)T的張量程序空間AT(也即該算法的搜索空間),目標(biāo)硬件H,搜索輪數(shù)L,以及初始性能預(yù)測模型C。輸出是在搜索過程遇到的在目標(biāo)硬件H上執(zhí)行時間最短的布局后程序P?。

在每一輪搜索過程中,主要經(jīng)過了以下過程。

(1)先根據(jù)目標(biāo)硬件H和性能預(yù)測模型C在搜索空間AT中采樣出候選的程序集合A_Cand(第3行)。

(2)對集合中每個程序進(jìn)行數(shù)據(jù)布局變換,得到布局變換程序集合L_Cand 和布局后程序集合P_Cand(第4 行)。

(3)根據(jù)自適應(yīng)策略選擇函數(shù)Select_Strategy確定本輪搜索應(yīng)使用的靜態(tài)數(shù)據(jù)布局策略(第5行)。該函數(shù)的輸入可以有當(dāng)前的搜索輪數(shù)Search_Idx、搜索總輪數(shù)L以及性能預(yù)測模型C等。該函數(shù)的實現(xiàn)可以是一種啟發(fā)式算法,一種預(yù)定義規(guī)則,或者是一個預(yù)訓(xùn)練的決策樹或神經(jīng)網(wǎng)絡(luò)模型。

(4)根據(jù)本輪搜索的靜態(tài)數(shù)據(jù)布局策略Strategy,對相應(yīng)的程序進(jìn)行編譯,并部署到目標(biāo)硬件H上執(zhí)行(第6 行)。當(dāng)Strategy 為NL 時,Build_Run函數(shù)只會對A_Cand 和P_Cand 中的程序進(jìn)行編譯并運行得到相應(yīng)的編譯及運行信息A_Infos 和P_Infos。對A_Cand 編譯并運行是為了訓(xùn)練性能預(yù)測模型C,使其能夠鎖定一個更小的搜索空間。對P_Cand 編譯并運行是為了發(fā)現(xiàn)具有最短執(zhí)行時間的布局后程序。當(dāng)Strategy 為LR 時,Build_Run 函數(shù)只會對L_Cand 和P_Cand 中的程序進(jìn)行編譯并運行得到相應(yīng)的編譯及運行信息L_Infos 和P_Infos。當(dāng)Strategy 為LC 時,Build_Run 函數(shù)只會對P_Cand 中的程序進(jìn)行編譯并運行得到相應(yīng)的編譯運行信息P_Infos。

(5)根據(jù)策略Strategy 選擇有效的編譯及運行信息去訓(xùn)練/微調(diào)性能預(yù)測模型C(第7 行)。

(6)根據(jù)布局后程序集合P_Cand 的編譯運行信息P_Infos 中的運行時間信息更新搜索過程中發(fā)現(xiàn)的最優(yōu)布局后程序P?(第8～11 行)。

與Ansor 所采用的單一數(shù)據(jù)布局策略相比,ALAnsor 有2 個特點:(1)一個任務(wù)T的不同搜索輪次之間可以使用不同的靜態(tài)數(shù)據(jù)布局策略(由Select_Strategy 決定);(2)無論何種靜態(tài)數(shù)據(jù)布局策略,每一輪都會編譯并運行布局后程序集合P_Cand。

3.3 基于AL 策略的AL-Ansor

將自適應(yīng)靜態(tài)數(shù)據(jù)布局策略AL 應(yīng)用于Ansor的整個搜索過程中,便得到了改進(jìn)的Ansor——ALAnsor,如圖5 所示。和圖4 的主要差別在于程序采樣和編譯執(zhí)行之間的數(shù)據(jù)布局變換過程。在ALAnsor 框架中,靜態(tài)數(shù)據(jù)布局策略的確定不再依靠一次性的輸入,而由內(nèi)部的自適應(yīng)策略選擇函數(shù)決定。輸出總是執(zhí)行時間最短的布局后的程序P?∈PT,而不是依賴于輸入的靜態(tài)數(shù)據(jù)布局策略的程序X?。

圖5 AL-Ansor 工作原理示意圖

4 實驗

4.1 實驗設(shè)置及任務(wù)負(fù)載

本文采用的基準(zhǔn)Ansor 框架即TVM 庫(版本為e718f5a8)的auto_scheduler 模塊。根據(jù)Ansor 的輸入靜態(tài)數(shù)據(jù)布局策略的不同,分別將輸入NL、LR 和LC 策略的Ansor 記為NL-Ansor、LR-Ansor 和LCAnsor 以方便后續(xù)實驗展示。Ansor 和AL-Ansor 搜索框架的目標(biāo)硬件為32 核Intel Xeon Gold 6226R CPU(@2.90 GHz),搜索次數(shù)均設(shè)置為1024 次,因為這一次數(shù)已足夠使性能預(yù)測模型收斂[14]。實驗中AL-Ansor 所采用的自適應(yīng)策略選擇函數(shù)Select_Strategy 為一種簡單的預(yù)定義規(guī)則:依次進(jìn)行NL、LR和LC 策略,且3 種策略搜索次數(shù)的占比為5 ∶5 ∶6。

本文采用的任務(wù)負(fù)載是一些常見于深度卷積神經(jīng)網(wǎng)絡(luò)中的卷積塊(包含輸入補(bǔ)pad、進(jìn)行2D 卷積、加偏置以及做ReLU 激活),因為卷積是其主要計算負(fù)載。這些卷積塊用TVM 的topi 接口進(jìn)行描述,然后分別調(diào)用NL-Ansor、LR-Ansor、LC-Ansor 和ALAnsor 進(jìn)行搜索優(yōu)化。所選卷積任務(wù)的輸入、輸出數(shù)據(jù)如表2 所示。

表2 實驗負(fù)載

從2 個方面對4 種框架進(jìn)行比較。一個是不同框架下的搜索時間(或者編譯時間,因從計算圖到生成張量程序的搜索過程可以視為編譯);另一個是它們搜索出的最優(yōu)程序的運行時間。在進(jìn)行實驗時,獨占目標(biāo)服務(wù)器測量時間。對于最優(yōu)程序運行時間的測量,采取以下方法:對搜索得到的最優(yōu)程序進(jìn)行3 組測量取平均值,每組測量100 輪取中位數(shù),每輪運行至少500 ms(根據(jù)程序執(zhí)行時間,每輪的執(zhí)行次數(shù)在幾百至幾千的范圍),以最小化可能存在的測量誤差。

4.2 搜索時間對比

圖6 展示了NL-Ansor、LR-Ansor、LC-Ansor 和AL-Ansor 的搜索時間對比。其中AL-Ansor 的搜索時間被拆分為3 部分:AL-Ansor-base、AL-Ansor-exb和AL-Ansor-exr。AL-Ansor-exb 和AL-Ansor-exr 分別表示AL-Ansor 的自適應(yīng)靜態(tài)數(shù)據(jù)布局策略AL 所引入的程序P的編譯開銷和運行開銷,因為在搜索過程中無論選取哪一種布局策略,都要生成P并對其進(jìn)行編譯和運行,這一點通過對比圖4 和圖5 可以看出。但這一開銷只會發(fā)生在自適應(yīng)策略選擇函數(shù)返回策略NL 的搜索輪次;對于返回策略LR 和策略LC 的搜索輪次,程序P本就包含在程序X之中,因此不會產(chǎn)生額外的編譯開銷和運行開銷。ALAnsor-base 所示的時間則為把這部分開銷去掉的時間。

圖6 NL-Ansor、LR-Ansor、LC-Ansor 和AL-Ansor的搜索時間對比

從圖6 可以觀察到以下現(xiàn)象:

(1)AL-Ansor-base 在10 個任務(wù)上的搜索時間相較于NL-Ansor、LR-Ansor 和LC-Ansor 平均提升了13.86%～18.24%。這些主要來自于布局變換后的程序X的編譯和運行時間提升,而程序X是采樣程序A經(jīng)過靜態(tài)數(shù)據(jù)布局變換得到的,這表明策略AL下決定采樣質(zhì)量的性能預(yù)測模型C要優(yōu)于其他單一策略。

(2)相較于AL-Ansor-base,在每個搜索輪次引入對程序P的評估(AL-Ansor-exb 和AL-Ansor-exr)會帶來平均41.82%的時間開銷。如前所述,該部分開銷主要是由策略NL 所在搜索輪次引入的。這部分開銷中87.67%來自于程序P在目標(biāo)硬件上的運行過程(AL-Ansor-exr),這是為了估計P部署到目標(biāo)硬件后的真實運行時間,以在搜索過程中選出最優(yōu)的布局后程序P?。AL-Ansor-exr 的開銷可以根據(jù)需要去調(diào)整:如果希望對程序P的評估盡可能準(zhǔn)確,則可以將這部分時間延長;如果希望整個搜索過程快一些,則可以將部分時間適當(dāng)縮短。

(3)在10 個任務(wù)上,相較于NL-Ansor、LR-Ansor和LC-Ansor,AL-Ansor 平均上有15.95%～22.17%的搜索時間開銷。但是,由于AL-Ansor 內(nèi)的靜態(tài)數(shù)據(jù)布局策略是自適應(yīng)選擇的,而Ansor 中無法預(yù)先確定哪一種策略下能找到當(dāng)前任務(wù)的最優(yōu)張量程序。所以,為了能找到最優(yōu)的張量程序,Ansor 需要嘗試每一種策略。從這個角度看,AL-Ansor 的搜索時間相較于Ansor 平均會有2.5 倍的搜索時間提升。

總體來看,AL-Ansor 在搜索時間上仍然是優(yōu)于最優(yōu)靜態(tài)數(shù)據(jù)布局策略未知的Ansor 的。雖然相較于確定采用某一種靜態(tài)數(shù)據(jù)布局策略的Ansor(NLAnsor、LR-Ansor 或LC-Ansor),AL-Ansor 會有一定的搜索時間開銷。但是,對于一個確定的任務(wù)而言,搜索過程往往是一次性的,而搜索得到的最優(yōu)張量程序在部署到目標(biāo)硬件后則會反復(fù)運行使用。因此部署后運行時間的提升更為關(guān)鍵,部署前搜索時間的少量開銷是可以接受的。下面將分析4 種自動生成框架得到的最優(yōu)張量程序運行時間。

4.3 最優(yōu)程序運行時間對比

圖7 中依次展示了由NL-Ansor、LR-Ansor 和AL-Ansor 搜索出的最優(yōu)張量程序X,相對于LC-Ansor 搜索出的最優(yōu)張量程序Y的性能提升百分比y,即y由式(4)得到:

圖7 NL-Ansor、LR-Ansor 和AL-Ansor 相對于LC-Ansor的性能提升

其中timeof(Y,H)為程序Y在硬件H上的執(zhí)行時間。

在圖7 中,條形圖所處方位展示了對應(yīng)框架采取策略與LC 的優(yōu)劣:位于x軸上方表示對應(yīng)策略優(yōu)于策略LC,位于x軸下方則表示劣于策略LC。條形圖的高度則展示了對應(yīng)策略與LC 的差距:高度越大,表示對應(yīng)策略所得最優(yōu)程序與策略LC 所得最優(yōu)程序之間的運行時間絕對值之差越大。從圖7中可以得出如下結(jié)論。

(1)Ansor 原本支持的3 種獨立策略NL、LR 和LC 并沒有絕對的優(yōu)劣之分。這再一次表明了由用戶預(yù)先指定單一的靜態(tài)數(shù)據(jù)布局策略是不恰當(dāng)?shù)摹１? 展示了3 種策略的不同優(yōu)劣排序在任務(wù)負(fù)載中均曾出現(xiàn)。

表3 3 種策略優(yōu)劣排序在T1～T10 上的發(fā)生情況

(2)策略NL、策略LR 和策略AL 相較于LC 分別平均提升了2.59%、4.13%和16.59%的性能。策略NL 和策略LR 相較于策略LC 其性能平均上有提升,這一點雖然反直覺,但也進(jìn)一步印證了在策略LC 下,搜索空間AT和優(yōu)化空間PT之間的差異對于張量程序搜索框架的不利影響。

(3)策略AL 在絕大多數(shù)任務(wù)上優(yōu)于其他策略。只在任務(wù)T4 和T6 上存在2 處例外,在T4 上,策略AL 相對于LC 損失了1.88%的性能;而在T6 上,AL相對于NL 損失了1.31%的性能。這些微小的損失可能來自于搜索過程中為了跳出局部最小點的探索機(jī)制,而此過程是隨機(jī)不可控的。但值得注意的是,由此為策略AL 帶來的影響是很小的。除去這2 個異常點,策略AL 相對于策略NL 提升了0.98%～26.40%(T10 上最小,T2 上最大,見圖8),相對于策略LR 提升了3.79%～27.91%(T4 上最小,T5 上最大,見圖8),相對于策略LC 提升了3.58%～36.80%(T3 上最小,T7 上最大)。

圖8 AL-Ansor 相對于NL-Ansor 和LR-Ansor 的性能提升

圖8 單獨展示了AL-Ansor 相對于NL-Ansor 和LR-Ansor 的性能提升百分比。在所選的10 個任務(wù)負(fù)載上,策略AL 相對于NL 平均提升了13.81%,相對于LR 平均提升了12.41%。

綜合來看,相較于使用單一靜態(tài)數(shù)據(jù)布局策略的Ansor,采用自適應(yīng)靜態(tài)數(shù)據(jù)布局策略的AL-Ansor 框架更加合理和高效。

5 結(jié)論

本文研究了靜態(tài)數(shù)據(jù)布局在深度學(xué)習(xí)張量程序自動生成框架中的應(yīng)用。發(fā)現(xiàn)應(yīng)用最廣泛、前景最佳的Ansor 框架存在2 個問題:(1)只能指定單一的靜態(tài)數(shù)據(jù)布局策略(NL 或LR 或LC)進(jìn)行張量程序的搜索,而單一布局策略具有非最優(yōu)問題;(2)對于專用于靜態(tài)數(shù)據(jù)布局的策略LC,其搜索空間和優(yōu)化空間因其中程序的訪存特征差別而存在較大差異,導(dǎo)致性能預(yù)測模型不準(zhǔn)確的問題。這些問題都會導(dǎo)致搜索得到的張量程序仍有較大優(yōu)化空間。針對這些問題,本文提出了自適應(yīng)的靜態(tài)數(shù)據(jù)布局策略AL及基于此策略的深度學(xué)習(xí)張量程序自動生成框架AL-Ansor。實驗表明,在張量程序搜索過程中,自適應(yīng)地使用多種不同的靜態(tài)數(shù)據(jù)布局策略,不僅能夠緩解如何選擇合適布局策略的難題,也能夠減小搜索空間和優(yōu)化空間的差異,因此搜索生成的張量程序具有更好的性能。這為深度學(xué)習(xí)的高效部署與應(yīng)用帶來了廣泛的好處。

雖然本文討論的主要對象是靜態(tài)數(shù)據(jù),目標(biāo)硬件平臺選取的是CPU,但對于動態(tài)數(shù)據(jù)或者其他目標(biāo)硬件平臺(如GPU、TPU 等),本文所揭示的原理也是適用的。另外,本文所提出的自適應(yīng)策略選擇函數(shù)是決定該方法好壞的一個因素,如何選擇它仍是有待研究的課題。