在BladeDISC正式開源三個月後，我們釋出了0.2.0版本，該更新包含了大量的效能最佳化與功能增強。

是目前業界領先的支援動態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上的論文 [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功能。我們提供了 以展示如何開啟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提供的 及 。

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的部分內容，更多關於本次版本更新的內容，請檢視完整的 。 歡迎加入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", 
3. "NVIDIA TensorRT",  https://developer.nvidia.com/tensorrt
4. "Putting the VM in TVM: The Relay Virtual Machine", 
5. "阿里 BladeDISC 深度學習編譯器正式開源", 
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推理加速器,  https://help.aliyun.com/document_detail/205128.html