Poly在深度學(xué)習(xí)領(lǐng)域中發(fā)揮的作用

從對上層應(yīng)用的約束角度來看，作為一種通用程序設(shè)計語言的編譯優(yōu)化模型，Poly本身對應(yīng)用是敏感的，只能處理滿足一定約束條件的、規(guī)則的應(yīng)用。Poly要求被分析的應(yīng)用中，循環(huán)邊界、數(shù)組下標都是仿射表達式，而且控制流必須是靜態(tài)可判定的，我們暫且把這種對應(yīng)用的要求稱為靜態(tài)仿射約束。實際上，對于通用語言而言，靜態(tài)仿射約束的限制對程序的要求不算低，但是深度學(xué)習(xí)領(lǐng)域的大部分核心計算卻恰好滿足這種靜態(tài)仿射約束，所以許多深度學(xué)習(xí)編譯軟件棧利用Poly來實現(xiàn)循環(huán)優(yōu)化。

而從充分發(fā)揮底層AI芯片架構(gòu)的能力角度來講，Poly也非常適合，這得益于Poly能夠自動判定和實現(xiàn)上層應(yīng)用中循環(huán)的tiling/blocking（分塊）變換并自動將軟件循環(huán)映射到并行硬件上。本系列文章第一篇中圖12和圖13就是Poly在GPU上自動實現(xiàn)分塊并將分塊后對應(yīng)的循環(huán)維度映射到GPU的線程塊和線程兩級并行硬件抽象上的實例。

為什么Poly需要自動實現(xiàn)分塊？這是由底層AI芯片的架構(gòu)導(dǎo)致的。以GPU為例，圖14[15]所示是GPU的架構(gòu)示意圖。每個GPU上擁有自己的全局緩存（Global/Device Memory），然后每個線程塊也有自己的局部緩存（Shared/Local Memory）。緩存越靠近計算單元，訪存的速度越快，但是緩存空間越小。因此，當計算數(shù)據(jù)量大于緩存空間的時候，就需要通過將原來的數(shù)據(jù)進行分塊的方式存儲到緩存上，以此來適應(yīng)目標架構(gòu)的硬件特征。

圖14 GPU架構(gòu)示意圖

而專用AI芯片的架構(gòu)可能更復(fù)雜，如圖15[16]所示是TPU v2和TPU v3的架構(gòu)示意圖，每個TPU有多種不同類型的計算單元，包括標量、向量以及矩陣計算單元，這些不同的計算單元對應(yīng)地可能會有各自不同的緩存空間，這就給分塊提出了更高的要求。只有通過對應(yīng)用的正確分塊才能充分利用好芯片上的架構(gòu)特征。

圖15 TPU架構(gòu)示意圖

當前一部分深度學(xué)習(xí)編譯軟件棧采用了手工調(diào)度和映射的方式來將上層應(yīng)用部署到底層芯片上。以TVM為例，圖16[17]中給出了一個TVM的調(diào)度示例。其中，調(diào)度過程首先將計算s進行分塊（對應(yīng)圖16中的split操作），然后將分塊后的維度映射到GPU的線程塊和線程上（對應(yīng)圖16中的bind操作）。

圖16 TVM調(diào)度示例

這部分工作需要由熟悉底層芯片架構(gòu)的人員來編寫，并且要人工分析分塊的合法性，映射也需要手工完成。而Poly的作用就是將上述手工調(diào)度的過程自動實現(xiàn)。為了實現(xiàn)自動調(diào)度，許多深度學(xué)習(xí)編譯軟件棧開始采用Poly來實現(xiàn)上述功能。那么，Poly在深度學(xué)習(xí)軟件棧上發(fā)揮的作用如何呢？

首先，Poly能夠計算精確的數(shù)據(jù)流信息。Poly通過將傳統(tǒng)的編譯器中語句之間的依賴關(guān)系細化到語句實例的粒度，分析的結(jié)果比傳統(tǒng)的方法更精確。計算精確的數(shù)據(jù)流信息有以下三點好處。

