Poly基本原理及卷積分析示例

**Poly基本原理介紹 **

考慮到許多讀者可能對(duì)Poly并不了解，而且許多Poly文獻(xiàn)讀起來(lái)也比較抽象，我們先簡(jiǎn)單介紹一下Poly的工作原理。我們力圖用最簡(jiǎn)單的代數(shù)與幾何描述來(lái)解釋Poly的基本原理。這部分內(nèi)容參考了文獻(xiàn)[12]的圖片，我們通過(guò)解讀這些圖片來(lái)解釋其中原理。

首先，我們假設(shè)有如圖1所示的一段簡(jiǎn)單的循環(huán)嵌套，其中N為常數(shù)。循環(huán)嵌套內(nèi)語(yǔ)句通過(guò)對(duì)A[i-1][j]和A[i][j-1]存儲(chǔ)數(shù)據(jù)的引用來(lái)更新A[i][j]位置上的數(shù)據(jù)。如果我們把語(yǔ)句在循環(huán)內(nèi)的每一次迭代實(shí)例抽象成空間上的一個(gè)點(diǎn)，那么我們可以構(gòu)造一個(gè)以(i,j)為基的二維空間，如圖2所示。圖中每個(gè)黑色的點(diǎn)表示寫(xiě)入A[i][j]的一次語(yǔ)句的迭代實(shí)例，從而我們可以構(gòu)造出一個(gè)所有黑色的點(diǎn)構(gòu)成的一個(gè)矩形，這個(gè)矩形就可以看作是二維空間上的一個(gè)Polyhedron（多面體），這個(gè)空間稱(chēng)為該計(jì)算的迭代空間。

圖1 一段簡(jiǎn)單的代碼示例

圖2 圖1示例代碼的迭代空間

我們可以用代數(shù)中的集合來(lái)對(duì)這個(gè)二維空間上的Polyhedron進(jìn)行表示，即{[i, j] : 1 <= i <= N - 1 and 1 <= j <= N – 1}，其中[i, j]是一個(gè)二元組，“:”后面的不等式表示這個(gè)集合的區(qū)間。我們可以給這個(gè)二元組做一個(gè)命名，叫做S，表示一個(gè)語(yǔ)句，那么這個(gè)語(yǔ)句的Polyhedron就可以表示成{S[i, j] : 1 <= i <= N - 1 and 1 <= j <= N – 1}。

由于語(yǔ)句S是先迭代i循環(huán)再迭代j循環(huán)，因此我們可以給語(yǔ)句S定義一個(gè)調(diào)度（順序），這個(gè)調(diào)度用映射表示，即{ S[i, j] -> [i, j] }，表示語(yǔ)句S[i, j] 先按i的順序迭代再按照j的順序迭代。

接下來(lái)，我們來(lái)分析語(yǔ)句和它訪存的數(shù)組之間的關(guān)系，在代數(shù)中我們用映射來(lái)表示關(guān)系。圖1中語(yǔ)句S對(duì)數(shù)組A進(jìn)行讀和寫(xiě)，那么我們可以用Poly來(lái)計(jì)算出S和A之間的讀訪存關(guān)系，可以表示成{ S[i, j] -> A[i - 1, j] : 1 <= i <= N -1 and 1 <= j <= N- 1; S[i, j] -> A[i, j - 1] : 1 <= i <= N - 1 and 1 <= j <= N -1 } 。同樣地，寫(xiě)訪存關(guān)系可以表示成{ S[i, j] -> A[i, j] : 1 <= i <= N - 1 and 1 <= j <= N -1 }。

基于這個(gè)讀寫(xiě)訪存關(guān)系，Poly就可以計(jì)算出這個(gè)循環(huán)嵌套內(nèi)的依賴關(guān)系，這個(gè)依賴關(guān)系可以表示成另外一種映射關(guān)系，即{ S[i, j] -> S[i, 1 + j] : 1 <= i <= N - 1 and 1 <= j <=N - 2; S[i, j] -> S[i + 1, j] : 1 <= i <= N - 2 and 1 <= j <= N- 1 }。

