AI芯片上利用Poly進(jìn)行軟硬件優(yōu)化的一些問(wèn)題

1. 關(guān)于IR

我們?cè)谶@篇文章中主要關(guān)注的是Poly的IR。在之前的內(nèi)容中，我們主要基于Poly傳統(tǒng)的schedule tree表示介紹了如何實(shí)現(xiàn)AI芯片上的軟硬件優(yōu)化。Google MLIR[29]針對(duì)schedule tree的不足，提出了一種簡(jiǎn)化的Poly IR。以圖22中所示的代碼為例，用schedule tree對(duì)該部分代碼進(jìn)行表示，可以得到如圖23所示的IR形式，而MLIR則是將其表示成圖24的形式。（注：圖23中的schedule tree與圖18-21表示的內(nèi)容一樣，只不過(guò)這里用文字的形式表示出圖的內(nèi)容，圖23中schedule后的標(biāo)量維度可以對(duì)應(yīng)為圖18-21中的sequence節(jié)點(diǎn)。）

圖22 另一個(gè)簡(jiǎn)單用例

圖23 圖22代碼的scheduletree表示

圖24MLIR對(duì)圖22中的簡(jiǎn)化Poly IR

通過(guò)對(duì)比可以看出，MLIR的Poly IR對(duì)循環(huán)進(jìn)行了更加顯式的表達(dá)，而省略了schedule tree中的domain、schedule等信息。這種Poly IR簡(jiǎn)化了Poly在實(shí)現(xiàn)循環(huán)變換之后的代碼生成過(guò)程。例如，在實(shí)現(xiàn)傾斜變換時(shí)，MLIR可以直接基于圖24生成如圖25所示的Poly IR。

圖25 經(jīng)過(guò)傾斜后的MLIR的Poly IR

但是相比于傳統(tǒng)的schedule tree表示，循環(huán)變換的過(guò)程更復(fù)雜了。在schedule tree上，循環(huán)變換，如這里的傾斜，可以直接修改schedule的仿射函數(shù)來(lái)實(shí)現(xiàn)，可參考圖6；但在MLIR中卻要對(duì)應(yīng)地修改顯式表達(dá)的循環(huán)變量及對(duì)應(yīng)的下標(biāo)信息。

筆者對(duì)這種方案存在兩點(diǎn)疑問(wèn)。一是在Poly的整個(gè)流程中，雖然代碼生成也比較復(fù)雜，但是循環(huán)變換的時(shí)間開(kāi)銷可能比代碼生成的開(kāi)銷更高，雖然簡(jiǎn)化了代碼生成，但是循環(huán)變換更加復(fù)雜，不知道這樣的代價(jià)是否值當(dāng)？當(dāng)然，Google MLIR團(tuán)隊(duì)集結(jié)了編譯領(lǐng)域最頂級(jí)的專家和最熟悉Poly的研究團(tuán)隊(duì)，筆者相信他們提出這種簡(jiǎn)化的Poly IR肯定是經(jīng)過(guò)深思熟慮的，可能后期需要向Google MLIR的Poly專家再請(qǐng)教之后才可能解答這個(gè)疑惑。二是這種簡(jiǎn)化的Poly IR是為了簡(jiǎn)化從ML Function到CFGFunction的代碼生成過(guò)程，那如果Poly變換之后的輸出不是基于LLVM IR的框架是否還有必要采用這種簡(jiǎn)化的Poly IR？畢竟，目前深度學(xué)習(xí)框架的“IR之爭(zhēng)”還沒(méi)有結(jié)束。

2. 關(guān)于循環(huán)變換

Poly的調(diào)度算法[19-22]基于線性整數(shù)規(guī)劃來(lái)求解新的調(diào)度仿射函數(shù)，而這個(gè)過(guò)程中會(huì)考慮到幾乎所有的循環(huán)變換及多個(gè)循環(huán)變換之間的組合。以圖26[30]為例是線性整數(shù)規(guī)劃問(wèn)題求解的示意圖，其中藍(lán)色的點(diǎn)表示整個(gè)空間上的整數(shù)，而圖中的斜邊可以看作是循環(huán)邊界等信息給出的約束，這些約束構(gòu)成了一個(gè)可行解區(qū)間（圖中綠色部分）。那么調(diào)度問(wèn)題可以抽象成在這個(gè)綠色的解空間內(nèi)尋找一個(gè)目標(biāo)問(wèn)題（紅色箭頭）的最優(yōu)解（在Poly 里，就是尋找按字典序最小的整數(shù)解）。