1.計算精確的緩存搬移數(shù)據(jù)量。Poly不僅能自動計算出從管理核心（如CPU）到加速芯片（如GPU）之間傳輸?shù)臄?shù)據(jù)總量，還負責(zé)計算加速芯片上多級緩存之間的數(shù)據(jù)搬移總量，例如從GPU的global memory到shared memory上的數(shù)據(jù)搬移?！按鎯Α眴栴}給我們揭示了一個道理：數(shù)據(jù)搬移是程序性能提升的關(guān)鍵，尤其是現(xiàn)在市場上越來越復(fù)雜的多級緩存架構(gòu)上，數(shù)據(jù)總量計算是否精確對程序性能的影響更加明顯。
2.降低內(nèi)存空間使用。通過計算精確的數(shù)據(jù)流信息，Poly可以計算出臨時tensor變量，這些臨時變量的聲明可在對應(yīng)的緩存級別上實現(xiàn)，從而降低加速芯片上數(shù)據(jù)的訪存開銷。例如，圖9所示的代碼段中，tensor b就可以看作是一個臨時tensor變量。
3.自動實現(xiàn)緩存上的數(shù)據(jù)部署。以華為剛公布的昇騰AI處理器芯片為例，如圖17[18]是該芯片的AI Core架構(gòu)示意圖。其中，UnifiedBuffer（輸出緩沖區(qū)）和L1 Buffer（輸入緩沖區(qū)）是低級緩存，離計算單元較遠；BufferA L0/B L0/C L0是高級緩存，靠近計算單元。在低級緩存上，Poly可以借助標記節(jié)點，將不同計算單元所需的數(shù)據(jù)分別流向UnifiedBuffer和L1 Buffer；同時，當數(shù)據(jù)到達高級緩存時，Poly仍然可以借助標記節(jié)點將數(shù)據(jù)自動部署到BufferA L0/B L0/C L0。（注：這里描述的是如何通過Poly來實現(xiàn)這樣的數(shù)據(jù)分流，只是為了說明Poly能夠?qū)崿F(xiàn)這樣的自動數(shù)據(jù)部署功能，與具體實現(xiàn)無關(guān)。至于昇騰AI處理器芯片的編譯團隊是否使用了Poly，或者是否使用了這種方法來實現(xiàn)數(shù)據(jù)的自動部署還請以官方公布為準。）

圖17 昇騰AI處理器的DaVinci Core架構(gòu)示意圖

其次，Poly能夠?qū)崿F(xiàn)幾乎全部的循環(huán)變換。Poly通過仿射函數(shù)來實現(xiàn)幾乎所有循環(huán)變換及其組合，這種仿射函數(shù)的計算過程不僅要考慮應(yīng)用程序的并行性和局部性，還要考慮底層加速芯片的硬件特征。從循環(huán)變換角度來講，Poly對編譯軟件棧的貢獻包括以下幾個方面。

1.Poly中的調(diào)度算法[19-22]能夠根據(jù)依賴關(guān)系分析的結(jié)果自動計算出變換后循環(huán)的并行性、循環(huán)維度是否可以實施分塊等特征，這些特征為后面硬件上的計算任務(wù)分配、緩存上的循環(huán)變換提供了理論依據(jù)。（這些信息保存在band節(jié)點（下面會介紹）的屬性中）而部分循環(huán)變換如skewing/shifting（傾斜/偏移）、interchange（交換）等都可以在調(diào)度階段自動完成。我們?nèi)匀灰詧D9中所示的例子來說明。對于圖10生成的代碼，Poly計算出來的調(diào)度用其中間表示（schedule tree）[23]后得到的結(jié)果如圖18所示，而圖11生成的代碼對應(yīng)的調(diào)度如圖19所示。（注：為方便說明，這里的schedule tree可能和實際在Poly中使用的有所不同，我們只是為了更直觀地表示schedule tree的表示方式。）其中，domain節(jié)點包含所有的語句實例集合，sequence節(jié)點表示其子節(jié)點按序執(zhí)行，而“[]”包含的節(jié)點稱為band節(jié)點，可以想象成循環(huán)。

圖18 圖10對應(yīng)的schedule tree表示

圖19 圖11對應(yīng)的schedule tree表示

2.自動實現(xiàn)深度學(xué)習(xí)應(yīng)用中最關(guān)鍵的tiling/blocking（分塊）和fusion（合并）變換。分塊的目的是為了充分利用加速芯片上的緩存，而合并的目的是為了生成更多的臨時緩存變量，降低訪存開銷。而且，Poly通過數(shù)學(xué)變換，能夠自動實現(xiàn)更復(fù)雜的、手工難以實現(xiàn)的分塊形狀[6, 24-26]。其中，合并可根據(jù)調(diào)度選項在調(diào)度變換過程實現(xiàn)，分塊則是在調(diào)度變換之后根據(jù)循環(huán)維度是否可分塊等特征來實現(xiàn)。如圖18和19就是根據(jù)不同的編譯選項實現(xiàn)的合并策略對應(yīng)的schedule tree，其中合并已經(jīng)通過sequence節(jié)點實現(xiàn)，而分塊的實現(xiàn)在Poly上很簡單，只需要將band節(jié)點中的仿射函數(shù)進行修改就可以得到分塊對應(yīng)的schedule tree。如圖20是圖18經(jīng)過分塊之后的調(diào)度樹，圖19的分塊也可以同樣的方式得到，我們就不再贅述了。

