简介：欢送走进走进阿里云机器学习PAI AICompiler编译器系列。近期，阿里云机器学习PAI团队全新上线一套Dynamic Shape Compiler框架，不仅作为AICompiler技术栈中原有的Static Shape Compiler框架的重要补充，更是减少了Compiler在企业级数据处理利用的有限可能，在晋升数据处理效率的同时，大幅晋升AI工程化效率。

挪动互联网的衰亡，不仅产生了海量数据，也对人机交互有了新的定义。企业如何动静解决不同规格图片数据，如何更灵活处理不同长度的对话语料等等，晋升企业经营效率，争取更多的商业机会和流量，成为泛滥企业摸索的热门技术利用。

近期，阿里云机器学习PAI团队全新上线一套Dynamic Shape Compiler框架，不仅作为AICompiler技术栈中原有的Static Shape Compiler框架的重要补充，更是减少了Compiler在企业级数据处理利用的有限可能，在晋升数据处理效率的同时，大幅晋升AI工程化效率。先来看看案例和成果数据。

性能后果
某TensorFlow语音辨认业务示例
以某业务方的语音辨认模型为例。过往咱们为这个业务提供的优化次要基于Static Shape Compiler，但因为shape变动范畴较大，只能采纳离线预编译的形式来实现部署，部署过程较为繁琐。

下图展现了基于Dynamic Shape Compiler在不同batchsize下的理论性能后果，其中纵轴为latency的晋升倍数。整体编译次数从之前的几千次升高到1次。从数字上来看，在只有一次编译的较小编译开销下，性能非常靠近Static Shape Compiler的性能优化后果。

某TensorFlow广告举荐业务示例
上面这个例子是某广告举荐业务方的推理模型，在线上零碎中，预估时的input shape变动十分频繁，比方：用户画像标签个数可能不同；用户的历史行为序列的长度会不同；召回广告的汇合大小会不同；广告所属类目标数量也会不同。这些变量会最终导致申请预估服务的input shape也是时刻变动的。
过往业务方的同学须要通过batching/手工干涉圈图等形式能力将编译次数管制到可承受范畴内，造成每次模型迭代后部署过程较为繁琐。从性能后果上看，比照基于XLA优化的性能，Dynamic Shape Compiler根本靠近甚至超过Static Shape Compiler的性能优化后果。

某TensorFlow语音合成业务示例
咱们以某业务方的TTS模型为例，具体看一下理论业务中对优化工具对动静shape反对的需要状况。在这个业务模型里，用户的输出sequence length输入sequence length都可能发生变化。此外，因为TTS类算法自身须要在推理过程中引入随机数的特点，即便输出同一条样本，外部子图的shape也会发生变化。咱们测试了如果基于static shape compiler来优化这个模型的话，输出数据数量以及累积编译次数的变动曲线。在这个业务模型中每个shape的编译开销约为20s左右，能够看到在通过几百轮迭代之后，编译cache都还没有收敛趋势，依据实践shape变动范畴测算总编译工夫至多在10个小时以上，因而这类业务如果应用XLA等动态shape编译器的话，无奈通明的实现可商用的部署。AICompiler外面的dynamic shape compiler组件很好的解决了这一问题，在一次编译的前提下帮忙用户取得均匀2X的性能收益，目前业界尚无其它AI编译器可能实现相似的成果。

某PyTorch公式辨认业务示例
下图是一个PyTorch业务的例子，比照的Baseline是基于libTorch执行导出后的TorchScipt脚本，能够看到AICompiler对PyTorch业务提供了同样的编译优化能力。

本文次要介绍这套动静shape编译框架，对更多技术细节趣味的读者能够参考DISC: A Dynamic Shape Compiler for Machine Learning Workloads.

从PAI团队三年前启动深度学习编译器方向的工作以来，“Dynamic Shape”问题始终是妨碍理论业务落地的重大问题之一。彼时，包含XLA在内的支流深度学习框架，都是基于Static Shape语义的编译器框架。即，just-in-time运行的编译器，会在运行时捕获待编译子图的理论输出shape组合，并且为每一个输出shape组合生成一份编译后果。

