移動網際網路的興起，不僅産生了海量資料，也對人機互動有了新的定義。企業如何動态處理不同規格圖檔資料，如何更靈活處理不同長度的對話語料等等，提升企業營運效率，争取更多的商業機會和流量，成為衆多企業探索的熱門技術應用。
近期，阿裡雲機器學習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 AICompiler已經內建在阿裡雲機器學習PAI的通用推理優化工具PAI-Blade靈活版中，使用者可以可以參考
開發文檔來快速體驗。
作者：朱凱，趙文益，楊軍