動態尺寸模型優化實踐之Shape Constraint IR Part II

語言: CN / TW / HK

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

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

本篇是關於Part II的介紹,Part I的介紹請參考這裡

 

設計和實現

shape constraint IR的設計

使用IR來建模shape constraint並不是一個很容易的事情。我們需要設計一種方案既方便我們做shape constraint分析,同時又不會使得IR的後續變換變的很複雜。經過多次迭代之後,我們選擇了type-based方案來構建shape constraint IR,基本思路如下段偽IR所示。

// original data-computation IR
func @main() {
  ...
  %0 = any_dialect.any_operation(...) : tensor<?x?xf32, [@S0, @S1]>
  ...
}

disc_shape.SymbolicDim @S0 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

disc_shape.SymbolicDim @S1 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

// A separated function to store shape constraint predicates between different symbolic dimensions.
// Each symbolic dim is either bound to a `disc_shape.dim` op or `disc_shape.bind_dim`
func @shape_constraint_graph(...) {
  %0 = disc_shape.dim() {ref: @S0} : index
  %1 = disc_shape.dim() {ref: @S1} : index
  disc_shape.tie_predicate_divisible(d0, d1) // d0 % d1 == 0

  // other tie_* ops
  //   disc_shape.tie_predicate_eq(d0, d1)  // d0 == d1
  //   disc_shape.tie_predicate_lt(d0, d1)  // dim less than
  //   disc_shape.tie_predicate_mul_eq(d0, d1, d2, ...) // d0 = d1 * d2 * ...
  //   // d0 * d1 = s0 * s1
  //   disc_shape.tie_predicate_product_eq([d0, d1, ..], [s0, s1, ...])
  //   // d0 = affine.apply(d1, d2, ...) {affine_attr = ...}
  //   disc_shape.tie_predicate_affine_eq(d0, d1, d2, ...) {affine_attr = ...}
}

在這個方案中,每一個symbolic dimension size(也即在編譯期間無法確定具體大小的dimension size)對應一個全域性的disc_shape.SymbolicDimIR物件。該IR物件中儲存了關於這個symbolic dimension size的分佈相關的約束,同時也儲存了對一個shape constraint function的引用。在上圖中最下面的部分是一個shape constraint function的例子。在這個function中儲存的是symbolic dimension dim之間的相關關係(結構化約束),每一種相關關係用一個op來抽象,比如說這個例子中說展示的整除等價關係便是由tie_predicate_divisibleop來描述。選擇type-based方案主data計算圖中並不會直接儲存shape constraint資訊。在上圖中的最上面是一個主data計算圖的例子。主data計算圖中,每一個tensor對應的type中都包含一個attribute,這個attribuet中儲存了這個tensor所對應的symblic dimension size的引用。

通過將描述shape constraint的IR和主data計算IR解耦開,一方面可以儘可能減少對已有pattern匹配的邏輯的干擾 (matmul+BiadAdd -> FusedMatmulBiasAdd這個pattern替換並不需要感知到shape constraint IR的存在),另外一方面,不同層級的data計算的IR,比如tensor level IR和buffer level IR,可以用同一套shape constraint的描述。從而可以緩解IR lowering過程中shape constraint資訊的丟失問題。

基於shape constraint IR的優化pipeline

將shape constraint IR作為第一等公民引入IR中之後,我們進一步構建了以shape constraint 為中心的優化 pipeline(如下圖所示)。通過對shape constraint的充分的挖掘,而非依賴於具體的shape的值來輔助完成各種優化,從而實現在動態shape語義下儘可能接近靜態shape優化工具的效能。

下圖中展示了目前BladeDISC中主要的幾個優化的階段。從最左邊開始看起。第一步是將前端AI框架的計算圖lower到MHLO的計算描述。這裡值得注意的是除了普通的data計算的lowering,還包含shape constraint的lowering,從而避免在動態shape語義資訊的丟失。到MHLO之後,我們首先會完成shape constraint的分析以及分析結果的IR化表示。分析得到的結果將可以指導我們完成計算圖上的一些基本化簡,比如冗餘broadcast op的消除,layout調整等。優化完之後的計算圖,我們會進一步對其中的訪存密集型運算元做融合優化,shape constraint將是決定那些運算元可以融合的很重要的判斷依據,通過更充分的挖掘shape constraint,我們可以找到更多融合的機會。最後在做程式碼生成的時候,我們發現在動態shape語義下index計算的開銷更加容易成為瓶頸,尤其是當運算元融合的數目比較多的時候,利用shape constraint我們可以大幅消除冗餘的index計算。以上限制於篇幅並未一一展開進行介紹,感興趣的同學可以通過這裡瞭解更多的細節。

image.png

初步測試

我們目前已經完成了shape constraint IR的第一階段開發,也即shape constraint IR的引入以及pass pipeline的適配性改造。第一階段的主要目標是搭建好整體的架子,還並未包含所有設計中優化的實現 (比如likely value的應用),我們將會在後續持續迭代完善shape constraint IR。在以上前提下,我們在一些比較典型的模型上完成初步的評測,下圖展示的是部分的評測結果。由於目前shape constraint IR還未應用到計算密集型運算元(GEMM/CONV等)的優化 ,故以下測評針對的是模型中訪存密集型部分,主要從兩個維度來衡量:

  • 訪存密集型部分kernel launch的次數,即衡量fusion的粒度,同等情況下次數越少,fusion的粒度越大;
  • 訪存密集型部分總消耗時間,即衡量生成的kernel的質量,同等情況下總時間越短,質量越高;

如下圖中所示,在CPU和GPU上我們都觀測到fusion粒度的明顯改善以及訪存密集型部分總消耗時間的減少。

合作討論

以上是近期我們在shape constraint IR Part II的分享,更多相關資訊歡迎加入BladeDISC使用者群交流討論。

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