简介：随着深度学习的一直倒退，AI模型构造在疾速演变，底层计算硬件技术更是层出不穷，对于宽广开发者来说不仅要思考如何在复杂多变的场景下无效的将算力施展进去，还要应答计算框架的继续迭代。深度编译器就成了应答以上问题广受关注的技术方向，让用户仅需专一于下层模型开发，升高手工优化性能的人力开发成本，进一步压迫硬件性能空间。阿里云机器学习PAI开源了业内较早投入理论业务利用的动静shape深度学习编译器 BladeDISC，本文将详解 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 什么是深度学习编译器

传统编译器是以高层语言作为输出，防止用户间接去写机器码，而用绝对灵便高效的语言来工作，并在编译的过程中引入优化来解决高层语言引入的性能问题，均衡开发效率与性能之间的矛盾。而深度学习编译器的作用相仿，其输出是比拟灵便的，具备较高形象度的计算图形容，输入包含CPU、GPU及其他异构硬件平台上的底层机器码及执行引擎。

传统编译器的使命之一是加重编程者的压力。作为编译器的输出的高级语言往往更多地是形容一个逻辑，为了便当编程者，高级语言的形容会更加形象和灵便，至于这个逻辑在机器上是否可能高效的执行，往往是考验编译器的一个重要指标。深度学习作为一个近几年倒退异样疾速的应用领域，它的性能优化至关重要，并且同样存在高层形容的灵活性和抽象性与底层计算性能之间的矛盾，因而专门针对深度学习的编译器呈现了。而传统编译器的另外一个重要使命是，须要保障编程者输出的高层语言可能执行在不同体系结构和指令集的硬件计算单元上，这一点也同样反馈在深度学习编译器上。面对一个新的硬件设施，人工的话不太可能有精力对那么多指标硬件从新手写一个框架所须要的全副算子实现，深度学习编译器提供中间层的IR，将顶层框架的模型流图转化成中间层示意IR，在中间层IR上进行通用的图层优化，而在后端将优化后的IR通用性的生成各个指标平台的机器码。

深度学习编译器的指标是针对AI计算工作，以通用编译器的形式实现性能优化和硬件适配。让用户能够专一于下层模型开发，升高用户手工优化性能的人力开发成本，进一步压迫硬件性能空间。

3 间隔大规模利用面临的瓶颈问题

深度学习编译器倒退到明天，尽管在指标和技术架构方面与传统编译器有颇多相似之处，且在技术方向上体现出了良好的后劲，然而目前的理论利用范畴却依然间隔传统编译器有肯定的差距，次要难点包含：

易用性

深度学习编译器的初衷是心愿简化手工优化性能和适配硬件的人力老本。然而在现阶段，大规模部署利用深度学习编译器的挑战依然较大，可能将编译器用好的门槛较高，造成这个景象的次要起因包含：

  • 与前端框架对接的问题。不同框架对深度学习工作的形象形容和API接口各有不同，语义和机制上有各自的特点，且作为编译器输出的前端框架的算子类型数量呈开放性状态。如何在不保障所有算子被残缺反对的状况下透明化的反对用户的计算图形容，是深度学习编译器可能易于为用户所宽泛应用的重要因素之一。
  • 动静shape问题和动静计算图问题。现阶段支流的深度学习编译器次要针对特定的动态shape输出实现编译，此外对蕴含control flow语义的动静计算图只能提供无限的反对或者齐全不可能反对。而AI的利用场景却恰好存在大量这一类的工作需要。这时只能人工将计算图改写为动态或者半动态的计算图，或者想方法将适宜编译器的局部子图提取进去交给编译器。这无疑减轻了利用深度学习编译器时的工程累赘。更重大的问题是，很多工作类型并不能通过人工的改写来动态化，这导致这些状况下编译器齐全无奈理论利用。
  • 编译开销问题。作为性能优化工具的深度学习编译器只有在其编译开销比照带来的性能收益有足够劣势的状况下才真正具备实用价值。局部利用场景下对于编译开销的要求较高，例如一般规模的须要几天工夫实现训练任务有可能无奈承受几个小时的编译开销。对于利用工程师而言，应用编译器的状况下不能疾速的实现模型的调试，也减少了开发和部署的难度和累赘。
  • 对用户透明性问题。局部AI编译器并非齐全主动的编译工具，其性能体现比拟依赖于用户提供的高层形象的实现模版。次要是为算子开发工程师提供效率工具，升高用户人工调优各种算子实现的人力老本。但这也对使用者的算子开发教训和对硬件体系结构的相熟水平提出了比拟高的要求。此外，对于新硬件的软件开发者来说，现有的形象却又经常无奈足够形容翻新的硬件体系结构上所须要的算子实现。须要对编译器架构足够相熟的状况下对其进行二次开发甚至架构上的重构，门槛及开发累赘依然很高。