可以注意到，Poly對(duì)程序的表示都是用集合和映射來(lái)完成的。當(dāng)我們把語(yǔ)句實(shí)例之間的依賴關(guān)系用藍(lán)色箭頭表示在迭代空間內(nèi)時(shí)，就可以得到如圖3所示的形式。根據(jù)依賴的基本定理[13]，沒(méi)有依賴關(guān)系的語(yǔ)句實(shí)例之間是可以并行執(zhí)行的，而圖中綠色帶內(nèi)（對(duì)角線上）的所有點(diǎn)之間沒(méi)有依賴關(guān)系，所以這些點(diǎn)之間可以并行執(zhí)行。但是我們發(fā)現(xiàn)這個(gè)二維空間的基是(i, j)，即對(duì)應(yīng)i和j兩層循環(huán)，無(wú)法標(biāo)記可以并行的循環(huán)，因?yàn)檫@個(gè)綠色帶與任何一根軸都不平行。所以Poly利用仿射變換把基(i, j)進(jìn)行變換，使綠色帶能與空間基的某根軸能夠平行，這樣軸對(duì)應(yīng)的循環(huán)就能并行，所以我們可以將圖3所示的空間轉(zhuǎn)化成如圖4所示的形式。

此時(shí)，語(yǔ)句S的調(diào)度就可以表示成{ S[i, j] -> [i + j, j]}的形式。所以Poly的變換過(guò)程也稱(chēng)為調(diào)度變換過(guò)程，而調(diào)度變換的過(guò)程就是變基過(guò)程、實(shí)現(xiàn)循環(huán)變換的過(guò)程。

圖3 帶依賴關(guān)系的迭代空間

圖4 變基之后的迭代空間

圖4中綠色帶和j軸平行，這樣在代碼中表示起來(lái)就方便了。我們說(shuō)Poly做循環(huán)變換的過(guò)程就是將基(i, j)變成(i + j, j)的一個(gè)過(guò)程，也就是說(shuō)，Poly的底層原理就是求解一個(gè)系數(shù)矩陣，這個(gè)系數(shù)矩陣能夠?qū)⑾蛄?i, j)轉(zhuǎn)換成向量(i + j, j)。

根據(jù)這樣的調(diào)度，Poly就可以利用它的代碼生成器，生成如圖5所示的代碼。此時(shí)，內(nèi)層循環(huán)就可以并行了。（注：這里示意的是“源到源”翻譯的Poly編譯器，也就是Poly生成的代碼還需要交給基礎(chǔ)編譯器如GCC、ICC、LLVM等編譯成機(jī)器碼才能運(yùn)行。也有內(nèi)嵌在基礎(chǔ)編譯中的Poly工具。）

圖5 Poly變換后生成的代碼

當(dāng)然，我們這里舉的例子是一個(gè)很簡(jiǎn)單的例子，在實(shí)際應(yīng)用中還有很多復(fù)雜的情況要考慮。Poly幾乎考慮了所有的循環(huán)變換，包括Interchange（交換）、Skewing/Shifting（傾斜/偏移）、Reversal（反轉(zhuǎn)）、Tiling/Blocking（分塊）、Stripe-mining、Fusion（合并）、Fission/Distribution（分布）、Peeling（剝離）、Unrolling（展開(kāi)）、Unswitching、Index-set splitting、Coalescing/Linearization等，圖6～8[14]中給出了幾種Poly中實(shí)現(xiàn)的循環(huán)變換示意圖，右上角的代碼表示原輸入循環(huán)嵌套，右下角的代碼表示經(jīng)過(guò)Poly變換后生成的代碼。圖中左邊的集合和映射關(guān)系的含義分別為：J代表原程序語(yǔ)句的迭代空間，S表示輸入程序時(shí)的調(diào)度，T表示目標(biāo)調(diào)度，ST就是Poly要計(jì)算的調(diào)度變換。

圖6 Poly中skewing變換示意圖

圖7 Poly中Fusion變換示意圖

圖8 Poly中Tiling變換示意圖

深度學(xué)習(xí)應(yīng)用的Poly優(yōu)化

讓我們以圖9中所示的二維卷積運(yùn)算（矩陣乘法）為例來(lái)簡(jiǎn)單介紹Poly是如何優(yōu)化深度學(xué)習(xí)應(yīng)用的。

圖9 一個(gè)2D卷積示例

