BladeDISC 0.2.0更新發布

• 2022 年 5 月 23 日
• AI
• 深度學習

在BladeDISC正式開源三個月後，我們發布了0.2.0版本，該更新包含了大量的性能優化與功能增強。

BladeDISC是目前業界領先的支援動態shape的深度學習優化編譯器。深度學習優化編譯器負責將上層的神經網路計算圖轉換為底層硬體可執行的程式，當前流行的深度學習優化編譯器（TVM[1]、XLA[2]、TensorRT[3]等）對靜態shape的支援力度較大，對動態shape的支援則有所欠缺。其中，XLA目前只支援靜態shape，TensorRT可以支援ranged shape（即指定尺寸範圍內的動態shape）。BladeDISC編譯器提供了對動態shape的完整支援，可以將包含動態shape語義的上層模型描述轉換為高效的底層可執行程式。更多關於BladeDISC的介紹內容，請查看此處[5]。

本文描述BladeDISC v0.2.0版本相對於最初開源版本（即v0.1.0版本）的主要更新內容。

性能優化

當前的神經網路計算圖主要由訪存密集型運算元（主要包括element-wise/elemental運算元及reduce運算元）和計算密集型運算元（主要包括GEMM和Convolution運算元）組成。BladeDISC v0.2.0版本相對v0.1.0版本在這兩方面都做了大量優化。

訪存密集運算元的優化

本章重點描述v0.2.0新增的GPU stitch優化（以及附帶的shape constraint功能增強）。

GPU上的stitch優化方法來源於我們發表在ASPLOS 2022上的論文AStitch[6]，我們正在逐步地將AStitch中的技術遷移到BladeDISC中，目前已完成部分遷移工作。Stitch的方法旨在解決當前深度學習模型在GPU上執行時遇到的大量的kernel調度（scheduling）和發射（launch）開銷以及片外訪存開銷，以及不同shape之下memory-intensive運算元的高效schedule問題，其基本方法是通過GPU上多層次的中間結果訪存管理（以暫存器，shared memory及global memory作為中間數據的buffer）、層次化的數據同步控制以及自動適配的parallelism-aware的codegen schedule生成，將原先多個不同的fusion kernel給stitch在一起，並生成高效的計算程式碼。

本次BladeDISC release實現了通過shared memory將原先被reduce運算元分割開的多個fusion進行stitch融合的功能，達到計算融合加速效果（跨CUDA thread block的global memory的stitch以及自動適配的parallelism-aware的codegen在本次release中暫未實現）。我們提供了一組示例來展示如何開啟stitch優化。

動態shape給stitch優化帶來了新的挑戰，其最重要的挑戰在於，在不知道確切shape的情況下，正確地判斷被stitch在一起的op之間的數據局部性，從而正確得為其中的數據傳輸分配片上存儲（即GPU shared memory）。本次更新包含了一系列的shape constraint功能增強，作為GPU stitch功能的基礎支撐。其基本思想是挖掘與傳播producer op和consumer op之間以及sibling op之間的shape constraint，通過symbolic的方式構建全局的shape等價關係；在此基礎上，進一步構建shape dimension之間的乘法關係，用以解析reshape等運算元的shape等價性關係。

值得說明的是，在前一個版本中（即BladeDISC開源版本v0.1.0），我們針對CPU上的模型，對部分訪存密集型計運算元圖實現了CPU端的stitch優化，從而更好地增強數據局部性，提升訪存效率。

計算密集運算元的優化

本章主要介紹本次更新中包含的GEMM合併優化及CPU上計算密集型運算元的pre-packing和layout優化。

GEMM 合併優化

本次release實現了以下兩種GEMM合併優化：

1.兩個GEMM運算元有公共的操作數，將其合併為一個運算元，比如 A x B 和 A x C 可以合併為 A x concat(B, C)，真實模型中的典型場景是Attention中的QKV合併；

2.兩個GEMM有相同的計算形狀，將其合併為一個batched GEMM，比如對於 A x B 和 C x D，如果A和C以及B和D的形狀相同，那麼就可以合併為一個batched GEMM。

GEMM合併帶來了兩個好處。其一是，GEMM合併可以增加計算尺寸，從而更好地打滿硬體的計算峰值，提升計算效率；其二是，GEMM合併可以減少GPU kernel數量，從而降低kernel調度和發射的開銷。

CPU上計算密集運算元的pre-packing和layout優化

