在本系列分享中咱们将介绍 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]）-> [ab, cd]
-
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