flowchart TD
subgraph Driver["GreedyPatternRewriteDriver"]
Queue["待遍历 Operation 工作队列 (Worklist)"]
Match{"Pattern::matchAndRewrite"}
Queue -->|Pop Op| Match
Match -->|匹配成功| Rewrite["Rewriter: replaceOp / eraseOp"]
Match -->|不匹配| Next["考察队列下一个 Operation"]
Rewrite -->|触发新修改| Requeue["将当前节点的用户与邻接 Op 重新压入队列"]
end
第 4 课:Pass、PatternRewriter 与贪心重写
参考答案版:包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。
本课目标与完成标准
学完后你应能:编写局部 rewrite,解释 benefit、收敛与 PatternRewriter 修改纪律。
通过必须同时满足:
- 独立补齐本课唯一的代码填空题,并通过给定检查;
- 三个问答题均说明因果链,而不是只报术语;
- 能指出至少一个正确性边界和一个性能取舍;
- 总分不低于 8/10,且没有一票否决级概念错误。
前置关系
- 课程前置:编译原理基础、C++ 阅读能力
- 本课在路线中的作用:Pass 决定遍历/阶段,RewritePattern 匹配局部 IR;PatternRewriter 记录替换,维护 def-use 与监听器不变量。
核心心智模型
1. 它是什么,解决什么问题
在编译器优化与方言转换中,代码变换绝大多数可以形式化为“树/图局部重写(DAG Pattern Matching & Rewriting)”:识别一个局部的子图模式(如代数恒等式 、乘 0 折叠 、算子融合与降级),并将其原子化替换为一个更简明或更底层的子图。 若采用传统手写嵌套循环遍历 AST,不仅代码极度臃肿,且在边遍历边修改 IR 时极易引发迭代器失效(Iterator Invalidation)与悬空指针。 MLIR 提供了极其工业级、自动维护图拓扑不变式的核心引擎——GreedyPatternRewriteDriver 与 PatternRewriter。 本课解决的核心工程问题是:理解基于工作队列(Worklist)的不动点迭代(Fixed-Point Iteration)机制,掌握严格通过 PatternRewriter 纪律修改 IR 的防崩溃契约,以及设计严格单调递减度量消除重写死循环(振荡)。
2. 它如何工作
贪心模式驱动器状态机与不动点收敛流程如图 4-1 与图 4-2 所示:
图 4-1:MLIR 贪心模式重写驱动器 (Greedy Pattern Driver) 核心状态机
flowchart LR
Init["初始化 Pass 并在 Op 上启动 Driver"] --> Loop["弹出工作队列顶部 Operation"]
Loop --> TryPatterns["遍历已注册 RewritePattern 列表"]
TryPatterns --> Matched{"match 成功?"}
Matched -->|是| Apply["执行 rewrite() 替换节点并通知监听器"]
Matched -->|否| More{"队列为空?"}
Apply --> More
More -->|否| Loop
More -->|是| Converged["IR 达到不动点 (Fixed Point),Pass 顺利收敛"]
图 4-2:MLIR 模式重写与死代码消除的不动点迭代流程
工作队列初始化与拓扑遍历: Pass 收集一组已注册的
RewritePattern,创建GreedyPatternRewriteDriver并将目标 Region 内的所有 Operation 压入工作队列(Worklist)。模式匹配与原子重写(
matchAndRewrite):- 驱动器从 Worklist 中 Pop 出一个候选根节点
op; - 按照预设的优先级(Benefit)从高到低依次尝试匹配注册的 Pattern;
- 若匹配成功,进入
rewrite阶段。必须通过传入的PatternRewriter &rewriter执行修改,如:rewriter.replaceOp(op, newValues);rewriter.eraseOp(op);
- 驱动器从 Worklist 中 Pop 出一个候选根节点
监听器通知与连锁反应(Listener Notification & Re-queuing): 如图 4-1 所示,
PatternRewriter内部挂载了监听器(RewriterListener)。当执行replaceOp时:- 驱动器自动将旧
op的所有使用方(Users)、操作数的生产者以及新创建的 Operation 重新压回 Worklist; - 因为替换旧节点可能使下游节点的某些代数化简条件瞬间满足(例如下游原本是
y = addi(%0, 0),当%0被消除后,下游可能进一步化简)。
- 驱动器自动将旧
不动点收敛(Fixed-Point Convergence): 循环往复推进,直到 Worklist 彻底清空(达到 IR 不动点),或者达到预设的最大迭代次数上限(避免死循环),整个 Pass 顺利收敛退役。
3. 正确性条件与常见误区
严禁绕过 Rewriter 直接调用底层指针 API(Direct Mutation Violation): 在 Pattern 内部绝对严禁直接调用
op->erase()或op->replaceAllUsesWith()! 因为底层指针被强行释放后,GreedyPatternRewriteDriver的 Worklist 并不知道该节点已死。当驱动器后续尝试从 Worklist 访问该指针时,会直接触发野指针访问段错误(Use-After-Free Crash)。必须全部通过rewriter.eraseOp()或rewriter.replaceOp()通知监听器同步更新队列。重写振荡与死循环(Rewriting Oscillation): 若 Pattern A 将 变换为 (例如将除法变为乘倒数:),而 Pattern B 又将 变换回 (将乘倒数合并为除法:),两个 Pattern 会在同一个节点上反复交替命中并互相重排,导致 Worklist 永远无法清空,编译器彻底卡死在死循环中。
单调性度量(Monotonicity Metric): 每个规范化 Pattern 必须具备明确的“单调性度量”(例如 Operation 总数严格减少、常量向右折叠排序、或类型层级单向降低),杜绝等价互换。
4. 性能与工程取舍
Benefit 优先级调优: 每个 Pattern 具备一个
PatternBenefit(无符号整数)。Benefit 越高的 Pattern 会被优先尝试匹配。将计算代价小、收益明确的简单折叠(如常量消除)赋予高 Benefit,能大幅减少复杂大子图 Pattern 的无效匹配计算。局部贪心 vs 全局最佳(Greedy vs Global Dynamic Programming): 贪心驱动器寻找的是局部不动点,执行速度极快;但局部贪心替换可能会错失需要跨作用域全局收益评估的最优代码生成方案。对于复杂的内存布局选择或张量降级,往往结合专用图搜索或 Transform 方言控制。
具体演示
以恒等式代数化简 x + 0 -> x 为例:
- 输入 IR:
%1 = arith.addi %x, %c0 : i32(其中%c0是常量 0)。 - Pattern 命中:
op.getLhs()获取到%x。 - 执行重写:
rewriter.replaceOp(op, x);。 - 状态转移:所有原本消费
%1的下游指令,其操作数指针被原子化无缝重定向(RAUW: Replace All Uses With)至%x。 - 清理与重入:旧的
arith.addi指令被安全移出 AST 并析构;%x的所有新旧使用者被重新压入 Worklist 触发连锁化简。 - 单调性:Operation 节点总数净减 1,严格单调收敛,绝无反向膨胀可能。
请在阅读后先合上这一节,用自己的语言复述“输入状态 → 中间状态 → 输出状态”,再做练习。
实践任务:唯一代码填空题
补齐 x + 0 -> x 的替换动作。
规则:只能修改 TODO/______ 所在位置;不要删除断言或放宽误差。代码注释说明了每个边界条件。
Code
// C++ 片段:假设已确认 rhs 是整数 0。
LogicalResult matchAndRewrite(arith::AddIOp op,
PatternRewriter &rewriter) const override {
Value x = op.getLhs();
rewriter.______(op, ______);
return success();
}检查方法
在完整插件环境运行 pass 两次,第二次 IR 应不再变化;当前无构建环境时做 API/def-use 静态审查。
提交时请给出:补齐后的代码、实际运行输出(环境不可用时注明“仅静态审查”)以及对失败用例的解释。
Q1
不要背定义:请从输入、状态变化和输出三个阶段解释“Pass、PatternRewriter 与贪心重写”的工作机制。
你的答案:
Q2
直接调用 op.erase() 为什么可能破坏 rewrite driver?
你的答案:
Q3
两个都正确但互相逆向的 canonicalization 如何避免振荡?
你的答案:
评分与通过规则
- 代码 4 分:正常输入 2 分,边界输入 1 分,解释实现 1 分;
- Q1~Q3 各 2 分;
- 一票否决:结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。
需要提示时按四级机制请求:概念区域 → 具体方向 → 关键局部 → 完整答案。