设想一家连锁餐饮集团:过去每个门店都有自己的“切菜哲学”,A 店把胡萝卜切丝,B 店切丁,C 店整根丢下锅。结果中央厨房的炒菜师傅每天都要面对三种完全不同的萝卜,只能一边炒一边吐槽,出菜速度和口味都像开盲盒。后来总部设了一个“标准切配间”:所有食材到店先按统一规格切配,胡萝卜一律切成 3 厘米小段,不允许自由发挥。看似只是多了一步整理,整个厨房的节奏却立刻稳定下来。
MLIR 的编译器世界也面临同样的混乱。来自 PyTorch、TensorFlow、手写 DSL 的模型被翻译成中间表示(IR)后,同一个计算常常有五花八门的写法:有人把 x*1 原样留着,有人把 1+2 放在运行期才算,还有人产生一个结果却根本没人用。如果后续优化 Pass 总是面对这些“百店百味”的输入,优化效果会变得难以预测。Canonicalization(规范化)Pass 就是那个“标准切配间”:它不负责让代码变快,而是把所有语义等价的操作统一成标准形式,为后续优化提供干净一致的原料,让整条编译流水线的行为又快又稳。
理解 Canonicalization,是读懂 MLIR 优化流水线的第一把钥匙。掌握了它,你才能写出可预测、可调试、可复用的编译器 Pass,这是 AI 编译器工程师的基本功。
目录
可视化演示
展示MLIR规范化Pass:模式匹配重写规则将等价IR操作统一为规范形式,算术简化分类(常量折叠/代数恒等/强度削减),编译流水线中操作数量缩减。
本节配有交互动效(进入视口后自动加载,约 3KB)。
模式匹配与重写规则 等价IR操作(x+0/x*1/x*2)匹配规范化模式 重写为规范形式(x/x/xCanonical(x+0) = x Pattern Matcher将等价变体统一为唯一表示
概念初识:什么是规范化
连锁餐饮的后厨总监发现,只要把“切法”统一,整个炒菜线就能稳定运作。编译器也是一样:MLIR 是 AI 编译器的“中央厨房”,各种前端框架生成的中间表示(IR)就是送到这里的食材。IR 是编译过程中表示程序结构的一种中间语言,Pass 则是这个厨房里的一道道工序。前端千奇百怪的写法让 IR 充满了“同义反复”,如果不先整理,后续优化就无从谈起。
Canonicalization 的定位正是这个“标准切配间”。它接收任意写法的操作图,通过规则库把语义等价的不同形式统一成一种标准形式。它不关注“哪段代码更快”,只关注“这段代码应该长成什么样”。例如 mul(x, 1) 被化简为 x,add(1, 2) 被折叠成 3,没有用户引用的节点被直接移除。整个过程保证语义不变,只改变表现形式。
要做到这一点,规范化实际上是在解决一个“选代表”的问题:在无数个语义等价的表达式中,挑一个结构代价最小的作为标准代表。
回到厨房的比喻,切配间并不决定这道菜好不好吃,它只把胡萝卜、土豆、洋葱都切成同样的小段。正因为所有原料先被“标准化”,后续的炒制、调味、摆盘才能像流水线一样稳定。Canonicalization 的价值不在于自己输出什么惊艳的优化,而在于让后面的 Pass 站在同一条起跑线上,得到可预测的结果。
- 规范化操作符
C(expr)表示对输入表达式expr执行规范化,目标是返回一个标准形式。 - 等价候选集
E收集了所有与expr语义等价的表达式,e是其中的候选者。 - 语义约束
s.t.表示“在……约束下”,这里要求候选形式与原始表达式计算结果完全一致,保证改写不会改变程序行为。 - 最小代价选择
argmin在所有合格候选中挑出cost(e)最小的那个,这个唯一的“最小代价代表”就是规范形式。
核心机制:模式匹配与 DAG 重写
切配间里有一台“配菜机器人”:它盯着流水线上每一份食材,对照标准手册找茬。看到一块不符合规格的胡萝卜,就按住它、按手册重新切一遍。切完后继续扫下一个,扫完一轮再回到起点重新扫,直到整条流水线上再也找不到一粒不规范的食材。这种“不达标就改,改完再查”的循环,就是 Canonicalization 的工作方式。
在 MLIR 中,这台机器叫作 DAG 重写器。它把 IR 看成一幅有向无环图,节点是操作,边是数据依赖。它按一种贪心策略遍历图上的每个节点,把当前节点依次提交给一组合法注册的“重写规则”(Rewrite Pattern),任一条规则命中,就立即原位替换。替换完成之后,新节点可能会创造新的匹配机会,所以重写器要反复扫描,直到任何规则都无法再匹配为止。
每条重写规则都是一种“从左往右”的化学置换。
配菜机器人看似只会“一根筋”地扫描,但正是这种机械的、反复的扫描保证了最终结果唯一。只要规则库设计得好,任何输入 IR 都会收敛到同一个标准形式,这就是规范化的“不动点”概念:改到不能再改,就是终点。
- 规则的左端
L描述待匹配的子图模式,比如mul(%x, %c),%x是任意值,%c是某个常量。 - 匹配条件
cond(L)对左端的模式提出额外要求,例如“%c必须等于 1”,这样规则就只对x*1生效。 - 规则的右端
R是命中后的替换结果,它和L语义等价,但更接近标准形式。
关键操作:切配间里的四把刀
标准切配间每天都做四件事:能算的提前算,没用的直接扔,花哨的改朴素,标签统一贴。这四件事对应 Canonicalization 中最常见的四类操作:常量折叠、死代码消除、结构简化和属性归一。
常量折叠是最直观的刀法。一个操作的所有输入如果都是常量,编译器可以直接算出结果,用一个常量节点替换整个子图。比如 add(1, 2) 直接被替换成 3,避免运行期白白计算。死代码消除则是“垃圾分类”:如果一个操作的结果没有任何其他操作引用,它就是在空转,完全可以安全删除,给后面的 Pass 减轻负担。
结构简化负责处理各种“花架子”写法。x + 0 化简为 x,x * 1 化简为 x,x - x 化简为 0。这些恒等式看似微不足道,却在图形神经网络、Transformer 等模型的 IR 中反复出现,积少成多就能大幅降低后续 Pass 的压力。属性归一则负责把等价但编码不同的属性统一起来,比如把来自不同方言的常量属性整理成同一种表示,让后续的匹配和判定逻辑不必为“两种写法”各写一份代码。
折叠操作有一个明确的触发条件。
这四把刀合在一起,正是规范化“删繁就简”的核心动作。回到厨房:能提前准备的食材提前切好打包,过期变质的直接丢掉,摆盘雕花全部换成朴素装盘,包装标签统一格式。切配间的每一次出手,都在让后厨变得更整齐、更高效。
- 全称遍历
∀ o ∈ operands(op)表示遍历操作的每一个操作数,一个都不放过。 - 常量判定
isConst(o)判断该操作数是否是编译器能静态求值的常量。 - 折叠触发当所有操作数都满足常量判定,
fold(op)就会执行求值,并把整个操作替换成常量结果。
编写方法:如何注册一条规范化规则
一位新主厨推出一道招牌菜,想把它纳入全集团的标准化流程。他不能口头说两句就完事,必须把食材规格、切配步骤、摆盘标准写进《后厨标准手册》,并经过收益评估。只有确定这道菜的新标准不会让后厨更乱,才会被正式收录。MLIR 中写规范化规则也是同样的道理。
编写一条模式的核心是定义 matchAndRewrite 两个阶段:匹配阶段负责检查当前节点是否符合某个模式,重写阶段负责把节点替换成新形式。下面是一段把整数乘法乘以 1 化简为 x 的典型规则:
写完规则后,还需要在该方言操作的 getCanonicalizationPatterns 中把它注册进 Canonicalizer,否则 DAG 重写器根本不知道这条规则存在。注册后的规则被放入规范模式库,在每次 --canonicalize 运行时被调用。
但并不是所有化简都值得收录,规则库需要一条收益准则来筛选。收益准则可以写成:
标准化手册不是越厚越好,收录的每一条菜谱都必须经得起推敲。MLIR 编译器社区正是靠着这套“收益可量化、规则可复用”的工程规范,让 Canonicalization 的模式库在几十年间不断沉淀,既丰富又克制。
- 规模变化
|size(L) − size(R)|度量规则应用前后 IR 的节点数变化,差值越大说明化简效果越明显。 - 风险惩罚
spec(r)表示这条规则可能引入的额外风险,比如触发更多轮重写、增加编译时间,λ是惩罚系数。 - 收益判定只有当总收益大于经验阈值
θ时,规则才值得写入规范模式库,否则只会增加重写器的负担。
struct SimplifyMulOne : public OpRewritePattern<arith::MulIOp> { LogicalResult matchAndRewrite(arith::MulIOp op, PatternRewriter &rewriter) const override { if (!matchPattern(op.getRhs(), m_One())) return failure(); rewriter.replaceOp(op, op.getLhs()); return success(); }};
流水线整合:让优化工序步调一致
标准化切配间不是孤岛。总部把切配工序放在整个后厨流水线的入口,要求所有食材先切配再烹饪;等炒完一道菜、需要复配时,还会再次回到切配间整理。Canonicalization 在 MLIR 流水线中的位置也是如此:它是各优化阶段之间的“常规净化器”,从头到尾反复出现。
一个典型流程是:某个优化 Pass 生成了一批新的 IR 结构,里面可能带着冗余的 add(x, 0) 或是重复的常量计算。此时跑一次 Canonicalize,就能把结构重新整理干净,再交给下一阶段。因此,Canonicalize 经常和“公共子表达式消除”(CSE)在循环中交替执行,直到 IR 不再变化。这种“优化—清理—再优化”的模式保证了每个优化 Pass 都面对标准输入。
Canonicalization 还能给流水线带来极强的可预测性:同样的输入 IR,无论在哪一次编译、哪一台机器上运行,规范化后都得到完全相同的标准形式。版本迭代时,新增一个优化 Pass 可能会引入新的结构,但 Canonicalize 会把它拉回统一的风格,保证老测试数据依然有效,回归对比变得可信。
整个流水线的形态可以用一张组合图描述:
切配间虽不负责炒菜,但整条流水线都依赖它维持节奏。Canonicalization 正是用这种“润物细无声”的反复整理,让 MLIR 流水线中的每一个 Pass 都能在稳定、统一的地基上构建,也让“优化效果可预测”从口号变成了工程现实。
- 前置清场
Canonicalize在优化 Pass 之前先跑一遍,清掉冗余和异形操作。 - 反复净化
Optimize执行后可能产生新的非标准结构,所以要再次Canonicalize。 - 稳定条件当第
k轮与第k+1轮的规范化结果完全相同时,说明已经收敛到不动点,可以进入后续降级阶段。
实战对策:调试、收敛与调优
某天切配间出了问题:机器人把胡萝卜切成了丁,又把丁切回整根,再切丁,永无止境。后厨经理立刻调出监控、对照手册、检查前后工序的产物,最后发现手册里写了两条矛盾的规则。调试 Canonicalization 也是一样的思路:打开日志、观察 IR、找出对抗模式。
平时调试编译器,我们最常用的是 mlir-opt 工具,它可以把单个 Pass 单独拿出来跑。对一段小 IR 执行 mlir-opt --canonicalize,就能在最小范围内观察规范化做了什么。配合 --print-ir-before/after,可以看到每个 Pass 执行前后的完整 IR 变化,比在完整流水线里大海捞针高效得多。
最棘手的实战问题是“不收敛”。如果规则之间存在对抗,比如规则 A 把 x*1 改写为 x,规则 B 又把 x 改写为 x*1,DAG 重写器就会在这两种形态之间反复横跳,永远到不了不动点。遇到这种情况,要优先审查规则的条件约束是否过弱,必要时直接移除收益为负的对抗规则。
判定收敛的依据很简单:连续两轮规范化后的 IR 是否完全相同。
切配间的故障排查经验告诉我们:规范化失控不是玄学,而是规则工程问题。在工程实践中,还可以用 -view-op-graph 可视化 IR 图、用 --debug 输出重写日志,甚至通过 Timer 统计每个 Pass 的耗时。掌握这些调试手段,你就能从“编译器黑盒用户”变成“编译器外科医生”。
- 状态跟踪
IR_k是第 k 轮规范化后的 IR 状态;我们逐轮观察它是否变化。 - 收敛判定若相邻两轮结构完全一致,说明重写器再也找不到可化简的子图,达到不动点。
- 循环检测如果某个历史状态再次出现,比如
IR_k = IR_m,则规则之间存在对抗,必须人工干预。
AI 应用场景
场景一:AI 编译器流水线的“标准入口”
阶段:模型优化与编译
解决的问题:从 PyTorch/TensorFlow 导出的模型转成 MLIR 后,算子写法千差万别,后续的算子融合、内核自动调优无法稳定生效。传统做法是手写各种“兜底逻辑”去兼容不同写法,维护成本极高。接入 Canonicalization 后,所有输入先被规范成统一形式,融合 Pass 只需面向一种 IR 形态编程。
具体案例:某团队把 HuggingFace 的 BERT 模型导出到 MLIR,先用 --canonicalize 清理冗余的 reshape 和 transpose,再做 --aggressive-fusion 算子融合。结果显示,融合规则命中率显著提升,生成的 GPU kernel 性能可复现,不再随输入前端的不同而波动。
场景二:芯片后端代码生成的前置清理
阶段:编译优化与代码生成
解决的问题:针对华为昇腾、NVIDIA GPU 等硬件生成代码时,后端 Pass 对 IR 形态有严格要求。Canonicalization 可以承载硬件厂商自定义的“本地化规则”,把通用标准形式进一步整理成硬件友好的形态,消除后端 Pass 的异常分支。
具体案例:在 GPU 后端中,某些 broadcast 操作会被规范化拆成更小的切片,或者把 truncf 与 add 的顺序重排,使得向量化 Pass 能生成更宽的内存访问指令。整条后端的代码生成路径因此变得更加稳定,代码体积也明显下降。
场景三:MLIR 方言开发中的规则调试
阶段:自定义 DSL 编译与方言落地
解决的问题:公司内部发明了一种新的张量方言,开发团队需要为它编写 Canonicalization 规则。如果不懂得“模式匹配 + 不动点收敛”的原理,很容易写出互相冲突的规则,导致编译卡死。
具体案例:团队为 tensor.pack 操作编写 fold 规则时,先用 mlir-opt --canonicalize --print-ir-after 观察每条规则的独立效果,再逐步合并。通过小步迭代,他们在一周内建立了一套稳定的规范化规则库,新方言在流水线中的表现从此可预测。
Python 演示
class Node: def __init__(self, op, args): self.op = op self.args = tuple(args) def __repr__(self): if self.op == 'const': return str(self.args[0]) if self.op == 'var': return str(self.args[0]) return self.op + '(' + ','.join(repr(a) for a in self.args) + ')' def __eq__(self, other): return (isinstance(other, Node) and self.op == other.op and self.args == other.args) def __hash__(self): return hash((self.op, tuple(repr(a) for a in self.args))) def mk_const(v): return Node('const', [v]) def mk_var(name): return Node('var', [name]) def is_const(node): return node.op == 'const' def fold_once(node): if node.op in ('const', 'var'): return node args = [fold_once(a) for a in node.args] node = Node(node.op, args) if len(node.args) == 2: a, b = node.args if node.op == 'add': if is_const(a) and a.args[0] == 0: return b if is_const(b) and b.args[0] == 0: return a if is_const(a) and is_const(b): return mk_const(a.args[0] + b.args[0]) if node.op == 'mul': if is_const(a) and a.args[0] == 1: return b if is_const(b) and b.args[0] == 1: return a if is_const(a) and is_const(b): return mk_const(a.args[0] * b.args[0]) if node.op == 'sub' and a == b: return mk_const(0) return node def canon(node): while True: new = fold_once(node) if repr(new) == repr(node): return new node = new def main(): expr = Node('add', [ Node('mul', [mk_var('x'), mk_const(1)]), Node('add', [mk_const(1), mk_const(2)]) ]) print('before:', expr) norm = canon(expr) print('after: ', norm) if __name__ == '__main__': main()
(x*1)+(1+2) 规范化为 x+3 的过程。def canon(node):这是规范化主循环,对应 MLIR 中 DAG 重写器反复扫描 IR 直到不再变化的行为。new = fold_once(node)每一轮只应用一层规则,先把子节点递归化简,再尝试化简当前节点。if repr(new) == repr(node):用结构字符串判断是否达到不动点;等价于编译器比较前后两轮 IR 是否完全一致。if node.op == 'add':分派到加法运算规则组,类似 MLIR 方言中为每个 Operation 注册一组 Pattern。if is_const(b) and b.args[0] == 0:命中"x + 0 → x"的恒等式规则,直接返回左侧操作数。if is_const(a) and is_const(b):两个操作数都是常量,触发常量折叠,把表达式替换为计算结果。expr = Node('add', [构造一个带冗余的初始表达式,模拟来自前端的不同写法。
常见误区
误区“Canonicalization 是个优化 Pass,把它跑一遍,模型推理速度应该立刻变快。”
正解规范化关注的是“形式统一”,不是“性能提升”。它确实会顺手化简 x*1、折叠常量等看似优化的操作,但真正的性能收益来自后续的算子融合、向量化、内存规划等 Pass。规范化的价值在于让这些 Pass 拥有简洁一致的输入,从而更高效、更可预测,而不是直接产出高性能代码。
为何常见:因为规范化的产物常常表现为“代码变短、节点变少”,看起来就像在做优化,所以很容易被误解为性能优化 Pass。
误区“把 Canonicalization 安排在流水线里跑一次,IR 就已经是规范形式了,后面不用再管。”
正解规范化是一个迭代到不动点的过程。一条规则改写完某个节点后,可能产生新的可匹配结构,必须继续扫描和应用规则。而且在真实流水线中,后续的优化 Pass 会不断制造新的冗余,所以 Canonicalize 需要被反复安排在多个阶段之间执行。
为何常见:初学时把 Pass 想象成“顺序执行一次就结束”的普通程序,忽略了 DAG 重写器自带的贪心循环特性。
误区“规范化会挑选全局最小的表达式,节点越少就说明规范化越成功。”
正解规范形式是“为了后续 Pass 的便利而约定的标准形式”,不一定是节点数最少的极简形式。某些情况下,为了保持结构规整、便于模式匹配,编译器甚至会保留一些稍显冗余的写法。
为何常见:直观上“优化”总是朝着更小更快的方向,于是把“标准化”错误理解为“最小化”。
误区“只要我多写几条化简规则,规范化效果就越好,编译器就越智能。”
正解规则越多,DAG 重写器的匹配开销越大;更危险的是,规则之间可能互相触发或对抗,导致重写不收敛。MLIR 社区对每条规则都要做收益评估,只有化简收益大于风险代价的规则才应被收录。
为何常见:把规则库误认为“知识库”,以为知识越多越好,忽略了编译器 Pass 的工程约束和收敛风险。
误区“规范化无论如何都保证语义不变,所以它对最终生成的代码没有任何影响。”
正解规范化确实不改变语义,但它决定了后续优化 Pass 看到的结构。不同的规范形式会让后续 Pass 走上截然不同的优化路径,最终生成的代码质量和性能可能差异很大。如果规范化行为不稳定,同一份模型在不同版本编译器上的表现就会不可预测。
为何常见:只看到“语义不变”这层底线,忽略了编译优化链条中每步决策都依赖于输入这个事实。
面试题
Q1. 规范化 Canonicalization 的核心目标是什么?
- A.将语义等价的不同形式统一为标准形式← 正确答案
- B.将深度学习模型转换为 ONNX 格式
- C.将 MLIR 源码编译成可执行文件
- D.将训练过程分布式部署到集群
答案:A。规范化像是厨师的“标准切配间”,它的目标是把语义等价但写法不同的操作统一成标准形式。B 是模型格式转换,C 是编译器后端工作,D 是分布式系统问题,都与规范化的核心目标无关。
Q2. 以下哪个命令可以单独运行 Canonicalization Pass?
- A.使用 --canonicalize 单独运行规范化← 正确答案
- B.使用 --cse 进行公共子表达式消除
- C.使用 --lower-to-llvm 做后端降级
- D.使用 --convert-func-to-llvm 转换函数
答案:A。mlir-opt --canonicalize 是单独执行规范化的标准方式,就像把配菜机器人单独开起来测试。其他选项分别对应 CSE、降级和转换,是不同的工序。
Q3. DAG 重写器中的规则应用会持续到什么时候?
- A.达到固定迭代次数后自动停止
- B.没有规则能匹配任何子图时到达不动点← 正确答案
- C.直到所有操作全部变为常量为止
- D.直到当前函数的节点数变为零
答案:B。规范化是贪心地反复应用规则,直到没有任何规则能继续匹配子图,这个状态就是不动点。A 不是标准的收敛条件;C 和 D 是极端假设,不符合真实 IR 的多样化结构。
Q4. 将 add(1, 2) 直接替换为常量 3 属于哪种规范化操作?
- A.死代码消除
- B.结构简化
- C.常量折叠← 正确答案
- D.属性归一
答案:C。所有操作数都是常量,编译器提前算出结果并用常量替换,这正是“能算的提前算”的常量折叠。死代码消除针对无引用的节点,结构简化处理恒等式,属性归一统一属性编码,都不符合题意。
Q5. 关于规范化,以下哪个说法是正确的?
- A.所有 IR 节点数一定会减少
- B.规范化是一种性能优化手段
- C.输入是语义等价的不同写法← 正确答案
- D.规范化只会执行一次
答案:C。规范化的典型输入是同一个语义对应的多种等价形式,把这些写法统一为标准形式是它的使命。A 过于绝对,规范化不一定让节点减少;B 混淆了标准化与性能优化;D 忽略了迭代到不动点的机制。
Q6. 在 MLIR 中注册规范化模式通常使用哪个入口?
- A.getCanonicalizationPatterns← 正确答案
- B.getOptimizationPatterns
- C.getRewritingPatterns
- D.getConversionPatterns
答案:A。MLIR 方言操作通过 getCanonicalizationPatterns 把规则交给 Canonicalizer。这就像把菜谱提交给切配间手册,而不是交给炒菜区或包装区。其他三个名称分别是通用优化、通用重写和类型转换的入口,不是 Canonicalization 的专用通道。
Q7. 如果规则 A 把 x*1 改写成 x,规则 B 又把 x 改写成 x*1,会发生什么?
- A.编译器会自动忽略冲突并继续
- B.规范化会直接报错并终止
- C.重写可能振荡导致无法收敛← 正确答案
- D.重写器会自动删除这两条规则
答案:C。这就是规则对抗的“死循环”,像机器人把胡萝卜切丁又切回整根。重写器会在两种形式间反复横跳,永远到不了不动点。编译器通常不会自动识别这种对抗,需要人工审查并移除规则。
Q8. 以下哪些操作属于 Canonicalization 的典型行为?
- A.删除无用户引用的死代码← 正确答案
- B.将常量表达式直接折叠成结果← 正确答案
- C.将存储布局改写成设备专用格式
- D.把等价结构统一成标准形式← 正确答案
答案:ABD。死代码消除、常量折叠、统一等价结构都是规范化的标准动作。C 是面向特定硬件的数据布局优化,属于后端优化 Pass 的职责,而不是通用的 Canonicalization 行为。
Q9. 关于规范化在编译流水线中的角色,以下说法正确的是?
- A.通常被安排在优化 Pass 之前执行← 正确答案
- B.能让优化结果更可预测← 正确答案
- C.可以取代所有硬件相关优化
- D.在流水线中可能多次出现← 正确答案
答案:ABD。规范化常作为前置清洗出现,也常与 CSE 等 Pass 循环配合,让优化结果稳定可预测。C 夸大了它的作用:硬件优化依赖架构相关知识,不是通用规范化能替代的。
Q10. 规范化可能导致编译问题的原因包括哪些?
- A.规则过多会放大重写开销← 正确答案
- B.规则对抗会导致不收敛← 正确答案
- C.规范选择会影响后续优化决策← 正确答案
- D.规范形式一定能提升算子性能
答案:ABC。规则过多会拖慢重写器,规则对抗会导致振荡,而规范形式的选择又会决定后续优化的路径,这些都是真实工程中需要警惕的。D 是典型误解:规范形式只是标准形式,并不代表性能最优,甚至可能为了让结构统一而保留非最优的写法。