简介： 在本系列分享中咱们将介绍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...本文为阿里云原创内容，未经容许不得转载。