共计 7184 个字符，预计需要花费 18 分钟才能阅读完成。
简介：欢送走进走进阿里云机器学习 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 麻利版中，用户能够能够参考开发文档来疾速体验。
作者：朱凯，赵文益，杨军
版权申明：本文内容由阿里云实名注册用户自发奉献，版权归原作者所有，阿里云开发者社区不领有其著作权，亦不承当相应法律责任。具体规定请查看《阿里云开发者社区用户服务协定》和《阿里云开发者社区知识产权爱护指引》。如果您发现本社区中有涉嫌剽窃的内容，填写侵权投诉表单进行举报，一经查实，本社区将立即删除涉嫌侵权内容。
