AICompiler動態shape編譯框架
移動互聯網的興起，不僅產生了海量數據，也對人機交互有了新的定義。企業如何動態處理不同規格圖片數據，如何更靈活處理不同長度的對話語料等等，提升企業運營效率，爭取更多的商業機會和流量，成為眾多企業探索的熱門技術應用。
近期，阿里云機器學習PAI團隊全新上線一套Dynamic Shape Compiler框架，不僅作為AICompiler技術棧中原有的Static Shape Compiler框架的重要補充，更是增加了Compiler在企業級數據處理應用的無限可能，在提升數據處理效率的同時，大幅提升AI工程化效率。先來看看案例和效果數據。
性能結果。
TensorFlow語音識別業務示例
以某業務方的語音識別模型為例。過往為這個業務提供的優化主要基于Static Shape Compiler，shape變化范圍較大，只能采用脫機預編譯的方式來完成部署，部署過程較為繁瑣。
下圖展示了基于Dynamic Shape Compiler在不同batchsize下的實際性能結果，其中縱軸為latency的提升倍數。整體編譯次數從之前的幾千次降低到1次。從數字上來看，在只有一次編譯的較小編譯開銷下，性能十分接近Static Shape Compiler的性能優化結果。

TensorFlow廣告推薦業務示例
下面這個例子是某廣告推薦業務方的推理模型，在在線系統中，預估時的input shape變化非常頻繁，比如：用戶畫像卷標個數可能不同；用戶的歷史行為序列的長度會不同；召回廣告的集合大小會不同；廣告所屬類目的數量也會不同。這些變量會最終導致請求預估服務的input shape也是時刻變化的。
過往業務方的同學需要通過batching/手工干預圈圖等方式才能將編譯次數控制到可接受范圍內，造成每次模型迭代后部署過程較為繁瑣。從性能結果上看，對比基于XLA優化的性能，Dynamic Shape Compiler基本接近甚至超過Static Shape Compiler的性能優化結果。

TensorFlow語音合成業務示例
以某業務方的TTS模型為例，具體看一下實際業務中對優化工具對動態shape支持的需求情況。在這個業務模型里，用戶的輸入sequence length輸出sequence length都可能發生變化。此外，由于TTS類算法本身需要在推理過程中引入隨機數的特點，即使輸入同一條樣本，內部子圖的shape也會發生變化。基于static shape compiler來優化這個模型的話，輸入數據數量，以及累積編譯次數的變化曲線。在這個業務模型中每個shape的編譯開銷約為20s左右，可以看到在經過幾百輪迭代之后，編譯cache都還沒有收斂趨勢，根據理論shape變化范圍測算總編譯時間至少在10個小時以上，這類業務如果使用XLA等靜態shape編譯程序的話，無法透明的實現可商用的部署。AICompiler里面的dynamic shape compiler組件很好的解決了這一問題，在一次編譯的前提下幫助用戶獲得平均2X的性能收益，目前業界尚無其它AI編譯程序能夠實現類似的效果。

PyTorch公式識別業務示例
下圖是一個PyTorch業務的例子，對比的Baseline是基于libTorch執行導出后的TorchScipt腳本，可以看到AICompiler對PyTorch業務提供了同樣的編譯優化能力。