圖26 線性整數(shù)規(guī)劃問(wèn)題示意圖

但是，一個(gè)重要的前提條件是Poly是面向通用程序設(shè)計(jì)語(yǔ)言的編譯數(shù)學(xué)模型，如果我們將Poly應(yīng)用到如深度學(xué)習(xí)這樣的特定領(lǐng)域，是否需要考慮和通用語(yǔ)言一樣的循環(huán)變換集合？一個(gè)很簡(jiǎn)單的例子是對(duì)于一個(gè)卷積算子，卷積核的循環(huán)嵌套會(huì)嵌套在輸入圖像的循環(huán)嵌套內(nèi)部，而卷積核的循環(huán)維度范圍可能會(huì)比輸入圖像的循環(huán)維度范圍小很多。當(dāng)Poly計(jì)算新的調(diào)度時(shí)，輸入圖像的循環(huán)維度和卷積核的循環(huán)維度可能發(fā)生傾斜變換，但這種傾斜似乎對(duì)卷積計(jì)算后面的變形、代碼生成等問(wèn)題都不太友好。所以，Poly調(diào)度算法考慮的循環(huán)變換，在深度學(xué)習(xí)領(lǐng)域是否都需要，還是只需要比較核心的、對(duì)性能提升比較關(guān)鍵的幾種變換？如果我們減少了需要考慮的循環(huán)變換個(gè)數(shù)及其組合，也就是說(shuō)在圖26中我們縮小了可行解的區(qū)間，那么求解起來(lái)是不是會(huì)更高效一些？

如果上面問(wèn)題的答案是肯定的，那么筆者認(rèn)為目前而言Poly能實(shí)現(xiàn)的循環(huán)變換中，對(duì)深度學(xué)習(xí)應(yīng)用最關(guān)鍵的循環(huán)變換應(yīng)該是分塊和合并（可能有待商榷）。假設(shè)只考慮分塊和合并這兩種循環(huán)變換，這種情況下問(wèn)題似乎簡(jiǎn)單一些。但是編譯優(yōu)化中還有一個(gè)比較關(guān)鍵的問(wèn)題就是如何決定實(shí)現(xiàn)的循環(huán)變換的順序。是先做合并后做分塊，還是先做分塊再做合并？事實(shí)上，對(duì)于循環(huán)變換的順序判定問(wèn)題，傳統(tǒng)的Poly中間表示沒(méi)有給出明確的答案，而不幸的是，MLIR也沒(méi)有解決這個(gè)問(wèn)題。當(dāng)然，這只是極簡(jiǎn)情況下的假設(shè)。只有分塊和合并顯然是不夠的，因?yàn)檠h(huán)變換后的代碼生成還要借助distribution（分布）來(lái)保證向量化等問(wèn)題的發(fā)掘。

3. 關(guān)于分塊

我們前面針對(duì)GPU、TPU以及昇騰910的架構(gòu)都進(jìn)行了分析，（多級(jí)）緩存是目前市場(chǎng)上AI芯片采用的架構(gòu)趨勢(shì)，而專用AI芯片如TPU、昇騰910等專用計(jì)算單元的設(shè)計(jì)似乎也引領(lǐng)了AI芯片的另一種方向?？梢哉f(shuō)，在當(dāng)前的AI芯片上，分塊是軟件棧必須實(shí)現(xiàn)的一種優(yōu)化手段了。

