全面支持 PyTorch 2.0：BladeDISC 5 月~11 月新功能发布

作者：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 linemodel = ...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: https://github.com/alibaba/EasyCV
  

• EasyNLP: https://github.com/alibaba/EasyNLP
  

BladeDISC 量化 (Experimental)

在这半年中我们完成了编译+量化结合的一系列初步摸索，在包括 X86、ARM 多个不同硬件平台上完成了早期的方案及性能验证，下表中展示了我们一些初步的成果。

近期我们将支持更多硬件平台（例如 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 以及服务端多路请求并发的场景下的可用性问题，相关的工作的最新进展可以参见这里。

推荐类稀疏模型

我们针对在业界中广泛应用的 Tensorflow 推荐类模型的性能热点部分--FeatureColumn中的稀疏算子进行了初步的优化支持。目前已经完成推理场景中常见的稀疏算子的 cpu codegen 支持，以及初步的算子融合支持，基于此目前已经可以在部分模型上获得了一定的收益。后续我们将支持更多种类，更大粒度的算子融合以及使用 CPU AVX 等指令集优化稀疏部分算子的计算性能，相关进展参见这里。

以上为本次 release 的主要内容。欢迎加入 BladeDISC 用户交流群。

项目开源地址：https://github.com/alibaba/BladeDISC

欢迎加入钉钉交流群，群号：44534789