Pass管理：模式重写——给编译器的“替换卡片”

可视化演示

演示MLIR Pass模式重写机制，展示DAG重写流程、算子融合与常量折叠关键优化。

模式重写概述

老师傅走进工厂，看到的不再是一堆零件，而是一张“积木装配图”：每个加工工位是一个节点，工位与工位之间用传送带连接，表示谁的结果被谁使用。工厂要做优化，不是推翻这张图，而是在局部插入“替换卡片”。这张卡片的学名，就叫模式重写。

在 MLIR 中，这个“装配图”就是中间表示（IR）。IR 里的每个操作（Operation）是一个节点，操作之间的数据依赖是边，整体构成一张有向无环图（DAG）。一次 Pass 就是一次“工厂巡检”，而模式重写则是让巡检员不用手工调整每一个齿轮，只需检查图上是否有卡片描述的结构；一旦发现，就按卡片完成替换。

传统的优化 Pass 往往需要手写遍历 IR、判断条件、创建新操作、删除旧操作的代码。模式重写把这些细节交给框架：你定义“看到什么、换成什么”，框架负责寻找匹配、安全改写、维护依赖。MLIR 的 RewritePattern 就是一个可复用的模式对象，PatternRewriter 则是那个严格按操作规程动手的工人。

回到玩具工厂：老师傅不需要认识每个螺丝钉，他只需要看图、认卡片、动手换。MLIR 工程师也一样，把 IR 看成 DAG，把优化看成局部替换，复杂系统就变得可拆解、可组合。

匹配与替换

老师傅手中的卡片必须有两部分：上半部分是“匹配条件”，下半部分是“替换动作”。比如“匹配：两个固定齿轮带动一个减速箱；替换：一体化减速电机”。如果匹配条件不满足，卡片绝不能用；如果满足，动手时还要保证接口一致，不能让下游工位找不到输入。

MLIR 的 matchAndRewrite 正是这个结构。第一步，match 检查当前操作及其操作数是否符合 DAG 模式；第二步，rewrite 用 PatternRewriter 创建新操作、替换旧操作。匹配失败就换下一张卡片，匹配成功则继续处理其他位置。这就是“局部作用域”：每次只影响被匹配的那一小块子图，不动整张图。

系统里往往有几十张卡片，它们的优先级由“模式收益（benefit）”决定。收益越高的卡片越先尝试，避免随机顺序导致永远走不到理想优化。收益只是个相对排序，真正对错还要靠语义等价来把关。

工厂老师傅的经验就这样沉淀成一张张卡片：匹配条件越精准，误伤越少；替换动作越规范，流水线越安全。编写 Pass 时，你也能把优化的“判断”和“动手”分开，让每一部分都能独立测试。

DRR声明式框架

给老师傅写操作说明有许多方式。最笨的方法是直接写机器人底层指令：先移动机械臂，再拧两颗螺丝，最后检查接口。另一种方法是写“卡片”：只需写“看到什么、换成什么”，再由一个文书员自动生成底层指令。MLIR 的 DRR（Declarative Rewrite Rules）就是后者。

DRR 使用 TableGen 描述模式。开发者只需要声明一个 Pat 规则，写下源 DAG 和目标 DAG，其余遍历、匹配、IR 更新代码都由工具生成。mlir-tblgen 就是这个“文书员”，它读取 TableGen 规则，自动生成 C++ 的 RewritePattern 子类。这就是 RewriterGen 的工作：把声明式描述翻译成命令式代码。

这种做法的好处是：优化规则的意图一目了然；C++ 样板代码不再需要手工维护；新增一条优化时，只需增加几行 TableGen，而不是一整个类。DRR 让模式重写拥有了“配置化”的形态，特别适合表达算子融合、常量折叠这类结构清晰的变换。

用“卡片”而不是“机器人指令”，意味着优化知识可以被阅读、评审和复用。DRR 之所以重要，是因为它把“优化意图”和“实现细节”分离，让 AI 编译器里的规则库像工厂手册一样清晰。

DAG到DAG变换

老师傅最讲究“接口对得上”。替换一个减速箱时，输入轴位置、输出轴转速都必须和原来一致；否则下一道工序就会卡住。模式重写的本质也是如此：不是改一个字符，而是把一个“子 DAG”换成另一个“子 DAG”，并重新接好边界上的边。

MLIR 中，匹配的单元通常是一个操作及其操作数构成的子图。替换时，新子图要接收原来的输入值，也要产出原来的输出值给下游使用。PatternRewriter::replaceOp 这类 API 会主动把旧结果的所有使用者重定向到新结果，所以开发者不用手工维护“谁在用我”的关系。这就是 DAG-to-DAG 的“接口映射”。

除此之外，重写还要考虑约束与副作用。如果一个操作有读写内存等副作用，就不能随意删除或改变执行顺序。合法的模式重写必须保证：匹配条件足够强，替换结果与原行为语义等价，副作用要么保留、要么被证明可消去。

如果把 IR 比作电路图，DAG-to-DAG 就是“替换一个元件并重新焊线”。MLIR 框架替我们检查导线是否接错，工程师则专注设计更好的元件。

典型优化应用

工厂常用的三类卡片，正好对应模式重写的三个经典应用。