Static Shape Compiler的劣势不言而喻，编译期齐全已知动态shape信息的状况下，Compiler能够作出更好的优化决策并失去更好的CodeGen性能，同时也可能失去更好的显存/内存优化plan和调度执行plan；然而，Static Shape Compiler的毛病也非常显著，具体包含：
• 编译开销的减少。对于训练业务，编译开销导致训练迭代速度不稳固，训练初期显著负优化，甚至整个训练过程的工夫开销负优化；对于Inference业务，很多业务理论部署和迭代时不容许呈现性能抖动，而离线的预编译预热又会使得部署的过程变简单。
• 内存显存占用的减少。除编译开销的问题之外，当shape变动范畴特地大的时候，编译缓存额定占用的内存显存，常常导致理论部署环境下的内存/显存OOM，间接妨碍业务的理论落地。
• 对于一部分业务场景，shape变动范畴可能十分大甚至是趋于无穷的，比拟常见的包含广告举荐类业务中常见的稠密化模型，还有例如分布式训练下的embedding切片等等。在这种状况下，编译缓存永远也无奈收敛，用户也就不可能通过compiler获取到性能收益了。
• 上述问题在局部状况下，能够通过人工干预Compiler的圈图过程来缓解，即，将shape变动激烈的子图排除在编译范畴之外。然而，这种解决办法对用户十分不敌对，大大降低了Compiler利用的通用性和透明性，这要求做部署和优化的同学同时对模型构造和compiler十分理解，且每一次模型构造迭代时，都须要破费额定的工作量来调整圈图取得能够承受的性能成果。

对于这一问题，已经呈现过若干类解决方案，包含，对Compiler在圈图过程中的自动化干涉；在编译期外部主动对变动维度做bucketing补齐并将子图计算结果做主动的slicing。然而这些解决方案都存在各自的局限，例如前者只能适配于小局部子图shape变动激烈的状况，后者在很多模型上都无奈失去主动slicing的齐备数学推导。

为彻底解决这一问题，咱们抉择基于MLIR（Multi Layer Intermediate Representation），联合团队过往对AICompiler中积攒的局部教训，打造一套齐备反对Dynamic Shape语义的AI编译器，心愿可能彻底解决深度学习编译器在这部分对灵活性要求较高的业务中无奈落地利用的问题。

整体架构
Dynamic Shape Compiler的整体架构，及其在AICompiler中的上下文关系如下图所示。

Compiler局部
MLIR Infrastruction
MLIR是由Google在2019年发动的我的项目，MLIR 的外围是一套灵便的多层IR基础设施和编译器实用工具库，深受 LLVM 的影响，并重用其许多优良理念。
这里咱们抉择基于MLIR的次要起因包含：
• 比拟丰盛的基础设施反对，使得实现编译器的惯例开发工作更为便捷，效率更好。TableGen，以及编写惯例pattern matching的graph optimization pass的简化等。
• Open for Extension的模块化设计架构，这里的外围是其Dialect形象的设计。除Dialect的concept自身，在架构设计上，基于LLVM在传统编译期畛域的成功经验，MLIR团队还是展现出了幼稚的架构设计能力，将整个MLIR架构的设计变得很具模块化。
• MLIR的胶水能力，使得其能够比拟灵便不便地与曾经存在的优化伎俩进行集成，而非拒斥。

具体实现
MLIR框架的上述个性，使得咱们能够比拟不便的有选择性的leverage局部社区已有组件，防止齐全的从新造轮子，也肯定水平上防止从头彻底重构XLA代码带来的微小工作量。
这里咱们依据过往对AI编译器的了解，抉择了4层比拟次要的中间层形象，包含：
• DHLO Dialect：可能齐备表白动静shape语义的算子层计算图形象，它的次要作用是可能用无限数量的算子类型来形容不同前端框架的大量算子定义，且表白足够灵便。
• DLHLO Dialect：引入Buffer语义的计算图形象，用于在编译器流程中进行内存/显存的治理和优化。
• Loop Dialect：用于将算子层的计算形容基于Loop等开展为指令集的计算形容，咱们在这一层上实现了算子fusion的CodeGen。
• GPU Dialect：为GPU编程模型中的kernel launching及各种底层原语提供中间层形象。

