全面支持 PyTorch 2.0：BladeDISC 5 月~11 月新功能發佈

2022 年 12 月 12 日
AI
深度學習

作者：BladeDISC研發團隊

BladeDISC 上一次更新主要發佈了 GPU AStitch 優化，方法來源於我們發表在 ASPLOS 2022上的論文AStitch。這一次，我們發佈了 0.3.0 版本。

本次更新中 BladeDISC 社區全面支持了 PyTorch 2.0 編譯，推進了和 Torch-MLIR 社區的合作；增加了 CPU 量化編譯和倚天新硬件支持；在編譯優化方面 BladeDISC 社區增加了一系列特性，包括改進了 GPU 訪存密集計算的性能，完成了 Shape Constaints IR 功能設計和支持。

本文描述 BladeDISC v0.3.0 版本對於 v0.2.0 的主要更新內容。

PyTorch 2.0 和動態編譯支持

BladeDISC 近半年來高度關注 PyTorch 2.0 相關特性的新動態，積极參与社區的協作，在此期間完成了 TorchBlade 編譯架構的調整，更好地 PyTorch 動態編譯和訓練支持。

TorchDynamo 優化

現在使用 PyTorch nightly 版本和 BladeDISC，只需要額外的兩行代碼改動即可完成 BladeDISC 的編譯加速:

import torch_blade # one more extra line
model = ...
compiled_model = torch.compile(model, backend='disc')

TorchBenchmark

BladeDISC 持續關注深度學習模型的通用優化，我們將 TorchBenchmark 作為優化的指南針評估和持續提升 BladeDISC 在不同模型上的魯棒性和優化效果。

TorchMLIR(MHLO) 和動態性貢獻

BladeDISC 一直是 MLIR/MHLO 動態性的主要貢獻和推動者。這個版本中我們與位元組 AML 團隊合作往 Torch-MLIR 貢獻了 Torch-to-MHLO 的模塊，特別是對動態性的支持，請參考 [RFC] Adding Torch to MHLO conversion #999，特別感謝位元組 AML 的同事的共同推動!

並且我們對 BladeDISC 中的 PyTorch 編譯流程做了架構調整，Torch-MLIR 成為了 BladeDISC 的基礎模塊。下圖中藍色線框表示了 Torch-to-MHLO 工作在 BladeDISC 中相對的位置。

我們號召社區對此模塊感興趣的朋友們一起合作推動 Torch-MLIR(MHLO) 模塊的迭代與演進，減少基礎工作的碎片化。

PyTorch 訓練支持

BladeDISC 正在逐步支持 PyTorch 模型的訓練優化，目前已經成功支持 BERT 等模型的編譯優化。BladeDISC 支持採用 PyTorch 2.0 的 TorchDynamo 和 Lazy Tensor Core 兩種方式開啟訓練優化(當前此模塊的接口沒有完全確定，鑒於目前 PyTorch 社區的活躍度，我們會更優先支持 TorchDynamo 的使用方式)。目前此模塊仍在持續更新和迭代，更多狀態請關注我們的開發動態。

EasyCV/NLP 推理加速編譯支持

BEVFormer: 是一種純視覺的自動駕駛感知算法，藉助 BladeDISC 優化，取得了 1.42 倍的端到端性能提升。
PAI-Diffusion Model: 在此版本中 BladeDISC 社區也對最近大熱的 AIGC Diffusion 相關模型提供了通用優化支持。BladeDISC 為 PAI-Diffusion 提供了接近 3 倍的端到端的性能提升。

更多信息請查看 PAI EasyCV/NLP 相關係列文章和開源地址:

EasyCV: //github.com/alibaba/EasyCV
EasyNLP: //github.com/alibaba/EasyNLP

BladeDISC 量化 (Experimental)

在這半年中我們完成了編譯+量化結合的一系列初步摸索，在包括X86、ARM多個不同硬件平台上完成了早期的方案及性能驗證，下表中展示了我們一些初步的成果。

模型	輸入	設備	優化前PyTorch/FP32	優化後Int8量化+編譯
bert-mini	8*64	g6r / Ampere Altra / 單核	135.9 ms	39.6 ms
bert-mini	8*64	g8m / YiTian /單核	127.8 ms	31.1 ms
bert-mini	8*64	hfg7 / Cooper Lake 8369 /單核	37.5 ms	21.5 ms

近期我們將支持更多硬件平台（例如CUDA），並提供如何量化PyTorch/TensorFlow模型的示例。此外，我們也將繼續優化提升量化模型的推理性能。

