Poly在深度學習領域中發揮的作用

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

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

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

圖14 GPU架構示意圖

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

圖15 TPU架構示意圖

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

圖16 TVM調度示例

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

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

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

圖17 昇騰AI處理器的DaVinci Core架構示意圖

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

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

圖18 圖10對應的schedule tree表示

圖19 圖11對應的schedule tree表示

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

圖20 圖18分塊之后的schedule tree

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

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

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

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

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

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