下图展现了咱们基于MLIR的Loop Dialect等基础设施，在CodeGen中实现最简略的Input fusion的基本原理。比照XLA中只有高层的HLO和底层的llvm两层两头示意，MLIR提供的Loop Dialect形象能够间接在中间层实现fusion，很好的简化了开发的复杂度。

篇幅起因，咱们在次不在赘述Compiler局部其它各个模块的具体实现细节，请感兴趣的同学请移步MLIR社区中发动的相干细节探讨：RFC，以及会议探讨。

此处想着重介绍下比照于XLA，Dynamic Shape Compiler须要额定思考的一些问题，包含：
• DHLO IR，咱们在XLA的HLO IR根底上，扩大了一套具备齐备动静shape表达能力的IR。动态场景下，HLO IR中的shape表白会被动态化，所有的shape计算会被固化为编译时常量保留在编译后果中；而在动静shape场景下，IR自身须要有足够的能力表白shape计算和动静shape信息的传递。
• Placer模块，对于Dynamic Shape Compiler来说，计算能够分为shape计算和data计算两类，对于GPU backend而言，通常shape计算的计算量较小，launch和拷贝开销相比较大因而通常更适宜在host侧实现计算。咱们实现了一个简略的单卡分图策略，对host侧和device侧计算执行不同的lowering pipeline。
• Buffer治理及Buffer优化模块，有别于动态Shape编译期可能比拟容易通过liveness剖析，实现Buffer的复用等优化，而在动静shape语境下，因为Buffer Size未知编译期则不容易做到完全一致的优化。咱们目前应用的是动静的Buffer申请和开释，优化申请和开释的工夫点，同时后盾应用应用层蕴含Cache的Allocator，来达到性能和灵活性之间的均衡。后续可思考在IR中充沛表白Shape Constraint信息的状况下，来尝试在编译期做精密的Buffer复用优化。

此外，咱们留神到在动静shape语境会为编译期的性能performance带来一些乏味的新挑战：
• 局部优化决策后置到运行期，以Implicit Broadcast为例，目前支流的前端AI框架都反对implicit broadcast语义，而在动静shape语义下，编译期无奈充沛晓得LHS/RHS是否须要执行Broadcast操作。为保障齐备性，如果所有状况下都稳当的执行Broadcast计算的话，则会带来比较严重的冗余计算/Fusion颗粒度等问题。其它与之相似问题还包含GPU Kernel的Launch Dimension抉择等，咱们解决这一问题的做法是编译期做多版本编译，运行期依据理论shape来抉择最优实现，保障灵活性的同时，缓解灵活性带来的性能损耗。
• Shape束缚信息的应用，咱们发现在Dynamic Shape Compiler中，即便Tensor的Shape信息未知，但Shape之间的束缚信息，例如两个Tensor之间的某两个维度的size是否相等等信息，依然会对编译后果的性能产生比拟重要的影响。次要起因包含：在图层面，这些信息带来了更大的图优化空间，而在CodeGen层面，这些信息可能更无效的领导低层Lowering做CSE等传统编译器优化，缩小冗余的计算指令数。

多前端框架反对
随着近年来PyTorch用户数量的继续减少，对PyTorch作业的性能优化需要也正在变得越来越重要。AICompiler框架在设计时也蕴含了扩大反对不同前端框架的思考。

从IR lowering的角度，这里咱们抉择相比于HLO更具泛化表达能力的DHLO Dialect作为不同前端框架的对立接入IR，而在PyTorch侧抉择用户部署时导出的TorchScript IR，通过实现一个轻量的Converter将TorchScript转换为DHLO IR实现了对PyTorch Inference作业的笼罩。MLIR绝对齐备的IR基础设施也为Converter的实现提供了便当。

