做签证网站,worldpress英文网站建设,公司注册资金最低是多少,肇庆企业自助建站在本系列分享中我们将介绍BladeDISC在动态shape语义下做性能优化的一些实践和思考。本次分享的是我们最近开展的有关shape constraint IR的工作#xff0c;鉴于篇幅较长#xff0c;为了提升阅读体验#xff0c;我们将分享拆分为两个部分#xff1a;
Part I 中我们将介绍问…在本系列分享中我们将介绍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} : indexdisc_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匹配的逻辑的干扰 (matmulBiadAdd - 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粒度的明显改善以及访存密集型部分总消耗时间的减少。 原文链接
本文为阿里云原创内容未经允许不得转载。
