阿里 BladeDISC 深度学习编译器正式开源

一 导读

随着深度学习的不断发展，AI模型结构在快速演化，底层计算硬件技术更是层出不穷，对于广大开发者来说不仅要考虑如何在复杂多变的场景下有效的将算力发挥出来，还要应对计算框架的持续迭代。深度编译器就成了应对以上问题广受关注的技术方向，让用户仅需专注于上层模型开发，降低手工优化性能的人力开发成本，进一步压榨硬件性能空间。阿里云机器学习PAI开源了业内较早投入实际业务应用的动态shape深度学习编译器 BladeDISC，本文将详解 BladeDISC的设计原理和应用。

二 BladeDISC是什么

BladeDISC是阿里最新开源的基于MLIR的动态shape深度学习编译器。

1 主要特性

• 多款前端框架支持：TensorFlow，PyTorch
• 多后端硬件支持：CUDA，ROCM，x86
• 完备支持动态shape语义编译
• 支持推理及训练
• 轻量化API，对用户通用透明
• 支持插件模式嵌入宿主框架运行，以及独立部署模式

三 深度学习编译器的背景

近几年来，深度学习编译器作为一个较新的技术方向异常地活跃，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，到后来呼声很高的MLIR以及其不同领域的延伸项目IREE、mlir-hlo等等。能够看到不同的公司、社区在这个领域进行着大量的探索和推进。

1 AI浪潮与芯片浪潮共同催生——从萌芽之初到蓬勃发展

深度学习编译器近年来之所以能够受到持续的关注，主要来自于几个方面的原因：

框架性能优化在模型泛化性方面的需求

深度学习还在日新月异的发展之中，创新的应用领域不断涌现，复杂多变的场景下如何有效的将硬件的算力充分发挥出来成为了整个AI应用链路中非常重要的一环。在早期，神经网络部署的侧重点在于框架和算子库，这部分职责很大程度上由深度学习框架、硬件厂商提供的算子库、以及业务团队的手工优化工作来承担。

上图将近年的深度学习框架粗略分为三代。一个趋势是在上层的用户API层面，这些框架在变得越来越灵活，灵活性变强的同时也为底层性能问题提出了更大的挑战。初代深度学习框架类似 Caffe 用 sequence of layer 的方式描述神经网络结构，第二代类似 TensorFlow 用更细粒度的 graph of operators 描述计算图，到第三代类似 PyTorch，TensorFlow Eager Mode 的动态图。我们可以看到框架越来越灵活，描述能力越来越强，带来的问题是优化底层性能变得越来越困难。业务团队也经常需要补充完成所需要的手工优化，这些工作依赖于特定业务和对底层硬件的理解，耗费人力且难以泛化。而深度学习编译器则是结合编译时图层的优化以及自动或者半自动的代码生成，将手工优化的原理做泛化性的沉淀，以替代纯手工优化带来的各种问题，去解决深度学习框架的灵活性和性能之间的矛盾。

AI框架在硬件泛化性方面的需求

表面上看近些年AI发展的有目共睹、方兴未艾，而后台的硬件算力数十年的发展才是催化AI繁荣的核心动力。随着晶体管微缩面临的各种物理挑战越来越大，芯片算力的增加越来越难，摩尔定律面临失效，创新体系结构的各种DSA芯片迎来了一波热潮，传统的x86、ARM等都在不同的领域强化着自己的竞争力。硬件的百花齐放也为AI框架的发展带来了新的挑战。

而硬件的创新是一个问题，如何能把硬件的算力在真实的业务场景中发挥出来则是另外一个问题。新的AI硬件厂商不得不面临除了要在硬件上创新，还要在软件栈上做重度人力投入的问题。向下如何兼容硬件，成为了如今深度学习框架的核心难点之一，而兼容硬件这件事，则需要由编译器来解决。

AI系统平台对前端AI框架泛化性方面的需求

当今主流的深度学习框架包括Tensoflow、Pytorch、Keras、JAX等等，几个框架都有各自的优缺点，在上层对用户的接口方面风格各异，却同样面临硬件适配以及充分发挥硬件算力的问题。不同的团队常根据自己的建模场景和使用习惯来选择不同的框架，而云厂商或者平台方的性能优化工具和硬件适配方案则需要同时考虑不同的前端框架，甚至未来框架演进的需求。Google利用XLA同时支持TensorFlow和JAX，同时其它开源社区还演进出了Torch_XLA，Torch-MLIR这样的接入方案，这些接入方案目前虽然在易用性和成熟程度等方面还存在一些问题，却反应出AI系统层的工作对前端AI框架泛化性方面的需求和技术趋势。

2 什么是深度学习编译器

传统编译器是以高层语言作为输入，避免用户直接去写机器码，而用相对灵活高效的语言来工作，并在编译的过程中引入优化来解决高层语言引入的性能问题，平衡开发效率与性能之间的矛盾。而深度学习编译器的作用相仿