在本系列分享中我們將介紹BladeDISC在動態shape語義下做效能最佳化的一些實踐和思考。本次分享的是我們最近開展的有關shape constraint IR的工作，鑑於篇幅較長，為了提升閱讀體驗，我們將分享拆分為兩個部分：

• Part I 中我們將介紹問題的背景，面臨的主要挑戰和以及我們做shape constraint IR的動機；
• Part II 中我們將介紹shape constraint IR的設計，實現以及一些初步的實驗結果；

本篇是關於Part I的介紹，Part II的介紹將後續發出。

背景和挑戰

主流的AI模型在部署時一般都具備不同程度的動態shape問題，比如輸入圖片尺寸，batch size 或者序列長度的變化等。與靜態shape語義下做最佳化相比，在動態shape語義下做最佳化由於缺少具體的shape資訊往往具有更大的難度，主要體現在以下幾個方面：

挑戰1：最佳化目標的轉變。在靜態shape下，我們的最佳化目標是希望在給定的shape下，儘可能逼近理論上限的速度，針對不同的shape可以使用不同的最佳化手段，而在動態shape語義下，我們最佳化目標是希望使用一套方法提升在整個shape分佈上的平均效能，更強調的是 最佳化方法的跨shape可遷移性和穩定性。因此很多在靜態shape下常用的最佳化，比如profling驅動的策略，將不再簡單可用。

挑戰2：更少的有效資訊。最佳化AI模型時很多常見的手段都把一些shape關係的斷言是否成立作為最佳化觸發的前置條件。比如在計算圖化簡時消除冗餘的broadcast op，需要依賴於判斷輸入和輸出是否具有相同的shape。在靜態shape語義下，判斷常量shape是否相等是顯然的，而動態shape語義下，判斷symbolic shape相等則困難的多，而一旦我們無法有效判斷這些前置的shape關係斷言是否成立，後續最佳化都無法進行，因而丟失很多最佳化機會，拉大與靜態shape情況下效能的差異。

挑戰3：更復雜的計算圖。在動態shape語義下，由於shape不再是編譯（或者最佳化）期間的常量，整個計算圖中混雜著計算shape的IR以及計算data的IR，使得整個計算圖的分析和最佳化都變得更復雜，同時也會引入更多shape相關計算的開銷。

下圖中展示了一個支援numpy語義(IB)的Add OP的例子以說明計算圖變複雜的具體過程。在IB語義下，Add OP在執行時會根據輸入shape之間的關係，隱式的插入broadcast運算，所以下圖左側中展示的三種可能輸入都是合法的。在靜態shape語義下我們很容易在編譯期就區分開實際是那種情況，故而在編譯時只需要對一種具體情況進行處理，而在動態shape語義下，由於編譯期間無法進行區分，我們需要確保編譯的結果在三種情況下都可以工作，因而會在計算圖中引入顯示的shape 計算的IR以及broadcast操作（如下圖右側所示）。在這個例子中上層框架中一個普通的Add OP在動態shape語義下，也會被展開成一個複雜的子圖。

也正因為上述的這些挑戰，目前成熟的最佳化工具（如TensorRT，XLA/TVM）對於動態shape支援都還比較有限。 BladeDISC是阿里雲端計算平臺PAI團隊自研的一款原生支援動態shape的AI編譯器。在過往一段時間我們開發BladeDISC最佳化動態shape模型的過程中，我們發現儘管不知道shape的具體的數值，但是透過充分發掘和利用Tensor的shape之間的結構化關係或者Tensor shape自身的分佈特點，可以有效的解決上述挑戰。 在這裡我們把Tensor的shape之間的結構化關係或者Tensor shape自身的分佈統稱為shape constraint。更進一步的，我們發現透過將shape constraint作為第一等公民引入到IR中，可以最大化的發揮shape constraint的效能。本文中我們將介紹BladeDISC中shape constraint IR的定義以及如何利用它來輔助完成一系列動態shape語意下的最佳化以解決上述挑戰，縮小與靜態shape之間的效能差異。

動機

為什麼選擇shape constraint？

在上一章節中我們分析了在動態shape語義下給做最佳化所帶來的一系列挑戰。我們發現使用shape constraint可以有效的解決這些最佳化上面臨的困難。以下我們將分別介紹使用shape constraint如何有效解決上述的三個挑戰。

應對挑戰1：跨shape可遷移性

在具體分析之前，我們先看shape constraint在本文中的定義。假設一個tensor的rank是N，則該tensor的shape可以記為 (d0, d1, ... dN-1)，其中 di表示的是該tensor在第 i個軸上所具有的大小，也記為第 i個軸的 dimension size。文字中討論的shape constraint可以分為以下兩類：

• shape結構化約束。該類約束描述的是dimension size之間的相關關係，比如：
• dimension size相等關係：某一個tensor的 di與另外一個tensor的 dj具有相同的大小，或者同一個tensor的 di與 dj具有相同的大小；
• tensor元素個數相等：即一個tensor和另外一個tensor具有相同數量的元素；
• dimension size乘積相等關係：比如 reshape（[a, b, c, d]） -> [a*b, c*d]
• shape分佈約束。該類約束描述的是某個或某幾個dimension size的（聯合）分佈，比如：
• di % 4 = 0， di != 0,  di * dj=10;
• likely values: 即描述某個dimension size更大機率可能取到的值；
• value range：即描述某個dimension size可能的取值區間；

由上述定義本身我們可以立刻得到一個結論： 由於無論是shape結構化約束還是分佈約束都不依賴於具體的shape值，因此基於shape constraint而構建的最佳化策略天然具備跨shape可遷移性。

應對挑戰2：shape關係斷言分析

