動態尺寸模型優化實踐之Shape Constraint IR Part I
在本系列分享中我們將介紹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語義implicit broadcast (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 siz。文本中討論的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]
- dimension size相等關係：某一個tensor的
- 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_reshapop的輸入和輸出應該具有相同的元素個數；mhlo.concatenatop 的輸入和輸出在非拼接的軸上應該具有相同的大小；
- 分析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
