简介：在本系列分享中咱们将介绍 BladeDISC 在动静 shape 语义下做性能优化的一些实际和思考。本次分享的是咱们最近发展的无关 shape constraint IR 的工作，Part II 中咱们将介绍 shape constraint IR 的设计，实现以及一些初步的试验后果在本系列分享中咱们将介绍 BladeDISC 在动静 shape 语义下做性能优化的一些实际和思考。本次分享的是咱们最近发展的无关 shape constraint IR 的工作，鉴于篇幅较长，为了晋升浏览体验，咱们将分享拆分为两个局部：Part I 中咱们将介绍问题的背景，面临的次要挑战和以及咱们做 shape constraint IR 的动机；Part II 中咱们将介绍 shape constraint IR 的设计，实现以及一些初步的试验后果；本篇是对于 Part II 的介绍，Part I 的介绍请参考这里。设计和实现 shape constraint IR 的设计应用 IR 来建模 shape constraint 并不是一个很容易的事件。咱们须要设计一种计划既不便咱们做 shape constraint 剖析，同时又不会使得 IR 的后续变换变的很简单。通过屡次迭代之后，咱们抉择了 type-based 计划来构建 shape constraint IR，基本思路如下段伪 IR 所示。// original data-computation IR
func @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 粒度的明显改善以及访存密集型局部总耗费工夫的缩小。
原文链接：https://click.aliyun.com/m/10… 本文为阿里云原创内容，未经容许不得转载。