模式重写Pattern Rewriting：让编译器学会“看菜改谱”

可视化演示

展示TVM模式重写中声明式子图匹配、算子融合改写与布局变换规则，交互切换查看图变换核心流程。

基础概念：从“手工改菜谱”到“规则手册”

在魔法厨房的旧时代，大厨想要优化一道菜，必须亲手把菜谱从头到尾读一遍，找到需要改的地方，用笔划掉，再写上新的步骤。传统Pass正是如此——编译器的每次优化都要你显式遍历中间表示（IR）的每个节点，手动判断是否满足条件，再手动修改。这种命令式的开发方式，写起来琐碎，读起来费劲，改起来还容易出错。

于是，魔法的理念升级了：为什么不让助手替我“盯”着菜谱呢？你只需在规则手册里写一条：“凡是出现'土豆+肉块'的组合，就替换为'南瓜+肉块'”。助手会自己满厨房找土豆炖肉，找到就改，改完接着找下一处。这套“声明式”的思维方式，正是Pattern Rewriting的核心。

在MLIR中，Pattern Rewriting框架允许开发者通过继承RewritePattern类来定义“匹配条件”和“替换逻辑”。当一个Pattern被注册到驱动引擎后，框架会自动在IR图中查找满足条件的子图（匹配），并将该子图替换为新的IR结构（替换）。这与传统Pass中手写遍历和修改逻辑形成鲜明对比——开发者从“命令者”变成了“规则制定者”。

回到厨房比喻：大厨从“亲手一行行改菜谱”变成“写一页规则手册”，而助手（框架）自动执行。这就是Pattern Rewriting带来革命性改变的第一步。

RewritePattern类：一条规则的诞生

一条完整的菜谱优化规则包含三要素：什么菜（匹配条件）、怎么改（替换逻辑）、省多少时间（收益值）。在Pattern Rewriting中，一条RewritePattern同样由这三要素构成。开发者不再需要关注“如何遍历”和“何时应用”，只需把这三要素声明清楚，框架就会替你执行。

在MLIR源码中，开发者通常继承mlir::RewritePattern基类，构造函数中传入PatternBenefit（收益值），并实现两个核心方法：match与rewrite。match方法负责检查某个Operation是否满足条件，如果满足则返回成功（并绑定变量），否则返回失败；rewrite方法负责实际的替换，将匹配到的旧操作替换为新的操作序列。

一条RewritePattern就像厨房规则手册中的一页：写明“看到什么、换成什么、多划算”,三要素齐全，助手才能准确执行而不打扰大厨。

模式匹配：助手的识别术

助手怎么在菜谱中认出“土豆+肉块”？它需要比对：是不是有一行写土豆，接下来有一行写肉块，而且它们之间存在正确的逻辑关系。在IR图中识别目标子图，也是同样的道理——match阶段需要完成操作数匹配、参数绑定和约束检查三件事。

操作数匹配要求待匹配Operation的类型、操作数数量、操作数类型满足模式设定的预期；参数绑定则将匹配到的子表达式绑定到变量上，供rewrite阶段引用；约束检查则会附加更严格的条件，例如“操作数必须是整数类型”或“属性的具体取值必须满足某范围”，保证模式不会被误触发。

助手的识别术练就得越精准，规则误伤的概率就越低——这既依赖匹配条件的完备，也依赖约束检查的严谨。

应用机制：规则的自动调度

规则手册里有一百条规则，助手该如何执行？是先处理所有土豆炖肉，再去处理红烧肉，还是交叉进行？MLIR给出了明确答案：通过驱动引擎自动调度。GreedyRewriteDriver是其中最常见的引擎，它会反复遍历IR图，对所有节点尝试应用所有模式，直到没有模式可匹配为止。

贪婪应用策略的核心是：每次迭代优先选择收益值（benefit）最高的模式进行替换，使得每次变换都朝着“最有利”的方向前进。模式排序则通过PatternBenefit比较器完成，在注册时提前排好优先级队列。这种自动化调度把开发者从繁琐的“何时调用哪个优化”中解放出来。

厨房助手按优先级逐条执行规则，每改完一轮再从头检查一遍，直到没有可改为止——这套调度机制，正是驱动引擎在IR上默默做的事。

重写编写：写一条好规则的艺术