本文主要介紹這套動態shape編譯框架，對更多技術細節興趣的讀者可以參考DISC: A Dynamic Shape Compiler for Machine Learning Workloads.
從PAI團隊三年前啟動深度學習編譯程序方向的工作以來，「Dynamic Shape」問題一直是阻礙實際業務落地的嚴重問題之一。包括XLA在內的主流深度學習框架，都是基于Static Shape語義的編譯程序框架。即，just-in-time運行的編譯程序，在運行時捕捉待編譯子圖的實際輸入shape組合，并且為每一個輸入shape組合生成一份編譯結果。
Static Shape Compiler的優勢顯而易見，編譯期完全已知靜態shape信息的情況下，Compiler可以作出更好的優化決策并得到更好的CodeGen性能，同時也能夠得到更好的顯存/內存優化plan和調度執行plan；然而，Static Shape Compiler的缺點也十分明顯，具體包括：
? 編譯開銷的增加。對于訓練業務，編譯開銷導致訓練迭代速度不穩定，訓練初期顯著負優化，甚至整個訓練過程的時間開銷負優化；對于Inference業務，很多業務實際部署和迭代時不允許出現性能抖動，而脫機的預編譯預熱又會使得部署的過程變復雜。
? 內存顯存占用的增加。除編譯開銷的問題之外，當shape變化范圍特別大的時候，編譯緩存額外占用的內存顯存，經常導致實際部署環境下的內存/顯存OOM，直接阻礙業務的實際落地。
? 對于一部分業務場景，shape變化范圍可能非常大甚至是趨于無窮的，比較常見的包括廣告推薦類業務中常見的稀疏化模型，還有例如分布式訓練下的embedding切片等等。在這種情況下，編譯緩存永遠也無法收斂，用戶也就不可能通過compiler獲取到性能收益了。
? 上述問題在部分情況下，可以通過人工干預Compiler的圈圖過程來緩解，即，將shape變化劇烈的子圖排除在編譯范圍之外。然而，這種解決辦法對用戶非常不友好，大大降低了Compiler應用的通用性和透明性，這要求做部署和優化的同學同時對模型結構和compiler非常了解，且每一次模型結構迭代時，都需要花費額外的工作量來調整圈圖，獲得可以接受的性能效果。

關于這一問題，曾經出現過若干類解決方案，包括，對Compiler在圈圖過程中的自動化干預；在編譯期內部自動對變化維度做bucketing補齊并將子圖計算結果做自動的slicing。然而這些解決方案都存在各自的局限，例如前者只能適配于小部分子圖shape變化劇烈的情況，后者在很多模型上都無法得到自動slicing的完備數學推導。
為徹底解決這一問題，選擇基于MLIR（Multi Layer Intermediate Representation），結合團隊過往對AICompiler中積累的部分經驗，打造一套完備支持Dynamic Shape語義的AI編譯程序，希望能夠徹底解決深度學習編譯程序在這部分對靈活性要求較高的業務中無法落地應用的問題。
整體架構
Dynamic Shape Compiler的整體架構，及其在AICompiler中的上下文關系如下圖所示。

Compiler部分
MLIR Infrastruction
MLIR是由Google在2019年發起的項目，MLIR 的核心是一套靈活的多層IR基礎設施和編譯程序實用工具庫，深受 LLVM 的影響，并重用其許多優秀理念。
這里選擇基于MLIR的主要原因包括：
? 比較豐富的基礎設施支持，使得完成編譯程序的常規開發工作更為便捷，效率更好。TableGen，以及編寫常規pattern matching的graph optimization pass的簡化等。
? Open for Extension的模塊化設計架構，這里的核心是其Dialect抽象的設計。除Dialect的concept本身，在架構設計上，基于LLVM在傳統編譯期領域的成功經驗，MLIR團隊還是展現出了老練的架構設計能力，將整個MLIR架構的設計變得很具模塊化。
? MLIR的膠水能力，使得其可以比較靈活方便地與已經存在的優化手段進行集成，而非拒斥。
具體實現
MLIR框架的上述特性，使得可以比較方便的有選擇性的leverage部分小區已有組件，避免完全的重新造輪子，也一定程度上避免從頭徹底重構XLA代碼帶來的巨大工作量。
這里根據過往對AI編譯程序的理解，選擇了4層比較主要的中間層抽象，包括：
? DHLO Dialect：能夠完備表達動態shape語義的操作數層計算圖抽象，主要作用是能夠用有限數量的操作數類型，描述不同前端框架的大量操作數定義，且表達足夠靈活。
? DLHLO Dialect：引入Buffer語義的計算圖抽象，用于在編譯程序流程中進行內存/顯存的管理和優化。
? Loop Dialect：用于將操作數層的計算描述基于Loop等展開為指令集的計算描述，在這一層上完成了操作數fusion的CodeGen。
? GPU Dialect：為GPU編程模型中的kernel launching及各種底層原語提供中間層抽象。
下圖展示了基于MLIR的Loop Dialect等基礎設施，在CodeGen中實現最簡單的Input fusion的基本原理。對比XLA中只有高層的HLO和底層的llvm兩層中間表示，MLIR提供的Loop Dialect抽象可以直接在中間層完成fusion，很好的簡化了開發的復雜度。