鲁棒性

目前支流的几个AI编译器我的项目大多数都还处于偏试验性质的产品，但产品的成熟度间隔工业级利用有较大的差距。这里的鲁棒性蕴含是否可能顺利完成输出计算图的编译，计算结果的正确性，以及性能上防止coner case下的极其badcase等。

性能问题

编译器的优化实质上是将人工的优化办法，或者人力不易探索到的优化办法通过泛化性的积淀和形象，以无限的编译开销来代替手工优化的人力老本。然而如何积淀和形象的方法学是整个链路内最为实质也是最难的问题。深度学习编译器只有在性能上真正可能代替或者超过人工优化，或者真正可能起到大幅度降低人力老本作用的状况下能力真正施展它的价值。

然而达到这一指标却并非易事，深度学习工作大都是tensor级别的计算，对于并行任务的拆分形式有很高的要求，但如何将手工的优化泛化性的积淀在编译器技术内，防止编译开销爆炸以及分层之后不同层级之间优化的联动，依然有更多的未知须要去摸索和开掘。这也成为以MLIR框架为代表的下一代深度学习编译器须要着重思考和解决的问题。

四 BladeDISC的次要技术特点

我的项目最早启动的初衷是为了解决XLA和TVM以后版本的动态shape限度，外部命名为 DISC (DynamIc Shape Compiler)，心愿打造一款可能在理论业务中应用的齐备反对动静shape语义的深度学习编译器。

从团队四年前启动深度学习编译器方向的工作以来，动静shape问题始终是妨碍理论业务落地的重大问题之一。彼时，包含XLA在内的支流深度学习框架，都是基于动态shape语义的编译器框架。典型的计划是须要用户指定输出的shape，或是由编译器在运行时捕获待编译子图的理论输出shape组合，并且为每一个输出shape组合生成一份编译后果。

动态shape编译器的劣势不言而喻，编译期齐全已知动态shape信息的状况下，Compiler能够作出更好的优化决策并失去更好的CodeGen性能，同时也可能失去更好的显存/内存优化plan和调度执行plan。然而，其毛病也非常显著，具体包含：

  • 大幅减少编译开销。引入离线编译预热过程，大幅减少推理工作部署过程复杂性；训练迭代速度不稳固甚至整体训练工夫负优化。
  • 局部业务场景shape变动范畴趋于无穷的，导致编译缓存永远无奈收敛，计划不可用。
  • 内存显存占用的减少。编译缓存额定占用的内存显存，常常导致理论部署环境下的内存/显存OOM，间接妨碍业务的理论落地。
  • 人工padding为动态shape等缓解性计划对用户不敌对，大幅升高利用的通用性和透明性，影响迭代效率。

在2020年夏天，DISC实现了仅反对TensorFlow前端以及Nvidia GPU后端的初版，并且正式在阿里外部上线投入理论利用。最早在几个受困于动静shape问题已久的业务场景上投入使用，并且失去了预期中的成果。即在一次编译且不须要用户对计算图做非凡解决的状况下，齐备反对动静shape语义，且性能简直与动态shape编译器持平。比照TensorRT等基于手工算子库为主的优化框架，DISC基于编译器主动codegen的技术架构在常常为非标准开源模型的理论业务上取得了显著的性能和易用性劣势。