第一张卡片是“合并工序”：如果看到“乘法工位把结果送到加法工位”，就换成一台能直接计算 a*b+c 的复合机器。这张卡片对应算子融合，减少了中间结果的搬运和内存占用。第二张卡片是“预计算固定件”：如果看到两个常数值相乘，就提前算出答案，贴上新标签。这张卡片对应常量折叠。第三张卡片是“标准化摆放”：把同类零件统一摆到固定位置，让其他卡片更容易命中。这张卡片对应规范化（Canonicalization）。

这三类优化都非常适合用 DRR 声明。算子融合把结构性的计算合并成一个整体；常量折叠把运行期计算变成编译期计算；规范化消除冗余形态，比如 x+0 直接变成 x。它们不改变结果，只改变产生结果的方式。

玩具工厂用这三张卡片，就能让流水线越来越紧凑。AI 编译器也一样，把常见的计算模式写成卡片，编译器就能自动把笨重的计算图优化成精悍的执行图。

Pass集成与调试

卡片再漂亮，如果没有人拿着它去车间巡视，也不会产生价值。Pass 就是那个“巡视计划”，而模式重写的驱动方式则是“不断巡逻直到没有卡片可用”。

在 MLIR 中，开发者先把一组 RewritePattern 注册到 RewritePatternSet，再交给贪婪重写驱动（Greedy Rewriter）。它会反复扫描 IR：只要某张卡片能命中，就执行替换；一轮结束后继续下一轮，直到没有任何模式能匹配。这种行为叫“迭代到固定点”。注册规则的 Pass 可以很小，比如 -canonicalize 就是一个以模式重写为主的规范化 Pass。

调试时，mlir-opt 是最常用的实验室：用一个测试文件跑指定 Pass，观察 IR 的变化。配合 -debug-only=mlir-rewriter 这类开关，还能看到每条规则尝试和命中的日志。测试和调试的价值，就像给工厂做“试运行”：先在小块流水线上验证卡片不出错，再放开到完整生产线。

“巡逻到没有卡片命中”看起来简单，却能保证优化结果稳定。Pass 集成到编译流程后，AI 模型的计算图就会在编译期间被自动打磨，开发者只要观察日志和测试就能验证效果。

收益与挑战

模式重写真正改变的是“优化开发的方式”。过去，一个优化 Pass 可能涉及复杂的 IR 遍历、节点插入、删除和依赖更新；现在，大多数优化可以写成一张张高可读性的卡片，开发和维护成本大大降低。

但它也有挑战。卡片多了，可能出现“两张卡片同时命中同一个子图”的冲突，需要用模式收益和 Pass 顺序来裁决。更麻烦的是规则之间可能互相触发，造成无限循环。比如一张规则把 a+b 变成 b+a，另一张又把 b+a 变回 a+b，贪婪重写就会来回打转。因此，模式重写不是“规则越多越好”，而是要在收益与复杂度之间找平衡。

MLIR 通过模式收益、规范化、调试开关和测试流水线来管理这些挑战。每个模式都带着一个相对优先级，重写框架按优先级决定尝试顺序；规范化则提前统一 IR 形态，减少规则的排列组合。

老师傅的手册越厚，越需要编号和索引。模式重写让 AI 编译器变得可扩展，但真正成熟的产品，还得靠纪律严明的规则管理、优先级设计和充分测试。

AI 应用场景

Python 演示

"""mini_drr.py: 用 Python 模拟 MLIR DRR 的 DAG-to-DAG 模式重写""" def match(pattern, ir):    if isinstance(pattern, str) and pattern.startswith("$"):        return {pattern: ir}    if isinstance(pattern, tuple):        if not isinstance(ir, tuple) or len(pattern) != len(ir):            return None        if pattern[0] != ir[0]:            return None        binds = {}        for p, n in zip(pattern[1:], ir[1:]):            sub = match(p, n)            if sub is None:                return None            binds.update(sub)        return binds    return {pattern: ir} if pattern == ir else None  def apply_rules(node):    # 规则1: 常量折叠 mul(const, const) -> const    binds = match(("mul", ("const", "$x"), ("const", "$y")), node)    if binds is not None:        return ("const", binds["$x"] * binds["$y"])    # 规则2: 算子融合 add(mul(a, b), c) -> fma(a, b, c)    binds = match(("add", ("mul", "$a", "$b"), "$c"), node)    if binds is not None:        return ("fma", binds["$a"], binds["$b"], binds["$c"])    return node  def rewrite(node):    if not isinstance(node, tuple):        return node    while True:        result = apply_rules(node)        if result == node:            break        node = result    new_args = [rewrite(arg) for arg in node[1:]]    return tuple([node[0]] + new_args)  if __name__ == "__main__":    dag = ("add", ("mul", ("const", 3), ("const", 4)), ("const", 5))    print("before:", dag)    print("after :", rewrite(dag))

1. def match(pattern, ir):递归比较“模式卡片”和当前 DAG，返回变量绑定。
2. if isinstance(pattern, str) and pattern.startswith("$"):$x 是通配符，可绑定任意子节点。
3. ("mul", ("const", "$x"), ("const", "$y"))这是一个 DRR 风格的源 DAG 模式。
4. return ("const", binds["$x"] * binds["$y"])常量折叠的替换动作：直接在编译期算出结果。
5. while True:贪婪重写直到没有规则能改变当前节点，对应 MLIR 的固定点循环。

常见误区

面试题