ramy  2022-01-28 16:49:05  深度学习 |   查看评论   

从2020年第二季度开始至今,DISC持续投入研发力量,针对前文提到的从云端平台方视角看到的深度学习编译器距离大规模部署和应用的几个瓶颈问题,在性能、算子覆盖率和鲁棒性、CPU及新硬件支持、前端框架支持等方面逐渐完善。目前在场景覆盖能力和性能等方面,已经逐渐替换掉团队过往基于XLA和TVM等静态shape框架上的工作,成为PAI-Blade支持阿里内部及阿里云外部业务的主要优化手段。2021年后,DISC在CPU及GPGPU体系结构的后端硬件上的性能有了显著的提升,同时在新硬件的支持上面投入了更多的技术力量。2021年底,为了吸引更多的技术交流和合作共建需要,以及更大范围的用户反馈,正式更名为BladeDISC并完成了初版开源。

关键技术

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

MLIR基础架构

MLIR是由Google在2019年发起的项目,MLIR 的核心是一套灵活的多层IR基础设施和编译器实用工具库,深受 LLVM 的影响,并重用其许多优秀理念。这里我们选择基于MLIR的主要原因包括其比较丰富的基础设施支持,方便扩展的模块化设计架构以及MLIR较强的胶水能力。

动态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][4],其基本思想可以概括为借助于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抽象一定程度上简化了这一复杂性,使整体性能进一步接近甚至超过手工算子实现。

多前端框架支持

AICompiler框架在设计时也包含了扩展支持不同前端框架的考虑。PyTorch侧通过实现一个轻量的Converter将TorchScript转换为DHLO IR实现了对PyTorch推理作业的覆盖。MLIR相对完备的IR基础设施也为Converter的实现提供了便利。BladeDISC包含Compiler以及适配不同前端框架的Bridge侧两部分。其中Bridge进一步分为宿主框架内的图层pass以及运行时Op两部分,以插件的方式接入宿主框架。这种工作方式使BladeDISC可以透明化的支持前端计算图,可以适配用户各种版本的宿主框架。

运行时环境适配

为将编译的结果能够配合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应用;覆盖行业包括互联网、电商、赢咖4注册、安全行业、在线娱乐、医疗和生物等等。

在推理场景下,BladeDISC与TensorRT等厂商提供的推理优化工具形成良好的技术互补,其主要差异性优势包括:

  • 应对动态shape业务完备的动态shape语义支持
  • 基于compiler based的技术路径的模型泛化性在非标准模型上的性能优势
  • 更为灵活的部署模式选择,以插件形式支持前端框架的透明性优势

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

在新硬件支持方面,目前普遍的情况是除了积累比较深厚的Nvidia等头部厂商之外,包括ROCM等其它GPGPU硬件普遍存在的情况是硬件的指标已经具备相当的竞争力,但厂商受制于AI软件栈上的积累相对较少,普遍存在硬件算力无法发挥出来导致硬件落地应用困难的问题。如前文所述,基于编译器的技术路径下天然对于硬件的后端具备一定的泛化能力,且与硬件厂商的技术储备形成比较强的互补。BladeDISC目前在GPGPU和通用CPU体系结构上的储备相对比较成熟。以GPGPU为例,在Nvidia GPU上的绝大部分技术栈可以迁移至海光DCU和AMD GPU等体系结构相近的硬件上。BladeDISC较强的硬件泛化能力配合硬件本身较强的通用性,很好的解决了新硬件适配的性能和可用性问题。

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

某识别类模型 推理 不同batchsize下 2.21X ~ 2.31X
某检测类模型A 推理 不同batchsize下 1.73X ~ 2.1X
某检测类模型B 推理 不同batchsize下 1.04X ~ 1.59X
某分子动力学模型 训练 2.0X

开源生态——构想和未来

我们决定建设开源生态主要有如下的考虑:

  • BladeDISC发源于阿里云计算平台团队的业务需求,在开发过程中与MLIR/MHLO/IREE等社区同行之间的讨论和交流给了我们很好的输入和借鉴。在我们自身随着业务需求的迭代逐渐完善的同时,也希望能够开源给社区,在目前整个AI编译器领域实验性项目居多,偏实用性强的产品偏少,且不同技术栈之间的工作相对碎片化的情况下,希望能够将自身的经验和理解也同样回馈给社区,希望和深度学习编译器的开发者和AI System的从业者之间有更多更好的交流和共建,为这个行业贡献我们的技术力量。
  • 我们希望能够借助开源的工作,收到更多真实业务场景下的用户反馈,以帮助我们持续完善和迭代,并为后续的工作投入方向提供输入。

后续我们计划以两个月为单位定期发布Release版本。BladeDISC近期的Roadmap如下:

  • 持续的鲁棒性及性能改进
  • x86后端补齐计算密集型算子的支持,端到端完整开源x86后端的支持
  • GPGPU上基于Stitching的大颗粒度自动代码生成
  • AMD rocm GPU后端的支持
  • PyTorch训练场景的支持

此外,在中长期,我们在下面几个探索性的方向上会持续投入精力,也欢迎各种维度的反馈和改进建议以及技术讨论,同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建。

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

除特别注明外,本站所有文章均为 赢咖4注册 原创,转载请注明出处来自阿里 BladeDISC 深度学习编译器正式开源

留言与评论(共有 0 条评论)
   
验证码:
[lianlun]1[/lianlun]