深入浅出 MLIR 核心优化规范化Canonicalize的代数简化艺术在 MLIR 的整个编译流水线Pass Pipeline中如果说CSE公共子表达式消除扮演的是“精简冗余”的清洁工那么Canonicalize规范化则是整个体系中最高频、最核心的“整容算术师”。在真实的编译器运行流程中各个高级方言在降级Lowering或图变换时往往会生成大量符合语法、但从数学或逻辑角度来看极其“愚蠢”的代码。例如将一个变量加上 0、乘以 1或者是连续做两次取反。Canonicalize 的核心使命就是将这些形态各异、但语义等价的代码统一收敛并改写为一种最简单、最标准、最易于后续处理的“规范形态Canonical Form”。规范化背后的核心哲学传统的编译器如经典的 LLVM IR通常把这类优化拆散在不同的 Pass 中比如死代码消除DCE、常数折叠Constant Folding和指令组合Instruction Combine。而 MLIR 将这些概念统一打包进了canonicalize这个超级 Pass 里面。它融合了两个最关键的优化技术常数折叠Constant Folding如果一个 Op 的所有输入在编译期就已经确定是常量规范化框架直接在编译期把结果算出来用一个arith.constant替换掉原来的计算 Op。模式匹配改写Pattern Rewriting通过预先定义好的数学规则将一种 Op 组合直接替换为另一种更优的 Op 组合。规范化的三大高频应用场景为了直观感受到 Canonicalize 的威力我们来看几个工业级编译中最常见、也最典型的代码改写案例场景一代数恒等式消除Algebraic Identities在代码生成或者算子展开时经常会引入大量的无效数学边界。规范化能够利用完美的代数定律将其抹平。原始代码某个计算过程产生了%res arith.addi %val, %c0变量加 0。规范化后整个加法 Op 直接消亡后续所有使用%res的地方被一气呵成地重定向到初始变量%val上。相同的逻辑同样适用于乘 1、异或 0、位移 0。场景二冗余操作抵消Cancellation当某些具备自反性质的操作连续作用在同一个变量上时规范化可以一眼看穿本质并将其成对消除。原始代码对一个布尔值连续做两次按位异或/取反或者高层图结构中对一个张量连续做两次完全对称的转置Transpose。规范化后两次操作全部被抹去直接还原最底层的数据引用。场景三类型/维度的精简还原Shape Minimization在处理高维张量Tensor时由于切片Slice和拼接Concat操作经常会产生大量的动态维度。原始代码从一个张量中动态切出一个子块但实际上切片的起始坐标和步长在编译期通过常数推导发现正好等于原张量的完整边界。规范化后这个原本需要耗费运行时计算的多维切片 Op如tensor.extract_slice会被直接判定为无效整体降级并替换为对原张量的直接引用。代码进化论看 Canonicalize 如何横扫冗余 IR我们来看一段在 AI 编译器中非常具有代表性的原始中间表示IR。这段代码在语法上完全合法但充满了低效的计算优化前的原始 IRfunc.func messy_pipeline(%arg0: i32) - (i32, i32) { // 1. 编译期完全已知的两个常数加法可触发常数折叠 %c10 arith.constant 10 : i32 %c32 arith.constant 32 : i32 %folded_val arith.addi %c10, %c32 : i32 // 2. 毫无意义的数学恒等式加0 %c0 arith.constant 0 : i32 %clean_val arith.addi %arg0, %c0 : i32 // 3. 乘法与加法的经典组合可能触发强度削减如乘以2变成自身相加 %c2 arith.constant 2 : i32 %scaled_val arith.muli %clean_val, %c2 : i32 func.return %folded_val, %scaled_val : i32, i32 }当我们对这段中间表示施加mlir-opt --canonicalize命令后整个代码流会被瞬间净化。优化后的纯净 IRfunc.func messy_pipeline(%arg0: i32) - (i32, i32) { // [脱胎换骨] 10 32 在编译期被直接算好精简为一条孤立的常量定义 %c42 arith.constant 42 : i32 // [消亡] 加0的操作被连根拔起%clean_val 的语义直接收敛到 %arg0 身上 // [强度削减] 乘以 2 被底层的规范化重写规则自动等价替换为了耗时更短、更亲和硬件的加法指令 %scaled_val arith.addi %arg0, %arg0 : i32 func.return %c42, %scaled_val : i32, i32 }规范化的底层驱动核心TableGen 与 ODS 模式匹配为什么 MLIR 的 Canonicalize 可以拥有如此强大的跨方言吞吐能力这得益于它和 MLIR 的表格驱动TableGen声明式重写规则DRR的完美绑定。当你在设计一个自定义方言时如果你想让某个 Op 具备规范化能力通常有两种高效的玩法玩法一使用 ODS 在.td文件中写声明式规则你可以直接在 TableGen 文件中声明重写模式Pattern不需要写任何繁琐的 C 代码// 声明一个将“乘以2”重写为“自身相加”的数学规范化规则 def MultiplyByTwoToHasAdd : Pat (MyDialect_MulOp $input, (ConstantOpConstantAttr:$attr, 2)), (MyDialect_AddOp $input, $input) ;MLIR 编译器在构建时会自动把这些声明式的拓扑关系编译成极其高效的图遍历和替换代码并自动打包进canonicalizePass 里面。玩法二在 C 中重写getCanonicalizationPatterns接口对于一些牵扯到动态维度或者复杂条件判断的极其复杂的规范化逻辑你可以直接在 Op 的定义中挂载 C 模式类voidMyCustomOp::getCanonicalizationPatterns(RewritePatternSetresults,MLIRContext*context){// 将你手写的、继承自 OpRewritePattern 的高级重写逻辑灌入结果集中results.addMyComplexOptimizationPattern(context);}总结在实际的生产环境部署中canonicalize通常不是孤立运行的。它像一个无处不在的润滑剂通常会以固定的节奏穿插在每一个核心转换 Pass 的前后例如Pass Pipeline: Optimize - Canonicalize - Lowering - Canonicalize。通过这种高频的、局部的代数和拓扑化简它能够不断扫清中间产物的代码迷雾确保后面的编译器后端看到的一直是最干净、最具表达力、也最容易转换为物理机器码的中间表示。