动态尺寸模型优化实践之 Shape Constraint IR Part I

在本系列分享中我们将介绍BladeDISC在动态shape语义下做性能优化的一些实践和思考。本次分享的是我们最近开展的有关shape constraint IR的工作，鉴于篇幅较长，为了提升阅读体验，我们将分享拆分为两个部分：

• Part I 中我们将介绍问题的背景，面临的主要挑战和以及我们做shape constraint IR的动机；
• Part II 中我们将介绍shape constraint IR的设计，实现以及一些初步的实验结果；

本篇是关于Part I的介绍，Part II的介绍将后续发出。

背景和挑战

主流的AI模型在部署时一般都具备不同程度的动态shape问题，比如输入图片尺寸，batch size 或者序列长度的变化等。与静态shape语义下做优化相比，在动态shape语义下做优化由于缺少具体的shape信息往往具有更大的难度，主要体现在以下几个方面：

挑战1：优化目标的转变。在静态shape下，我们的优化目标是希望在给定的shape下，尽可能逼近理论上限的速度，针对不同的shape可以使用不同的优化手段，而在动态shape语义下，我们优化目标是希望使用一套方法提升在整个shape分布上的平均性能，更强调的是优化方法的跨shape可迁移性和稳定性。因此很多在静态shape下常用的优化，比如profling驱动的策略，将不再简单可用。

挑战2：更少的有效信息。优化AI模型时很多常见的手段都把一些shape关系的断言是否成立作为优化触发的前置条件。比如在计算图化简时消除冗余的broadcast op，需要依赖于判断输入和输出是否具有相同的shape。在静态shape语义下，判断常量shape是否相等是显然的，而动态shape语义下，判断symbolic shape相等则困难的多，而一旦我们无法有效判断这些前置的shape关系断言是否成立，后续优化都无法进行，因而丢失很多优化机会，拉大与静态shape情况下性能的差异。

挑战3：更复杂的计算图。在动态shape语义下，由于shape不再是编译（或者优化）期间的常量，整个计算图中混杂着计算shape的IR以及计算data的IR，使得整个计算图的分析和优化都变得更复杂，同时也会引入更多shape相关计算的开销。

下图中展示了一个支持numpy语义implicit broadcast (IB)的Add OP的例子以说明计算图变复杂的具体过程。在IB语义下，Add OP在运行时会根据输入shape之间的关系，隐式的插入broadcast