而針對(duì)分塊這一種變換而言，仍然還有很多值得研究的問(wèn)題。比如，以圖27[31]中的二維卷積為例，卷積核（kernel）通過(guò)在輸入圖像（input）上進(jìn)行“滑動(dòng)”來(lái)計(jì)算輸出圖像（output）的結(jié)果。當(dāng)卷積核針對(duì)輸入圖像的某一個(gè)像素點(diǎn)（圖中深藍(lán)色的方塊）進(jìn)行計(jì)算時(shí)，需要通過(guò)對(duì)其周邊特定區(qū)域的像素點(diǎn)（圖中淺藍(lán)色的方塊）進(jìn)行加權(quán)（即卷積操作）后得到輸出圖像上的一個(gè)像素點(diǎn)（圖中紅色方塊）。由于卷積核需要在輸入圖像上進(jìn)行滑動(dòng)，這種滑動(dòng)的過(guò)程在大多數(shù)情況下會(huì)導(dǎo)致輸入圖像的數(shù)據(jù)被多次訪問(wèn)。以圖27為例，當(dāng)滑動(dòng)的步長(zhǎng)（stride）為1時(shí)，除輸入圖像上第一列和最后一列的像素點(diǎn)之外所有的像素點(diǎn)都會(huì)被重復(fù)計(jì)算。如果我們按照卷積核的大小對(duì)輸入圖像進(jìn)行分塊，那么分塊之后輸入圖像的每個(gè)分塊之間都會(huì)存在overlap（數(shù)據(jù)重疊）問(wèn)題。如何利用Poly在深度學(xué)習(xí)應(yīng)用中自動(dòng)實(shí)現(xiàn)這種滿足數(shù)據(jù)重疊的分塊？

圖27 卷積操作示例

一種方式是采用PolyMage[6]類似的方法利用Poly的調(diào)度來(lái)求解這樣的overlap的區(qū)間，但是這種方式有可能會(huì)導(dǎo)致過(guò)多的冗余計(jì)算，而且用調(diào)度來(lái)求解分塊的形狀在某種程度上會(huì)使Poly的過(guò)程變得更加復(fù)雜，代碼生成亦如是；另一種方式是在schedule tree上利用特殊的節(jié)點(diǎn)來(lái)實(shí)現(xiàn)，但是目前這種方式的代碼實(shí)現(xiàn)都還沒(méi)有公開(kāi)。

另外一個(gè)問(wèn)題是在上一期的內(nèi)容中有讀者提問(wèn)到關(guān)于分塊和冗余計(jì)算的問(wèn)題。冗余計(jì)算的確會(huì)給性能的提升帶來(lái)一定的影響，但是這種冗余計(jì)算的引入是為了實(shí)現(xiàn)分塊之間的并行。我們?cè)谇拔奶岬竭^(guò)，并行性和局部性有的時(shí)候是沖突的，為了達(dá)到兩者之間的平衡，往往是需要作出一些其它的犧牲來(lái)達(dá)到目的[32]。而更重要的是，這種帶有冗余計(jì)算的分塊形狀是目前幾種分塊形狀中，實(shí)現(xiàn)降低內(nèi)存開(kāi)銷最有效的一種形狀。以圖28[6]為例，圖中列出來(lái)三種不同的分塊形狀，其中最左側(cè)的梯形分塊引入了冗余計(jì)算，但是這種分塊在一次分塊計(jì)算完成（水平方向）后，分塊內(nèi)需要傳遞給下一次計(jì)算的活躍變量（紅色圓圈）總數(shù)最少，而其它形狀如中間的分裂分塊和最右側(cè)的平行四邊形分塊剩余的活躍變量總數(shù)都很多，無(wú)法實(shí)現(xiàn)有效降低內(nèi)存開(kāi)銷的目的。（注：圖中未分析鉆石分塊[25]和六角形分塊[26]，但是這兩種分塊可以看作是分裂分塊的一種特殊形式。）

圖28 不同分塊形狀分塊內(nèi)活躍變量的分析對(duì)比

除此之外，分塊面臨的問(wèn)題還有很多。比如，Poly中實(shí)現(xiàn)的分塊都是計(jì)算的分塊，而數(shù)據(jù)分塊只是通過(guò)計(jì)算分塊和計(jì)算與數(shù)據(jù)之間的仿射函數(shù)來(lái)計(jì)算得到，這種結(jié)果能夠保證數(shù)據(jù)的分塊是最優(yōu)的嗎？在分布式結(jié)構(gòu)上呢？而針對(duì)TPU和昇騰910等專用AI加速芯片，多級(jí)分塊應(yīng)該如何實(shí)現(xiàn)才能更好的發(fā)揮這些加速芯片的特征呢？

4. 關(guān)于合并

