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

• Part I 中我们将介绍问题的背景，面临的主要挑战和以及我们做 shape constraint IR 的动机；

• Part II 中我们将介绍 shape constraint IR 的设计，实现以及一些初步的实验结果；

本篇是关于 Part II 的介绍，Part I 的介绍请参考https://xie.infoq.cn/article/be6f6596f2b52839e1cfce88c

设计和实现

shape constraint IR 的设计

使用 IR 来建模 shape constraint 并不是一个很容易的事情。我们需要设计一种方案既方便我们做 shape constraint 分析，同时又不会使得 IR 的后续变换变的很复杂。经过多次迭代之后，我们选择了 type-based 方案来构建 shape constraint IR，基本思路如下段伪 IR 所示。

// original data-computation IRfunc @main() {  ...  %0 = any_dialect.any_operation(...) : tensor<?x?xf32, [@S0, @S1]>  ...}
disc_shape.SymbolicDim @S0 {  range list : [[...], [...], ...]  likely_values : [...]  ...  symbolic_shape_graph: @shape_constraint_graph}
disc_shape.SymbolicDim @S1 {  range list : [[...], [...], ...]  likely_values : [...]  ...  symbolic_shape_graph: @shape_constraint_graph}
// A separated function to store shape constraint predicates between different symbolic dimensions.// Each symbolic dim is either bound to a `disc_shape.dim` op or `disc_shape.bind_dim`func @shape_constraint_graph(...) {  %0 = disc_shape.dim() {ref: @S0} : index  %1 = disc_shape.dim() {ref: @S1} : index  disc_shape.tie_predicate_divisible(d0, d1) // d0 % d1 == 0
  // other tie_* ops  //   disc_shape.tie_predicate_eq(d0, d1)  // d0 == d1  //   disc_shape.tie_predicate_lt(d0, d1)  // dim less than  //   disc_shape.tie_predicate_mul_eq(d0, d1, d2, ...) // d0 = d1 * d2 * ...  //   // d0 * d1 = s0 * s1  //   disc_shape.tie_predicate_product_eq([d0, d1, ..], [s0, s1, ...])  //   // d0 = affine.apply(d1, d2, ...) {affine_attr = ...}  //   disc_shape.tie_predicate_affine_eq(d0, d1, d2, ...) {affine_attr = ...}}

在这个方案中，每一个 symbolic dimension size（也即在编译期间无法确定具体大小的 dimension size）对应一个全局的disc_shape.SymbolicDimIR 对象。该 IR 对象中存储了关于这个 symbolic dimension size 的分布相关的约束，同时也存储了对一个 shape constraint function 的引用。在上图中最下面的部分是一个 shape constraint function 的例子。在这个 function 中存储的是 symbolic dimension dim 之间的相关关系（结构化约束），每一种相关关系用一个 op 来抽象，比如说这个例子中说展示的整除等价关系便是由tie_predicate_divisibleop 来描述。选择 type-based 方案主 data 计算图中并不会直接存储 shape constraint 信息。在上图中的最上面是一个主 data 计算图的例子。主 data 计算图中，每一个 tensor 对应的 type 中都包含一个 attribute，这个 attribuet 中存储了这个 tensor 所对应的 symblic dimension size 的引用。

通过将描述 shape constraint 的 IR 和主 data 计算 IR 解耦开，一方面可以尽可能减少对已有 pattern 匹配的逻辑的干扰 (matmul+BiadAdd -> FusedMatmulBiasAdd这个 pattern 替换并不需要感知到 shape constraint IR 的存在)，另外一方面，不同层级的 data 计算的 IR，比如 tensor level IR 和 buffer level IR，可以用同一套 shape constraint 的描述。从而可以缓解 IR lowering 过程中 shape constraint 信息的丢失问题。

基于 shape constraint IR 的优化 pipeline

将 shape constraint IR 作为第一等公民引入 IR 中之后，我们进一步构建了以 shape constraint 为中心的优化 pipeline（如下图所示）。通过对 shape constraint 的充分的挖掘，而非依赖于具体的 shape 的值来辅助完成各种优化，从而实现在动态 shape 语义下尽可能接近静态 shape 优化工具的性能。

下图中展示了目前 BladeDISC 中主要的几个优化的阶段。从最左边开始看起。第一步是将前端 AI 框架的计算图 lower 到 MHLO 的计算描述。这里值得注意的是除了普通的 data 计算的 lowering，还包含 shape constraint 的 lowering，从而避免在动态 shape 语义信息的丢失。到 MHLO 之后，我们首先会完成 shape constraint 的分析以及分析结果的 IR 化表示。分析得到的结果将可以指导我们完成计算图上的一些基本化简，比如冗余 broadcast op 的消除，layout 调整等。优化完之后的计算图，我们会进一步对其中的访存密集型算子做融合优化，shape constraint 将是决定那些算子可以融合的很重要的判断依据，通过更充分的挖掘 shape constraint，我们可以找到更多融合的机会。最后在做代码生成的时候，我们发现在动态 shape 语义下 index 计算的开销更加容易成为瓶颈，尤其是当算子融合的数目比较多的时候，利用 shape constraint 我们可以大幅消除冗余的 index 计算。以上限制于篇幅并未一一展开进行介绍，感兴趣的同学可以通过这里了解更多的细节。

初步测试

我们目前已经完成了 shape constraint IR 的第一阶段开发，也即 shape constraint IR 的引入以及 pass pipeline 的适配性改造。第一阶段的主要目标是搭建好整体的架子，还并未包含所有设计中优化的实现 (比如 likely value 的应用)，我们将会在后续持续迭代完善 shape constraint IR。在以上前提下，我们在一些比较典型的模型上完成初步的评测，下图展示的是部分的评测结果。由于目前 shape constraint IR 还未应用到计算密集型算子(GEMM/CONV 等)的优化 ，故以下测评针对的是模型中访存密集型部分，主要从两个维度来衡量：

• 访存密集型部分 kernel launch 的次数，即衡量 fusion 的粒度，同等情况下次数越少，fusion 的粒度越大；

• 访存密集型部分总消耗时间，即衡量生成的 kernel 的质量，同等情况下总时间越短，质量越高；

如下图中所示，在 CPU 和 GPU 上我们都观测到 fusion 粒度的明显改善以及访存密集型部分总消耗时间的减少。

合作讨论

以上是近期我们在 shape constraint IR Part II 的分享，更多相关信息欢迎加入 BladeDISC 用户群交流讨论。