在CPU上，對於GEMM及Convolution計算，BladeDISC支援對GEMM的操作數進行pre-packing優化，通過packing的數據layout轉換，使得矩陣乘操作對操作數的訪問能夠更好地利用數據局部性（比如使得數據訪問更好地適配cache line）。具體來說，BladeDISC封裝了CPU上的計算庫，通過對其提供的packing函數的封裝與調用來實現pre-packing功能。我們提供了基於Albert的示例以展示如何開啟pre-packing優化。

對於Convolution函數，不同硬體vendor的不同計算庫可能需要不同的數據layout以得到最佳性能，不同的數據類型在不同layout下的性能表現也會不同（比如，NVIDIA GPU上的FP16在TensorCore上和FP32在SIMT core上對layout的要求有所不同）。本次release針對CPU和GPU上的計算庫，結合數據類型的考慮，為Convolution計算自動適配最佳的數據layout。數據layout的轉換通過transpose操作來進行，我們實現了前後的transpose的抵消，最大限度減少額外帶來的transpose的影響。

上圖展示了BladeDISC在四個當下流行的模型上的性能效果（在T4 GPU上進行驗證，更多模型還在驗證中）。圖中的Framework表示原始的深度學習框架（FastSpeech2使用了TensorFlow 2.4框架，其他模型使用了PyTorch 1.7.1框架），Static Compiler表示該框架下接入靜態優化編譯器後的性能（TensorFlow使用XLA，PyTorch通過轉onnx來利用TensorRT 8.2進行優化，TensorRT 8.2在優化T5和S2T的過程中失敗，因此沒有性能數據）。可以看到，BladeDISC相對於基本的深度學習框架可以取得最高達到8倍的性能加速效果，在BERT和FastSpeech2上，BladeDISC取得了與業界先進的靜態優化編譯器相近的優化效果。

重要功能支援

本次release也包括一系列的重要功能更新，包括：

X86和AArch64 CPU硬體的支援

本次release在X86和AArch64架構的CPU平台上都做了大量更新和支援。
X86平台方面，在已有的訪存密集運算元codegen支援的基礎上（v0.1.0包含的功能），本次release增加了計算密集運算元的支援。具體來說，BladeDISC同時支援了MKL和oneDNN兩種不同的計算庫後端，並支援運行時按需選擇。在功能支援之外，本次releaes也包括對計算密集運算元的性能優化（如前面章節提到的layout優化和weight pre-packing優化）。

AArch64平台方面，本次release完成了訪存密集運算元codegen對於AArch64平台的適配，以及計算密集運算元庫ACL的支援（通過oneDNN的形式）。

在上述功能的支援下，BladeDISC在X86平台以及AArch64平台上都已經端到端可用。具體使用方式及性能效果參見BaldeDISC提供的TF示例及PyTorch示例。

Blade推理加速器TensorRT的圈圖支援

本次release開源了Blade推理加速器[7]兩個重要的功能：TorchBlade和TensorFlowBlade。這兩部分是Blade推理加速器面向兩個最為廣泛使用的深度學習框架所做的接入層，旨在提升模型優化的體驗和完整度。Blade推理加速器在接入BladeDISC之外，也接入了TensorRT。具體來說，對於PyTorch和TensorFlow的模型，Blade推理加速器會自動識別出可以被TensorRT優化的計運算元圖，並送給TensorRT優化引擎進行優化。一定程度上提升了使用TensorRT的轉換成功率，並且提供了與BladeDISC發揮聯合優化作用的可能性。

PyTorch Training的Proof-of-Concept跑通

BladeDISC正在逐步支援PyTorch模型的訓練優化，目前已經成功跑通mnist的簡單模型。在實現層面，BladeDISC利用PyTorch的Lazy Tensor Core機制，將TorchScript子圖優化為高效的可執行程式。詳細的設計文檔見此處。

以上為本次release的部分內容，更多關於本次版本更新的內容，請查看完整的release note。歡迎加入BladeDISC用戶交流群：

項目開源地址：//github.com/alibaba/BladeDISC

歡迎加入釘釘交流群，群號：44534789

參考文獻

1. “TVM: An Automated End-to-End Optimizing Compiler for Deep Learning”, Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy. OSDI 2018

2. “XLA: Optimizing Compiler for Machine Learning”, //www.tensorflow.org/xla

3. “NVIDIA TensorRT”, //developer.nvidia.com/tensorrt

4. “Putting the VM in TVM: The Relay Virtual Machine”, //tvm.apache.org/docs//arch/virtual_machine.html

5. “阿里 BladeDISC 深度學習編譯器正式開源”, //zhuanlan.zhihu.com/p/462641670

6. “AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures”, Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Jidong Zhai, Shuaiwen Leon Song, and Wei Lin. ASPLOS 2022

7. Blade推理加速器, //help.aliyun.com/document_detail/205128.html