圖20 圖18分塊之后的schedule tree

3.通過代碼生成方式自動實現(xiàn)不改變語句順序、但只改變循環(huán)結(jié)構(gòu)的變換。這類循環(huán)變換包括peeling（剝離）、unrolling（展開）等。因為這些循環(huán)變換不改變語句的執(zhí)行順序，而只是對循環(huán)的結(jié)構(gòu)進行修改來實現(xiàn)。這些循環(huán)變換對特殊加速芯片上的代碼生成有十分重要的作用，例如一些架構(gòu)可能并不喜歡循環(huán)上下界中有min/max這樣的操作，此時就需要實現(xiàn)這類循環(huán)變換。這類循環(huán)變換可以通過在schedule tree中的band節(jié)點上添加特殊的options屬性來實現(xiàn)。（注：我們的圖中沒有標出options，但實際使用的schedule tree中有options，而options中的內(nèi)容是一個集合或者映射表達式，計算起來也很方便。）

第三，Poly能夠自動實現(xiàn)存儲系統(tǒng)的管理。在越來越復(fù)雜的加速芯片架構(gòu)上，復(fù)雜的存儲系統(tǒng)是實現(xiàn)芯片上計算部署的難點，即便是硬件開發(fā)人員來手工實現(xiàn)程序在存儲結(jié)構(gòu)上的管理，也是一個十分耗時且易出錯的任務(wù)。而Poly借助中間表示自動實現(xiàn)了在多級緩存結(jié)構(gòu)上的存儲管理[27]，使得底層優(yōu)化和硬件開發(fā)人員從這些瑣碎的工作中脫離出來。這種自動管理存儲系統(tǒng)的實現(xiàn)包括以下兩個方面。

1.自動計算緩存之間傳遞數(shù)據(jù)需要插入的位置。由于數(shù)據(jù)傳輸指令在原程序中是不存在的，所以Poly要能夠?qū)崿F(xiàn)這種從無到有的指令生成過程，并且正確計算出相應(yīng)的位置。Poly借助schedule tree上的特殊節(jié)點和仿射函數(shù)，實現(xiàn)了數(shù)據(jù)傳輸指令位置的準確計算和自動插入。
2.自動生成數(shù)據(jù)傳輸指令的循環(huán)信息。確定數(shù)據(jù)傳輸指令的位置后，Poly可以根據(jù)數(shù)學(xué)關(guān)系計算出當前指令所在循環(huán)的層次和維度信息，并自動為數(shù)據(jù)傳輸指令計算對應(yīng)的調(diào)度關(guān)系，然后交給后端代碼生成器生成代碼。

如圖21所示，是圖20的schedule tree經(jīng)過插入特殊的extension節(jié)點之后，得到的帶有數(shù)據(jù)傳輸指令的中間表示。其中，kernel0和kernel1分別對應(yīng)圖20中最上面sequence節(jié)點下的兩棵子樹，而to_device_B和to_device_a表示從CPU的內(nèi)存上拷貝tensor B和a到GPU的global memory，這兩個語句在計算之前。from_device_c表示將GPU上的tensor c從global memory傳輸回CPU內(nèi)存上，這個語句在計算之后。Poly并沒有傳輸tensor b，而是在GPU的global memory上創(chuàng)建和使用了tensor b。（注：to_device_B和to_device_a也可以顛倒順序執(zhí)行，為了便于說明我們在這里假設(shè)按序執(zhí)行。）

圖21 圖20的schedule tree插入數(shù)據(jù)傳輸指令之后的中間表示

最后，Poly還能夠自動計算出變換之后循環(huán)到硬件上的映射。在提供多級并行硬件抽象和按計算的類型提供不同計算單元的加速芯片上，軟件循環(huán)要實現(xiàn)到硬件上的映射，而這種映射關(guān)系也可以借助Poly的仿射函數(shù)和schedule tree上的標記來自動實現(xiàn)。這可以通過在kernel0和kernel1的子樹內(nèi)的band節(jié)點上添加特殊標記來實現(xiàn)。（注：圖中未標出。）