从2020年第二季度开始至今，DISC继续投入研发力量，针对前文提到的从云端平台方视角看到的深度学习编译器间隔大规模部署和利用的几个瓶颈问题，在性能、算子覆盖率和鲁棒性、CPU及新硬件反对、前端框架反对等方面逐步欠缺。目前在场景笼罩能力和性能等方面，曾经逐步替换掉团队过往基于XLA和TVM等动态shape框架上的工作，成为PAI-Blade反对阿里外部及阿里云内部业务的次要优化伎俩。2021年后，DISC在CPU及GPGPU体系结构的后端硬件上的性能有了显著的晋升，同时在新硬件的反对下面投入了更多的技术力量。2021年底，为了吸引更多的技术交换和单干共建须要，以及更大范畴的用户反馈，正式更名为BladeDISC并实现了初版开源。

五 关键技术

BladeDISC的整体架构，及其在阿里云相干产品中的上下文关系如下图所示：

1 MLIR基础架构

MLIR是由Google在2019年发动的我的项目，MLIR 的外围是一套灵便的多层IR基础设施和编译器实用工具库，深受 LLVM 的影响，并重用其许多优良理念。这里咱们抉择基于MLIR的次要起因包含其比拟丰盛的基础设施反对，不便扩大的模块化设计架构以及MLIR较强的胶水能力。

2 动静shape编译

上图为BladeDISC的主体Pass Pipeline设计。比照目前支流的深度学习编译器我的项目，次要技术特点如下：

图层IR设计

BladeDISC抉择基于HLO作为外围图层IR来接入不同的前端框架，然而HLO是本来为XLA设计的纯动态shape语义的IR。动态场景下，HLO IR中的shape表白会被动态化，所有的shape计算会被固化为编译时常量保留在编译后果中；而在动静shape场景下，IR自身须要有足够的能力表白shape计算和动静shape信息的传递。BladeDISC从我的项目建设开始始终与MHLO社区放弃严密的单干，在XLA的HLO IR根底上，扩大了一套具备齐备动静shape表达能力的IR，并减少了相应的基础设施以及前端框架的算子转换逻辑。这部分实现目前曾经残缺upstream至MHLO社区，确保后续其它MHLO相干我的项目中IR的一致性。

运行时Shape计算、存储管理和Kernel调度

动静shape编译的次要挑战来自于须要在动态的编译过程中可能解决动静的计算图语义。为齐备反对动静shape，编译后果须要可能在运行时做实时的shape推导计算，不仅要为数据计算，同时也须要为shape计算做代码生成。计算后的shape信息用于做内存/显存治理，以及kernel调度时的参数抉择等等。BladeDISC的pass pipeline的设计充分考虑了上述动静shape语义反对的需要，采纳了host-device联结codegen的计划。以GPU Backend为例，包含shape计算、内存/显存申请开释、硬件治理、kernel launch运行时流程全副为主动代码生成，以期失去齐备的动静shape端到端反对计划和更为极致的整体性能。

动静shape下的性能问题

在shape未知或者局部未知的状况下，深度学习编译器在性能上面临的挑战被进一步放大。在大多数支流硬件backend上，BladeDISC采纳辨别计算密集型局部和访存密集型局部的策略，以期在性能与复杂性和编译开销之间获取更好的均衡。

对于计算密集型局部，不同的shape要求更加精密的schedule实现来取得更好的性能，pass pipeline在设计上的次要思考是须要反对在运行时依据不同的具体shape抉择适合的算子库实现，以及解决动静shape语义下的layout问题。

