011、MLIR的Pattern Rewrite框架：DRR与C++ Rewrite

MLIR的Pattern Rewrite框架：DRR与C++ Rewrite上周帮团队排查一个MLIR自定义Pass的bug，现象很诡异：同一个IR经过两次Pass后，某些op的operand顺序莫名其妙变了，导致下游的bufferization直接崩掉。我盯着dump出来的IR看了半小时，最后发现是Pattern Rewrite的匹配优先级在作祟——一个DRR生成的pattern和一个手写的C++ Rewrite pattern互相覆盖了。这种问题在MLIR的pattern rewrite框架里太容易踩坑了，今天就把这块掰开揉碎讲清楚。为什么需要两套Rewrite机制MLIR的pattern rewrite框架本质上解决的是“如何在IR上做局部变换”的问题。你写一个Pass，无非就是遍历op，匹配某种结构，然后替换成另一种结构。但MLIR的IR是SSA形式的DAG，不是简单的树，所以匹配和替换的复杂度比AST rewrite高一个量级。MLIR给了两条路：DRR（Declarative Rewrite Rules）和C++ Rewrite。DRR用TableGen写声明式规则，自动生成C++代码；C++ Rewrite则是手写继承OpRewritePattern的类。两条路各有适用场景，但混用的时候坑特别多。DRR：看起来很美，但别被它骗了DRR的语法确实简洁，比如你想把addi和muli合并成一个fused_mul_add：def FuseMulAdd : Pat (addi (muli $x, $y), $z), (fused_mul_add $x, $y, $z) ;三行搞定。但实际项目里DRR有几个隐藏的坑。第一个坑：DRR生成的pattern是“贪婪”的。它默认会匹配所有可能的子图，而且匹配顺序由TableGen生成的代码决定，不是你写的顺序。我遇到过DRR把一个muli同时匹配到两个不同的pattern里，导致IR被重复改写。解决办法是在DRR里显式指定PatternBenefit，但很多人不知道DRR也支持这个：def FuseMulAdd : Pat (addi (muli $x, $y), $z), (fused_mul_add $x, $y, $z) ; // 这里踩过坑：不加benefit的话，默认benefit是1，容易被其他pattern覆盖 let benefit = 10;第二个坑：DRR对op的约束检查是“静态”的。它只能检查op的类型和operand数量，没法做动态的运行时检查。比如你想只在x和y都是float类型时才做融合，DRR里写不了这种条件。你必须在生成的C++代码里手动加NativeCodeCall，但那样又失去了DRR的简洁性。第三个坑：DRR生成的pattern无法处理“副作用”。如果你的rewrite需要修改op的属性、插入新的op、或者做复杂的类型推导，DRR基本无能为力。这时候就得老老实实写C++ Rewrite。C++ Rewrite：手写才是硬道理C++ Rewrite的核心是继承OpRewritePattern，重写matchAndRewrite方法。这里有个关键点：matchAndRewrite返回LogicalResult，成功返回success()，失败返回failure()。别写成return true或return false，MLIR的框架会检查返回值类型，写错了编译不过，但运行时行为可能诡异。一个典型的C++ Rewrite长这样：structFuseMulAddPattern