BladeDISC 編譯優化改進

新硬件支持：AArch64（倚天）

在這半年中我們進一步完善了對AArch64類硬件（尤其是倚天）的支持，進行了一系列針對性改進：

增加對BF16/int8 GEMM/Conv的支持，充分利用倚天硬件的能力；
對Arm Compute Library進行了一系列的定製和改進，解決其在dynamic shape及高並發場景下的各種可用性問題；
訪存密集型算子CodeGen質量改進，包括Stitch-CPU對reshape類算子的支持從而支持更大粒度的kernel，以及op duplication策略的引入進一步減少訪存量；

GPU上訪存密集計算codegen性能增強

針對GPU上訪存密集計算子圖代碼生成提供了一系列的深度優化，單個LayerNorm子圖在常規推理shape下可帶來最高2X的性能提升。主要的優化功能包括：

Fusion中若干獨立控制流的合併。比如一個stitch fusion中包含多個獨立的且shape相同的row reduce計算，則將這幾個獨立row reduce計算的控制流合併為同一個，一方面減少不必要的計算，另一方面增大ILP；
Row-reduce的schedule選擇邏輯優化。針對不同的shape，選擇更加合適的row reduce計算的的schedule；
優化element-wise fusion的向量化優化。通過指令交叉來支持element-wise fusion的數據讀和計算的向量化；
Loop相關優化。添加loop unroll及instruction interleave優化，增大ILP；添加循環不變量外移優化，減少不必要的計算；
消除kernel的無效argument ，減小kernel launch開銷。

以上功能可以通過設置變量DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL=true來打開。

Shape Constraint IR

在這半年中我們完成Shape Constraint IR的設計和開發，通過將shape constraint作為第一等公民引入到IR中，可以方便我們充分挖掘計算圖中蘊含的結構化約束，並以此來輔助完成一系列動態shape語意下的優化，進一步縮小與static shape compiler在優化空間上的差異。感興趣的讀者可以在這裡了解我們的設計文檔，也可以閱讀我們在知乎上分享的技術文章 (link1, link2)。

對二次定製開發支持的增強

我們基於MLIR社區PDL的工作重構了BladeDISC中接入一個custom library call流程，極大的簡化了相關的開發工作量。在新的方式下，用戶只需要提供一個PDL的pattern描述文件，以及一個符合BladeDISC runtime接口的kernel，便可以在不重新編譯BladeDISC的情況下，實現pattern的替換及對應kernel的調用，明顯改進了BladeDISC對二次定製開發的支持。我們在量化這個場景下檢驗了新的基礎設施的可用性及工程效率，感興趣的同學可以參考這裡和這裡的例子。後續我們還將藉助PDL和transform dialect進行進一步拓展，使得不僅僅是對custom kernel，也能對特定pattern的CodeGen策略進行定製。

Runtime Abstraction Layer 改進

大模型權重的支持：我們重構了常量（比如模型的權重參數）編譯結果的存儲格式，從基於protobuf改成了自定義的格式，從而去除了對const上限的約束，方便我們支持大模型的優化。

並發性能改進：針對高並發場景（比如同時服務數百路並發的推理請求）進行了一系列的優化，進一步縮小了kernel 調度以及共享資源鎖同步的開銷，在某語音識別200+並發的場景中進一步取得20%+的性能改進。

Ongoing Work

CUTLASS Gemm CodeGen

在GPU上接入了CUTLASS進行計算密集型算子的算子融合與代碼生成，自動化地將GEMM及後續的element-wise（如GELU等激活函數子圖、Tranpose算子等）進行計算融合與代碼生成。目前GEMM + GELU及GEMM + Transpose的通路已經走通，且在BERT模型上取得了加速效果，魯棒性正在提升中，可使用DISC_ENABLE_COMPUTE_INTENSIVE_FUSE=true設置來嘗試使用。

MLIR Transform Dialect Based CodeGen

我們目前正在基於MLIR社區的transform dialect進行計算密集型相關pattern在dynamic shape語意下的代碼生成，目前第一個目標是在AArch64平台上GEMM相關pattern達到與ACL相當的性能，以期通過白盒的方式徹底解決ACL在dynamic shape以及服務端多路請求並發的場景下的可用性問題，相關的工作的最新進展可以參見這裡。

全面支持 PyTorch 2.0：BladeDISC 5 月~11 月新功能發佈