在 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 用户交流群：

参考文献

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", https://www.tensorflow.org/xla
   

3. "NVIDIA TensorRT", https://developer.nvidia.com/tensorrt
   

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

5. "阿里 BladeDISC 深度学习编译器正式开源", https://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 推理加速器, https://help.aliyun.com/document_detail/205128.html