循環(huán)的合并是一個(gè)挖掘局部性的過(guò)程。我們?nèi)匀幌霃?qiáng)調(diào)的問(wèn)題是，局部性和并行性是加速芯片上兩個(gè)非常重要的變換目標(biāo)，但是這兩個(gè)目標(biāo)有的時(shí)候是互斥的，就如我們?cè)趫D10和圖11中所示的例子一樣。合并的循環(huán)越多，破壞計(jì)算并行性的可能性越大；而如果要保持計(jì)算的并行性，可能就要放棄一些循環(huán)的合并。

然而，在不同的架構(gòu)上哪些合并是最優(yōu)的，似乎靜態(tài)判定是不太可能的。就如我們?cè)诘谝徊糠址窒淼哪菢樱?a href="http://www.ttokpm.com/v/tag/132/" target="_blank">CPU上生成OpenMP代碼可能一層并行就足矣，這時(shí)局部性的效果可能就比并行性的效果更好；而在GPU上，由于有兩層并行硬件的抽象，可能并行性的收益比局部性的效果更佳。所以，現(xiàn)在許多深度學(xué)習(xí)軟件棧也采用了Auto-tuning的方式來(lái)通過(guò)實(shí)際的多次運(yùn)行來(lái)判定哪種策略是最優(yōu)的。然而，即便是Auto-tuning的方式，能夠保證遍歷到所有的合并形式嗎？如何選擇一個(gè)合適的合并策略，是必須要通過(guò)調(diào)優(yōu)的方式來(lái)確定嗎？利用靜態(tài)分析的方式來(lái)遍歷所有合并策略的工作也研究過(guò)[33]，但是這種動(dòng)態(tài)規(guī)劃的方式是不是又會(huì)帶來(lái)時(shí)間復(fù)雜度的問(wèn)題？

5. 關(guān)于Poly時(shí)間復(fù)雜度和對(duì)Poly的擴(kuò)展

關(guān)于Poly的時(shí)間復(fù)雜度問(wèn)題，我們?cè)谏衔闹幸呀?jīng)提到Poly的調(diào)度實(shí)質(zhì)是線性整數(shù)規(guī)劃問(wèn)題的求解過(guò)程，而實(shí)際上Poly的代碼生成過(guò)程也會(huì)涉及到線性整數(shù)規(guī)劃問(wèn)題的求解。我們?cè)谟懻撋疃葘W(xué)習(xí)領(lǐng)域是否需要所有的循環(huán)變換及其組合的時(shí)候，設(shè)想從減少循環(huán)變換的個(gè)數(shù)來(lái)減小解空間，以此來(lái)加速調(diào)度的過(guò)程；另外，MLIR的初衷也是為了降低代碼生成的復(fù)雜度。而文獻(xiàn)[34, 35]也試圖對(duì)特定情況或者將線性整數(shù)規(guī)劃問(wèn)題簡(jiǎn)化為線性規(guī)劃問(wèn)題來(lái)降低時(shí)間復(fù)雜度。不過(guò)，諸此種種都沒(méi)有從實(shí)質(zhì)上解決掉問(wèn)題的關(guān)鍵，因?yàn)閱?wèn)題的實(shí)質(zhì)仍然是NP級(jí)別的難題。想要從質(zhì)上改變這個(gè)現(xiàn)狀，可能還需要一段比較長(zhǎng)的時(shí)間，其它的計(jì)算機(jī)科學(xué)領(lǐng)域的方法比如constraint programming說(shuō)不定也能是一個(gè)解決的方法。

另外，Poly的靜態(tài)仿射約束對(duì)稀疏tensor等領(lǐng)域的擴(kuò)展也提出了挑戰(zhàn)。關(guān)于稀疏tensor的工作目前也有了一定的研究[36, 37]，但Poly無(wú)法直接應(yīng)用于含有非規(guī)則下標(biāo)的tensor的情況，怎么樣在這個(gè)領(lǐng)域?qū)oly進(jìn)行擴(kuò)展也可能是深度學(xué)習(xí)利用Poly優(yōu)化的另一個(gè)需要解決的問(wèn)題。因?yàn)镻oly在解決稀疏矩陣問(wèn)題的研究時(shí)，有了一定的進(jìn)展[38-41]，這說(shuō)明Poly的non-affine擴(kuò)展還是可行的，而深度學(xué)習(xí)框架的可定制性給這個(gè)問(wèn)題也創(chuàng)造了更多的機(jī)會(huì)。