动态尺寸模型优化实践之 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 运算，所以下图左侧中展示的三种可能输入都是合法的。在静态 shape 语义下我们很容易在编译期就区分开实际是那种情况，故而在编译时只需要对一种具体情况进行处理，而在动态 shape 语义下，由于编译期间无法进行区分，我们需要确保编译的结果在三种情况下都可以工作，因而会在计算图中引入显示的 shape 计算的 IR 以及 broadcast 操作（如下图右侧所示）。在这个例子中上层框架中一个普通的 Add OP 在动态 shape 语义下，也会被展开成一个复杂的子图。

也正因为上述的这些挑战，目前成熟的优化工具（如 TensorRT，XLA/TVM）对于动态 shape 支持都还比较有限。BladeDISC是阿里云计算平台 PAI 团队自研的一款原生支持动态 shape 的 AI 编译器。在过往一段时间我们开发 BladeDISC 优化动态 shape 模型的过程中，我们发现尽管不知道 shape 的具体的数值，但是通过充分发掘和利用 Tensor 的 shape 之间的结构化关系或者 Tensor shape 自身的分布特点，可以有效的解决上述挑战。在这里我们把 Tensor 的 shape 之间的结构化关系或者 Tensor shape 自身的分布统称为 shape constraint。更进一步的，我们发现通过将 shape constraint 作为第一等公民引入到 IR 中，可以最大化的发挥 shape constraint 的效能。本文中我们将介绍 BladeDISC 中 shape constraint IR 的定义以及如何利用它来辅助完成一系列动态 shape 语意下的优化以解决上述挑战，缩小与静态 shape 之间的性能差异。

动机

为什么选择 shape constraint？

在上一章节中我们分析了在动态 shape 语义下给做优化所带来的一系列挑战。我们发现使用 shape constraint 可以有效的解决这些优化上面临的困难。以下我们将分别介绍使用 shape constraint 如何有效解决上述的三个挑战。

应对挑战 1：跨 shape 可迁移性

在具体分析之前，我们先看 shape constraint 在本文中的定义。假设一个 tensor 的 rank 是 N，则该 tensor 的 shape 可以记为(d0, d1, ... dN-1)，其中di表示的是该 tensor 在第i个轴上所具有的大小，也记为第i个轴的dimension size。文本中讨论的 shape constraint 可以分为以下两类：

• shape 结构化约束。该类约束描述的是 dimension size 之间的相关关系，比如：dimension size 相等关系：某一个 tensor 的di与另外一个 tensor 的dj具有相同的大小，或者同一个 tensor 的di与dj具有相同的大小；tensor 元素个数相等：即一个 tensor 和另外一个 tensor 具有相同数量的元素；dimension size 乘积相等关系：比如reshape（[a, b, c, d]） -> [a*b, c*d]
  

• shape 分布约束。该类约束描述的是某个或某几个 dimension size 的（联合）分布，比如：di % 4 = 0，di != 0, di * dj=10;likely values: 即描述某个 dimension size 更大概率可能取到的值；value range：即描述某个 dimension size 可能的取值区间；

由上述定义本身我们可以立刻得到一个结论：由于无论是 shape 结构化约束还是分布约束都不依赖于具体的 shape 值，因此基于 shape constraint 而构建的优化策略天然具备跨 shape 可迁移性。

应对挑战 2：shape 关系断言分析

很多重要的优化都将一些 shape 关系的断言是否成立作为优化触发的前置条件，比如：

• 计算图化简。比如上文提到的消除冗余的 broadcast 节点的例子。

• 计算图 layout 全局优化。计算密集型算子的性能和其输入输出的数据排布（layout）有很强的关系，一个更合适的 layout 往往具有更好的性能。而一般最优的 layout 是随着 shape 而变化的，导致在动态 shape 语意下无法静态确定最优的 layout；这里一种优化策略是：将 shape 兼容的计算密集算子分到一组，只需要在组间插入 layout 转化的算子，而组内由于可以确认使用同一种 layout 而不必再插入 layout 转化的算子，从而提升性能；这里 shape 兼容的断言是优化触发的前置条件。

• fusion 决策。在做算子融合时并不是越多越好，只有将 shape 兼容的算子进行融合才会取得比较好的效果，同样里 shape 兼容的断言的判定是做算子融合的前置条件。