而访存密集型局部的主动算子交融作为深度学习编译器次要的性能收益起源之一，同样面临shape未知状况下在性能上的挑战。许多动态shape语义下比拟确定性的问题，例如指令层的向量化，codegen模版抉择，是否须要implicit broadcast等等在动静shape场景下都会面临更大的复杂性。针对这些方面的问题，BladeDISC抉择将局部的优化决策从编译时下沉到运行时。即在编译期依据肯定的规定生成多个版本的kernel实现，在运行时依据理论shape主动抉择最优的实现。这一机制被称作speculation，在BladeDISC内基于host-device的联结代码生成来实现。此外，在编译期没有具体shape数值的状况下，会很容易在各个层级失落掉大量的优化机会，从图层的线性代数简化、fusion决策到指令层级的CSE、常数折叠等。BladeDISC在IR及pass pipeline的设计过程中着重设计了shape constraint在IR中的形象和在pass pipeline中的应用，例如编译期未知的不同dimension size之间的束缚关系等。在优化整体性能方面起到了比拟显著的作用，保障可能足够靠近甚至超过动态shape编译器的性能后果。

大颗粒度算子交融

团队在开启BladeDISC我的项目之前，已经基于动态shape编译器在大颗粒度算子交融及主动代码生成方面有过若干摸索3，其根本思维能够概括为借助于GPU硬件中低访存开销的shared memory或CPU中低访存开销的Memory Cache，将不同schedule的计算子图缝合进同一个kernel内，实现多个parallel loop复合，这种codegen办法称之为fusion-stitching。这种访存密集型子图的主动代码生成突破了惯例的loop fusion，input/output fusion对fusion颗粒度的限度。在保障代码生成品质的同时，大幅减少fusion颗粒度，同时防止复杂性及编译开销爆炸。且整个过程齐全对用户通明，无需人工指定schedule形容。


在动静shape语义下实现fusion-stitching比照动态shape语义下同样须要解决更大的复杂性，动静shape语义下的shape constraint形象肯定水平上简化了这一复杂性，使整体性能进一步靠近甚至超过手工算子实现。

3 多前端框架反对

AICompiler框架在设计时也蕴含了扩大反对不同前端框架的思考。PyTorch侧通过实现一个轻量的Converter将TorchScript转换为DHLO IR实现了对PyTorch推理作业的笼罩。MLIR绝对齐备的IR基础设施也为Converter的实现提供了便当。BladeDISC蕴含Compiler以及适配不同前端框架的Bridge侧两局部。其中Bridge进一步分为宿主框架内的图层pass以及运行时Op两局部，以插件的形式接入宿主框架。这种工作形式使BladeDISC能够透明化的反对前端计算图，能够适配用户各种版本的宿主框架。

4 运行时环境适配

为将编译的后果可能配合TensorFlow/PyTorch等宿主在各自的运行环境中执行起来，以及治理运行时IR层不易表白的状态信息等等，咱们为不同的运行时环境实现了一套对立的Compiler架构，并引入了运行时形象层，即RAL(Runtime Abstraction Layer)层。

RAL实现了多种运行环境的适配反对，用户能够依据须要进行抉择，具体包含：

  • 全图编译，独立运行。当整个计算图都反对编译时，RAL提供了一套繁难的runtime以及在此之上RAL Driver的实现，使得compiler编译进去后果能够脱离框架间接运行，缩小框架overhead。
  • TF中子图编译运行。
  • Pytorch中子图编译运行。

以上环境中在诸如资源管理，API语义等上存在差别，RAL通过形象出一套最小汇合的API ，并清晰的定义出它们的语义，将编译器与运行时隔来到来，来达到在不同的环境中都可能执行编译进去的后果的目标。此外RAL层实现了无状态编译，解决了计算图的编译之后，编译的后果可能被屡次执行时的状态信息处理问题。一方面简化了代码生成的复杂度，另一方面也更容易反对多线程并发执行（比方推理）的场景，同时在错误处理，回滚方面也更加容易反对。

六 利用场景

BladeDISC的典型利用场景能够不太严格的分为两类：其一是在支流的硬件平台上（包含Nvidia GPU，x86 CPU等）上作为通用、通明的性能优化工具，升高用户部署AI作业的人力累赘，进步模型迭代效率；另一个重要的利用场景是帮忙新硬件做AI场景的适配和接入反对。