好规则要精准：既不能太泛（把鸡鸭鱼肉全变成白菜豆腐），也不能太窄（只针对某一道特制的菜）。在Pattern Rewriting中，编写一条高质量的重写规则需要把握定义模式、创建新操作、删除旧操作三个环节。

定义模式时，匹配侧要具体描述目标IR结构，避免过度泛化；创建新操作时，使用PatternRewriter提供的replaceOp或create方法生成替代Operations，框架会维护IR的一致性；删除旧操作时，调用erazeOp安全移除被替换的节点，确保没有悬挂引用。这些操作必须通过PatternRewriter完成，直接修改操作指针容易破坏IR的内部一致性。

一道标准菜讲究“营养均衡、口味稳定”，一条好规则讲究“精确匹配、安全替换”。写规则的艺术，就是在精准与通用之间找到平衡点。

框架优势：为什么厨房需要这套系统

魔法厨房从“大厨亲手改菜谱”进化到“助手按规则自动改菜谱”，收益是巨大的。最直接的变化是：新来的帮厨不需要从零学习大厨的全部手艺，只需读懂规则手册就能上手——这就是“降低门槛”的力量。

Pattern Rewriting框架的优势在工程实践中尤为突出。降低门槛：新开发者只需理解“匹配-替换”的思维，无需掌握复杂的IR遍历细节；代码简化：上百行的手工遍历代码可能浓缩为几十行的规则定义；可维护性：每条规则独立成类，便于单元测试、复用和替换。这些优势让MLIR生态中的Pass开发效率大幅提升。

这套系统让魔法厨房真正实现了“一份规则、处处生效”。对AI编译器工程师来说，Pattern Rewriting不仅仅是工具，更是一种优雅的工程思维。

AI 应用场景

Python 演示

"""模拟MLIR Pattern Rewriting的简易表达式重写引擎""" class Expr:    """表达式节点：op是操作名，args是操作数列表"""    def __init__(self, op, *args):        self.op = op        self.args = args     def __repr__(self):        if not self.args:            return self.op        inner = ", ".join(repr(a) for a in self.args)        return f"{self.op}({inner})"  class Pattern:    """一条重写规则：匹配函数 + 替换函数 + 收益值"""    def __init__(self, name, match_fn, rewrite_fn, benefit):        self.name = name        self.match_fn = match_fn        self.rewrite_fn = rewrite_fn        self.benefit = benefit  def match_add_zero(node):    """匹配形如 x + 0 或 0 + x 的子表达式，返回绑定字典或 None"""    if node.op == "add" and len(node.args) == 2:        if isinstance(node.args[1], Expr) and node.args[1].op == "0":            return {"x": node.args[0]}        if isinstance(node.args[0], Expr) and node.args[0].op == "0":            return {"x": node.args[1]}    return None  def rewrite_add_zero(bindings):    """x + 0 -> x，直接返回被绑定的一侧"""    return bindings["x"]  def apply_rules(expr, patterns):    """自底向上递归应用规则，等于一个简化版GreedyRewriteDriver"""    # 1. 先重写所有子表达式    new_args = []    for arg in expr.args:        if isinstance(arg, Expr):            new_args.append(apply_rules(arg, patterns))        else:            new_args.append(arg)    expr = Expr(expr.op, *new_args)    # 2. 在当前节点上尝试每条模式    for p in patterns:        bindings = p.match_fn(expr)        if bindings is not None:            new_expr = p.rewrite_fn(bindings)            return apply_rules(new_expr, patterns)    return expr  if __name__ == "__main__":    # 构造表达式 (x + 0) * 2 + 0    tree = Expr(        "add",        Expr("mul", Expr("add", Expr("x"), Expr("0")), Expr("2")),        Expr("0"),    )    print("原始表达式:", tree)    patterns = [Pattern("add-zero", match_add_zero, rewrite_add_zero, 1)]    result = apply_rules(tree, patterns)    print("重写后结果:", result)

1. class Pattern:定义一条重写规则，对应MLIR中的RewritePattern类。
2. self.match_fn = match_fn存储匹配函数，相当于重写match方法。
3. def apply_rules(expr, patterns):核心驱动引擎，相当于GreedyRewriteDriver。
4. new_args.append(apply_rules(arg, patterns))自底向上的递归重写，先处理子表达式。
5. bindings = p.match_fn(expr)调用匹配函数，返回{"x": ...}这样的参数绑定。
6. return apply_rules(new_expr, patterns)替换成功后对新表达式继续检查，直到稳定。

常见误区

面试题