在动态 shape 语义下，由于不知道 shape 具体的值，symbolic shape 关系的断言的分析往往是困难的。symbolic shape 关系的断言可以看成是 symbolic dimension size 间的关系断言的逻辑关系表达式，因而问题可以转换成对于 symbolic dimension size 间的关系断言的判定。而由前述 shape constraint 定义可知，symbolic dimension size 间的关系断言本身是 shape constraint 一个实例。在这里我们把判定 symbolic dimension size 间的关系断言是否成立的这个问题转换成该断言是否是已知原子 shape constraint 的组合。这里“原子”的定义是不能够通过其他 shape constraint 的实例的组合得到。举个例子，假设我们需要判定tensor A: tensor<axaxf32>是否比tensor B: tensor<bxbxf32>具有更多的元素个数，这个问题经过转换过之后可以变成 dimension size 关系a > b是否成立。

在完成上述问题的转换之后，目前剩下的未解决的问题是：如何获得已知结果的原子 shape constraint。具体来说有以下几种方式：

• 由用户提供或在 JIT 编译期间自动捕获。比如用户可以提供关于模型输入 shape range，JIT 编译期间可以将捕获到的一组输入 shape 当作 likely value 注入编译流程等。

• 由 op 定义所携带的 shape consraint 信息。比如：elementwise op 的输入和输出应该具有相同的大小；mhlo.dynamic_reshape op 的输入和输出应该具有相同的元素个数；mhlo.concatenate op 的输入和输出在非拼接的轴上应该具有相同的大小；

• 分析 shape 计算的 IR 或者通过传播已有的 shape constraint 来获得新的信息，比如：

在充分利用上述来源原子 shape contraint 的基础上，我们可以大幅减少由于 shape 未知导致一些优化前置条件无法判断，进而导致优化无法生效的问题。

应对挑战 3：shape 计算开销及更复杂的计算图

在动态 shape 语义下，会引入更多的 shape 计算的开销，整个计算图会变得更复杂。

对于 shape 计算开销，shape 结构化约束我们可以抵消大量重复的 symbolic 计算，从而尽可能减小额外的开销。

对于计算图化简而言，我们一方面可以通过利用 shape 结构化约束消除一部分的冗余计算，比如前文中提到的由于 IB 问题引入的大量 broadcast op 可以在计算图化简中消除。剩下的无法利用 shape 结构化约束消除的 broadcast 可以进一步利用以下 shape 分布约束进行优化：IB 触发（即需要插入隐式的 broadcast）的概率远小于 IB 不触发的概率。通过生成带 IB 和不带 IB 两个版本的（如下图所示），让 common case 变得更快，从而提升期望性能；

为什么需要 shape constraint IR？

shape constraint 在动态 shape 语义下很有用，但是要用好它却并不容易。由于在整个 pass pipeline 中都可能会使用到 shape constraint 信息，因此我们需要一种方式将 shape constraint 信息在不同的 pass 之间进行传递。BladeDISC 早期时使用是一种隐式传递的策略，即每个 pass 在需要使用 shape constraint 信息时会通过分析一遍 IR 来重建 shape constraint 信息。不难看出在上述的方案中 shape constraint 本身并不是 IR 一部分，这种策略带来的问题是：

• 通过分析 IR 的方式一般只能够重建（部分）结构化约束信息，大部分的分布约束信息无法通过直接分析 data 计算 IR 来重建，而分布约束对于动态 shape 语义下的性能优化同样很重要；

• 在整个 pass pipeline 中 IR 会经历多次的 lowering （如 TF/Torch dialect -> mhlo -> lmhlo -> ...），在每次 lowering 的过程中都可能会丢失一部分 shape constraint 信息，比如下图中所展示的tf.SplitOp的例子。在这个例子中上层的tf.SplitOp会被 lower 成一系列的 mhlo 的SliceOp。根据tf.SplitOp的定义我们可以知道它的所有输出（假设有N个）应该具有相同的 shape，且输入在被拆分的轴上的 dimension size 可以被N整除，这些信息在我们 lower 到 mhlo 时如果只进行 data computation 的转换，而不进行 shape constraint 信息的转换，将会被丢失，从而使得后续的 pass 无法再利用相应的信息进行优化；

为了解决上述的问题，我们更进一步将 shape constraint 作为第一等公民引入到 IR 中，使得可以对结构化约束信息和分布约束信息进行统一的建模，同时也确保在整个 pass pipeline 中各个 pass 可以看到一致的 shape constraint 信息，进而更好的支持在不同的阶段完成各自适合的优化。

合作讨论

以上是近期我们在 shape constraint IR Part I 的分享，Part II 后续也会发出，敬请期待，更多相关信息欢迎加入 BladeDISC 用户群交流讨论。

欢迎大家加入 BladeDISC 用户群交流讨论，与我们一起共建。钉钉群号：44534789