很多重要的最佳化都將一些shape關係的斷言是否成立作為最佳化觸發的前置條件，比如：

• 計算圖化簡。比如上文提到的消除冗餘的broadcast節點的例子。
• 計算圖layout全域性最佳化。計算密集型運算元的效能和其輸入輸出的資料排布（layout）有很強的關係，一個更合適的layout往往具有更好的效能。而一般最優的layout是隨著shape而變化的，導致在動態shape語意下無法靜態確定最優的layout；這裡一種最佳化策略是：將shape相容的計算密集運算元分到一組，只需要在組間插入layout轉化的運算元，而組內由於可以確認使用同一種layout而不必再插入layout轉化的運算元，從而提升效能；這裡shape相容的斷言是最佳化觸發的前置條件。
• fusion決策。在做運算元融合時並不是越多越好，只有將shape相容的運算元進行融合才會取得比較好的效果，同樣裡shape相容的斷言的判定是做運算元融合的前置條件。

在動態shape語義下，由於不知道shape具體的值，symbolic shape關係的斷言的分析往往是困難的。symbolic shape關係的斷言可以看成是symbolic dimension size間的關係斷言的邏輯關係表示式，因而問題可以轉換成對於symbolic dimension size間的關係斷言的判定。而由前述shape constraint定義可知，symbolic dimension size間的關係斷言本身是shape constraint一個例項。 在這裡我們把判定symbolic dimension size間的關係斷言是否成立的這個問題轉換成該斷言是否是已知原子shape constraint的組合 。這裡“原子”的定義是不能夠透過其他shape constraint 的例項的組合得到 。舉個例子，假設我們需要判定 tensor A: tensor<axaxf32>是否比 tensor B: tensor<bxbxf32>具有更多的元素個數，這個問題經過轉換過之後可以變成dimension size關係 a > b是否成立。

在完成上述問題的轉換之後，目前剩下的未解決的問題是：如何獲得 已知結果的原子shape constraint。具體來說有以下幾種方式：

• 由使用者提供或在JIT編譯期間自動捕獲。比如使用者可以提供關於模型輸入shape range，JIT編譯期間可以將捕獲到的一組輸入shape當作likely value 注入編譯流程等。
• 由op定義所攜帶的shape consraint資訊。比如：
• elementwise op的輸入和輸出應該具有相同的大小；
• mhlo.dynamic_reshape op的輸入和輸出應該具有相同的元素個數；
• mhlo.concatenate op 的輸入和輸出在非拼接的軸上應該具有相同的大小；
• 分析shape計算的IR或者透過傳播已有的shape constraint來獲得新的資訊，比如：

在充分利用上述來源原子shape contraint的基礎上，我們可以大幅減少由於shape未知導致一些最佳化前置條件無法判斷，進而導致最佳化無法生效的問題。

應對挑戰3：shape計算開銷及更復雜的計算圖

在動態shape語義下，會引入更多的shape計算的開銷，整個計算圖會變得更復雜。

對於shape計算開銷，shape結構化約束我們可以抵消大量重複的symbolic計算，從而儘可能減小額外的開銷。

對於計算圖化簡而言，我們一方面可以透過利用shape結構化約束消除一部分的冗餘計算，比如前文中提到的由於IB問題引入的大量broadcast op可以在計算圖化簡中消除。剩下的無法利用shape結構化約束消除的broadcast可以進一步利用以下shape分佈約束進行最佳化：IB觸發（即需要插入隱式的broadcast）的機率遠小於IB不觸發的機率。透過生成帶IB和不帶IB兩個版本的（如下圖所示），讓common case變得更快，從而提升期望效能；

為什麼需要shape constraint IR？

shape constraint在動態shape語義下很有用，但是要用好它卻並不容易。由於在整個pass pipeline中都可能會使用到shape constraint資訊，因此我們需要一種方式將shape constraint資訊在不同的pass之間進行傳遞。BladeDISC早期時使用是一種隱式傳遞的策略，即每個pass在需要使用shape constraint資訊時會透過分析一遍IR來重建shape constraint資訊。不難看出在上述的方案中shape constraint本身並不是IR一部分，這種策略帶來的問題是：

• 透過分析IR的方式一般只能夠重建（部分）結構化約束資訊，大部分的分佈約束資訊無法透過直接分析data計算IR來重建，而分佈約束對於動態shape語義下的效能最佳化同樣很重要；
• 在整個pass pipeline中IR會經歷多次的lowering （如TF/Torch dialect -> mhlo -> lmhlo -> ...），在每次lowering的過程中都可能會丟失一部分shape constraint資訊，比如下圖中所展示的 tf.SplitOp的例子。在這個例子中上層的 tf.SplitOp會被lower成一系列的mhlo的 SliceOp。根據 tf.SplitOp的定義我們可以知道它的所有輸出（假設有 N個）應該具有相同的shape，且輸入在被拆分的軸上的dimension size可以被 N整除，這些資訊在我們lower到mhlo時如果只進行data computation的轉換，而不進行shape constraint資訊的轉換，將會被丟失，從而使得後續的pass無法再利用相應的資訊進行最佳化；

為了解決上述的問題，我們更進一步將shape constraint 作為第一等公民引入到IR中，使得可以對結構化約束資訊和分佈約束資訊進行統一的建模，同時也確保在整個pass pipeline中各個pass可以看到一致的shape constraint資訊，進而更好的支援在不同的階段完成各自適合的最佳化。

合作討論

以上是近期我們在shape constraint IR Part I的分享，Part II後續也會發出，敬請期待，更多相關資訊歡迎加入BladeDISC使用者群一起交流討論。

歡迎大家加入BladeDISC使用者群一起交流討論，與我們一起共建。釘釘群號：44534789