目前BladeDISC曾经广泛应用在阿里外部和阿里云上内部用户的多个不同利用场景下，笼罩模型类型波及NLP、机器翻译、语音类ASR、语音TTS、图像检测、辨认、AI for science等等多种典型AI利用；笼罩行业包含互联网、电商、主动驾驶、平安行业、在线娱乐、医疗和生物等等。

在推理场景下，BladeDISC与TensorRT等厂商提供的推理优化工具造成良好的技术互补，其次要差异性劣势包含：

  • 应答动静shape业务齐备的动静shape语义反对
  • 基于compiler based的技术门路的模型泛化性在非标准模型上的性能劣势
  • 更为灵便的部署模式抉择，以插件模式反对前端框架的透明性劣势

下图为Nvidia T4硬件上几个实在的业务例子的性能收益数字：

在新硬件反对方面，目前广泛的状况是除了积攒比拟深厚的Nvidia等头部厂商之外，包含ROCM等其它GPGPU硬件普遍存在的状况是硬件的指标曾经具备相当的竞争力，但厂商受制于AI软件栈上的积攒绝对较少，普遍存在硬件算力无奈施展进去导致硬件落地利用艰难的问题。如前文所述，基于编译器的技术门路下人造对于硬件的后端具备肯定的泛化能力，且与硬件厂商的技术储备造成比拟强的互补。BladeDISC目前在GPGPU和通用CPU体系结构上的储备绝对比拟成熟。以GPGPU为例，在Nvidia GPU上的绝大部分技术栈能够迁徙至海光DCU和AMD GPU等体系结构相近的硬件上。BladeDISC较强的硬件泛化能力配合硬件自身较强的通用性，很好的解决了新硬件适配的性能和可用性问题。

下图为海光DCU上几个实在业务例子上的性能数字：

七 开源生态——构想和将来

咱们决定建设开源生态次要有如下的思考：

  • BladeDISC发祥于阿里云计算平台团队的业务需要，在开发过程中与MLIR/MHLO/IREE等社区同行之间的探讨和交换给了咱们很好的输出和借鉴。在咱们本身随着业务需要的迭代逐步欠缺的同时，也心愿可能开源给社区，在目前整个AI编译器畛域实验性我的项目居多，偏实用性强的产品偏少，且不同技术栈之间的工作绝对碎片化的状况下，心愿可能将本身的教训和了解也同样回馈给社区，心愿和深度学习编译器的开发者和AI System的从业者之间有更多更好的交换和共建，为这个行业奉献咱们的技术力量。
  • 咱们心愿可能借助开源的工作，收到更多实在业务场景下的用户反馈，以帮忙咱们继续欠缺和迭代，并为后续的工作投入方向提供输出。

后续咱们打算以两个月为单位定期公布Release版本。BladeDISC近期的Roadmap如下：

  • 继续的鲁棒性及性能改良
  • x86后端补齐计算密集型算子的反对，端到端残缺开源x86后端的反对
  • GPGPU上基于Stitching的大颗粒度主动代码生成
  • AMD rocm GPU后端的反对
  • PyTorch训练场景的反对

此外，在中长期，咱们在上面几个探索性的方向上会继续投入精力，也欢送各种维度的反馈和改良倡议以及技术探讨，同时咱们非常欢送和期待对开源社区建设感兴趣的同行一起参加共建。

  • 更多新硬件体系结构的反对和适配，以及新硬件体系结构下软硬件协同方法学的积淀
  • 计算密集型算子主动代码生成和动静shape语义下全局layout优化的摸索
  • 稠密子图的优化摸索
  • 动静shape语义下运行时调度策略、内存/显存优化等方面的摸索
  • 模型压缩与编译优化联结的技术摸索
  • 图神经网络等更多AI作业类型的反对和优化等

原文链接
本文为阿里云原创内容，未经容许不得转载。