作者 | 姜来
来历 | 阿里技能大众号

一 导读

跟着深度学习的不断开展，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作业类型的支撑和优化等

开源地址：

https://github.com/alibaba/BladeDISC

参考文献

1. "DISC: A Dynamic Shape Compiler for Machine Learning Workloads", Kai Zhu, Wenyi Zhao, Zhen Zheng, Tianyou Guo, Pengzhan Zhao, Feiwen Zhu, Junjie Bai, Jun Yang, Xiaoyong Liu, Lansong Diao, Wei Lin
2. Presentations on MLIR Developers' Weekly Conference: 1, 2
3. "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. The 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2022. [to appear]
4. "FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads", Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang, and Wei Lin. arXiv preprint