在本系列分享中咱们将介绍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