這次不再贅述Compiler部分其它各個模塊的具體實現細節，移步MLIR小區中發起的相關細節討論：RFC，以及會議討論。
此處想著重介紹下對比于XLA，Dynamic Shape Compiler需要額外考慮的一些問題，包括：
? DHLO IR，在XLA的HLO IR基礎上，擴展了一套具有完備動態shape表達能力的IR。靜態場景下，HLO IR中的shape表達會被靜態化，所有的shape計算會被固化為編譯時常量保留在編譯結果中；而在動態shape場景下，IR本身需要有足夠的能力表達shape計算和動態shape信息的傳遞。
? Placer模塊，對于Dynamic Shape Compiler來說，計算可以分為shape計算和data計算兩類，對于GPU backend而言，通常shape計算的計算量較小，launch和拷貝開銷相比較大因此通常更適合在host側完成計算。實現了一個簡單的單卡分圖策略，對host側和device側計算執行不同的lowering pipeline。
? Buffer管理及Buffer優化模塊，有別于靜態Shape編譯期能夠比較容易通過liveness分析，實現Buffer的復用等優化，而在動態shape語境下，由于Buffer Size未知編譯期則不容易做到完全一致的優化。目前使用的是動態的Buffer申請和釋放，優化申請和釋放的時間點，同時后臺使用應用層包含Cache的Allocator，來達到性能和靈活性之間的平衡。后續可考慮在IR中充分表達Shape Constraint信息的情況下，來嘗試在編譯期做精細的Buffer復用優化。
此外，注意到在動態shape語境會為編譯期的性能performance帶來一些有趣的新挑戰：
? 部分優化決策后置到運行期，以Implicit Broadcast為例，目前主流的前端AI框架都支持implicit broadcast語義，而在動態shape語義下，編譯期無法充分知道LHS/RHS是否需要執行Broadcast操作。為保證完備性，如果所有情況下都穩妥的執行Broadcast計算的話，則會帶來比較嚴重的冗余計算/Fusion顆粒度等問題。其它與之類似問題還包括GPU Kernel的Launch Dimension選擇等，解決這一問題的做法是編譯期做多版本編譯，運行期根據實際shape來選擇最優實現，保證靈活性的同時，緩解靈活性帶來的性能損耗。
? Shape約束信息的使用，發現在Dynamic Shape Compiler中，即使Tensor的Shape信息未知，但Shape之間的約束信息，例如兩個Tensor之間的某兩個維度的size是否相等等信息，仍然會對編譯結果的性能產生比較重要的影響。主要原因包括：在圖層面，這些信息帶來了更大的圖優化空間，而在CodeGen層面，這些信息能夠更有效的指導低層Lowering做CSE等傳統編譯程序優化，減少冗余的計算指令數。
多前端框架支持
隨著近年來PyTorch用戶數量的持續增加，對PyTorch作業的性能優化需求也正在變得越來越重要。AICompiler框架在設計時，包含了擴展支持不同前端框架的考慮。
從IR lowering的角度，這里選擇相比于HLO更具泛化表達能力的DHLO Dialect作為不同前端框架的統一接入IR，而在PyTorch側選擇用戶部署時導出的TorchScript IR，通過實現一個輕量的Converter將TorchScript轉換為DHLO IR實現了對PyTorch Inference作業的覆蓋。MLIR相對完備的IR基礎設施也為Converter的實現提供了便利。
RAL (Runtime Abstraction Layer）
除編譯本身的問題之外，還面臨其它一些問題，例如如何將編譯的結果能夠配合TensorFlow/LibTorch等宿主在各自的運行環境上下文中執行起來，如何管理運行時IR層不易表達的狀態信息等等。希望為不同的運行時環境實現一套統一的Compiler架構，為此引入了運行時抽象層，即RAL層。RAL層主要負責解決如下問題：
Compile Once and Run Anywhere
RAL實現了多種運行環境的適配支持，用戶可以根據需要進行選擇。
? 全圖編譯，獨立運行。當整個計算圖都支持編譯時，RAL提供了一套簡易的runtime以及在此之上RAL Driver的實現，使得compiler編譯出來結果可以脫離框架直接運行，減少框架overhad，比如在支持某語音ASR模型（類transformer網絡）推理優化時，使用全圖編譯將框架開銷從TF的22ms減小到4ms；
? TF中子圖編譯運行。RAL目前實現了TF Driver，可以支持在訓練/推理場景中對圈出的子圖進行編譯執行；
? Pytorch中子圖編譯運行。RAL目前實現了Libtorch Driver，可以支持在推理場景中對圈出子圖進行編譯執行；
以上環境中在諸如資源（e.g. memory）管理，API語義等上存在差異，希望能夠引入一層抽象對compiler側屏蔽這些差異。RAL通過抽象出一套最小集合的API （RAL中稱為Driver），并清晰的定義出它們的語義，這樣compiler和runtime就可以在一定層度上隔離開來，簡化compiler的開發，同時通過提供這套API在不同環境下的實現，來達到在不同的環境中都能夠執行編譯出來的結果的目的。
Stateless編譯
dynamic shape compiler完成一個計算圖的編譯之后，編譯的結果可能被多次執行，而有些op的執行是帶狀態的：
? 在device（e.g. gpu）上執行時，對const op希望只在第一次執行時加載并常駐device，而不是每次都引入一次host-to-device的拷貝；
? 對于需要根據具體shape信息進行tuning的op （e.g. gemm/conv），tuning cache需要一個地方存儲；
RAL將資源初始化等帶狀態的部分抽取出來，封裝成context來管理生命周期。在代碼生成的過程中，通過簡單的注入context，將狀態量隱藏在context之后，使得compiler側看到的是一個純計算的過程。無狀態的設計一方面簡化了代碼生成的復雜度，另一方面也更容易支持多線程并發執行（比如推理）的場景，同時在錯誤處理，回滾方面也更加容易支持。
對用戶透明的編譯模式切換
對于Dynamic Shape Compiler在AICompiler中的定位是：與原Static Shape Compiler并列的一套框架，在允許適度犧牲性能的情況下，提供對于強Dynamic Shape類業務的通用透明支持。
然而從用戶的角度來說，通常并不容易判斷一個Workload的更適合Dynamic Shape Compiler還是Static Shape Compiler，為此結合接耦和全量打開[link]中的工作，設計了一套編譯模式自動切換的狀態機。其基本思路是，在任務初期先選擇較為安全的Dynamic Shape Compiler，結合后臺編譯讓用戶能夠在運行時盡早得到有性能提升的編譯執行，并在后續執行過程中結合資源的實際占用情況和實際運行時的shape變化范圍來有選擇性的切換到Static Shape Compiler的執行。
兩套compiler在運行時的切換關系如下圖所示：

阿里云機器學習PAI平臺面向企業客戶及開發者，提供輕量化、高性價比的云原生機器學習平臺，涵蓋交互式建模、拖拽式可視化建模、分布式訓練到模型在線部署的全流程覆蓋，百余種落地場景，全面提升機器學習工程效率。目前，PAI AICompiler已經集成在阿里云機器學習PAI的通用推理優化工具PAI-Blade敏捷版中，用戶可以可以參考開發文檔來快速體驗。

總結

以上是生活随笔為你收集整理的AICompiler动态shape编译框架的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。