RAL (Runtime Abstraction Layer）
除编译自身的问题之外，咱们还面临其它一些问题，例如如何将编译的后果可能配合TensorFlow/LibTorch等宿主在各自的运行环境上下文中执行起来，如何治理运行时IR层不易表白的状态信息等等。咱们心愿为不同的运行时环境实现一套对立的Compiler架构，为此咱们引入了运行时形象层，即RAL层。RAL层次要负责解决如下问题：
Compile Once and Run Anywhere
RAL实现了多种运行环境的适配反对，用户能够依据须要进行抉择。
• 全图编译，独立运行。当整个计算图都反对编译时，RAL提供了一套繁难的runtime以及在此之上RAL Driver的实现，使得compiler编译进去后果能够脱离框架间接运行，缩小框架overhad，比方咱们在反对某语音ASR模型（类transformer网络）推理优化时，应用全图编译将框架开销从TF的22ms减小到4ms；
• TF中子图编译运行。RAL目前实现了TF Driver，能够反对在训练/推理场景中对圈出的子图进行编译执行；
• Pytorch中子图编译运行。RAL目前实现了Libtorch Driver，能够反对在推理场景中对圈出子图进行编译执行；

以上环境中在诸如资源（e.g. memory）治理，API语义等上存在差别，心愿可能引入一层形象对compiler侧屏蔽这些差别。RAL通过形象出一套最小汇合的API （RAL中称为Driver），并清晰的定义出它们的语义，这样compiler和runtime就能够在肯定层度上隔离开来，简化compiler的开发，同时通过提供这套API在不同环境下的实现，来达到在不同的环境中都可能执行编译进去的后果的目标。

Stateless编译
dynamic shape compiler实现一个计算图的编译之后，编译的后果可能被屡次执行，而有些op的执行是带状态的：
• 在device（e.g. gpu）上执行时，对const op心愿只在第一次执行时加载并常驻device，而不是每次都引入一次host-to-device的拷贝；
• 对于须要依据具体shape信息进行tuning的op （e.g. gemm/conv），tuning cache须要一个中央存储；

RAL将资源初始化等带状态的局部抽取进去，封装成context来治理生命周期。在代码生成的过程中，通过简略的注入context，将状态量暗藏在context之后，使得compiler侧看到的是一个纯计算的过程。无状态的设计一方面简化了代码生成的复杂度，另一方面也更容易反对多线程并发执行（比方推理）的场景，同时在错误处理，回滚方面也更加容易反对。

对用户通明的编译模式切换
咱们对于Dynamic Shape Compiler在AICompiler中的定位是：与原Static Shape Compiler并列的一套框架，在容许适度就义性能的状况下，提供对于强Dynamic Shape类业务的通用通明反对。

然而从用户的角度来说，通常并不容易判断一个Workload的更适宜Dynamic Shape Compiler还是Static Shape Compiler，为此咱们联合接耦和全量关上[link]中的工作，设计了一套编译模式主动切换的状态机。其基本思路是，在工作初期先抉择较为平安的Dynamic Shape Compiler，联合后盾编译让用户可能在运行时尽早失去有性能晋升的编译执行，并在后续执行过程中联合资源的理论占用状况和理论运行时的shape变动范畴来有选择性的切换到Static Shape Compiler的执行。

两套compiler在运行时的切换关系如下图所示：

阿里云机器学习PAI平台面向企业客户及开发者，提供轻量化、高性价比的云原生机器学习平台，涵盖交互式建模、拖拽式可视化建模、分布式训练到模型在线部署的全流程笼罩，百余种落地场景，全面晋升机器学习工程效率。目前，
PAI AICompiler曾经集成在阿里云机器学习PAI的通用推理优化工具PAI-Blade麻利版中，用户能够能够参考开发文档来疾速体验。

作者：朱凯，赵文益，杨军

版权申明：本文内容由阿里云实名注册用户自发奉献，版权归原作者所有，阿里云开发者社区不领有其著作权，亦不承当相应法律责任。具体规定请查看《阿里云开发者社区用户服务协定》和《阿里云开发者社区知识产权爱护指引》。如果您发现本社区中有涉嫌剽窃的内容，填写侵权投诉表单进行举报，一经查实，本社区将立即删除涉嫌侵权内容。