Poly會(huì)將循環(huán)嵌套內(nèi)的計(jì)算抽象成一個(gè)語(yǔ)句。例如圖9中S1語(yǔ)句表示卷積初始化，S2代表卷積歸約；而S0和S3則分別可以看作卷積操作前后的一些操作，比如S0可以想象成是量化語(yǔ)句，而S3可以看作是卷積后的relu操作等。

為了便于理解，我們以CPU上的OpenMP程序?yàn)槟繕?biāo)對(duì)圖9中的示例進(jìn)行變換。Poly在對(duì)這樣的二維卷積運(yùn)算進(jìn)行變換的時(shí)候，會(huì)充分考慮程序的并行性和局部性。如果我們對(duì)變換后的程序并行性的要求大于局部性的要求，那么Poly會(huì)自動(dòng)生成如圖10所示的OpenMP代碼；如果我們對(duì)局部性的要求高于并行性，那么Poly會(huì)自動(dòng)生成如圖11所示的OpenMP代碼。（注：不同的Poly編譯器生成的代碼可能會(huì)因采用的調(diào)度算法、編譯選項(xiàng)、代碼生成方式等因素而不同。）

圖10 Poly生成的OpenMP代碼——并行性大于局部性

圖11 Poly生成的OpenMP代碼——局部性大于并行性

通過(guò)對(duì)比圖10和圖11，兩種生成的代碼采用的循環(huán)fusion（合并）策略不同：圖10中所示的代碼采用了({S0}, {S1, S2, S3})的合并策略，圖11中生成的代碼則使用了({S0,S1, S2, S3})的合并策略，但是必須通過(guò)對(duì)S2向右偏移99次、S3向右偏移148次，以及循環(huán)層次的interchange（交換）來(lái)實(shí)現(xiàn)這樣的合并。顯然，圖11所示的代碼局部性更好。而并行性上，仔細(xì)研究后不難發(fā)現(xiàn)，圖11生成的代碼中，只有最外層c0循環(huán)是可以并行的，而圖10代碼中，S0語(yǔ)句的c0、c1循環(huán)都可以并行，并且包含S1、S2、 S3三條語(yǔ)句的循環(huán)嵌套的c0、c1循環(huán)也都可以并行，相對(duì)于圖11代碼，圖10生成的代碼可并行循環(huán)的維度更多。

當(dāng)然，在面向CPU生成OpenMP代碼時(shí)，多維并行的優(yōu)勢(shì)沒(méi)有那么明顯，但是當(dāng)目標(biāo)架構(gòu)包含多層并行硬件抽象時(shí)，圖9中的代碼能夠更好地利用底層加速芯片。例如，當(dāng)面向GPU生成CUDA代碼時(shí)，而圖10對(duì)應(yīng)的CUDA代碼（如圖12所示）由于合并成了兩個(gè)部分，會(huì)生成2個(gè)kernel，但是每個(gè)kernel內(nèi)c0維度的循環(huán)被映射到GPU的線程塊上，而c1維度的循環(huán)被映射到GPU的線程上；圖11對(duì)應(yīng)的CUDA代碼（如圖13所示）只有1個(gè)kernel，但是只有c0維度的循環(huán)被映射到GPU的線程塊和線程兩級(jí)并行抽象上。為了便于閱讀，我們并未開(kāi)啟GPU上shared memory和private memory自動(dòng)生成功能。從圖中也不難發(fā)現(xiàn)，Poly也可以自動(dòng)生成線程之間的同步語(yǔ)句。（注：圖中循環(huán)分塊大小為32，圖12中線程塊上線程布局為3216，圖13中為324*4。）

圖12 Poly生成的CUDA代碼——并行性大于局部性

圖13 Poly生成的CUDA代碼——局部性大于并行性

值得注意的是，為了充分挖掘程序的并行性和局部性，Poly會(huì)自動(dòng)計(jì)算出一些循環(huán)變換來(lái)實(shí)現(xiàn)有利于并行性和局部性的變換。例如，為了能夠達(dá)到圖11和圖13中所有語(yǔ)句的合并，Poly會(huì)自動(dòng)對(duì)S2和S3進(jìn)行shifting（偏移）和interchange（交換）。