深入浅出 MLIR 核心优化公共子表达式消除CSE的内在机理在编译优化的漫长历史中公共子表达式消除CSECommon Subexpression Elimination绝对算得上是最经典、最基础的全局优化手段之一。简单来说如果一段代码中多次出现了完全相同的计算并且这些计算所依赖的输入变量在期间没有发生任何改变编译器就会聪明地只计算一次然后把结果缓存起来供后续的代码复用。传统的编译器如经典的 LLVM IR通常是在扁平的基本块Basic Blocks或通过复杂的支配树Dominance Tree分析来实现全局的 CSE。然而在MLIR多级中间表示的世界里由于万物皆 Op、高度嵌套的 Region 结构以及独特的方言Dialect生态实现 CSE 的方式既极其优雅又充满挑战。为什么 MLIR 形式下的 CSE 如此高效在传统的 IR 中基础的算术运算如加、减通常是硬编码在核心 IR 里的原生指令。而 MLIR 严格贯彻了解耦哲学所有的计算甚至是基础的arith.addi或arith.mulf都是彼此平等的Operation操作。这就赋予了 MLIR 的 CSE 一种近乎降维打击的优势跨方言的通用性。MLIR 内部的cse优化 Pass 不需要关心你当前的 Op 是底层的寄存器加法arith.addi还是高级的张量算子如tosa.add甚至是你在自定义方言里写的一个专用矩阵乘法。只要这个 Op 满足了特定的静态规约MLIR 统一的基础设施就能以极高的效率自动识别并将其消除。MLIR 判定公共子表达式的硬性标准MLIR 是如何确定两个看似独立的 Operation 是可以被合并的“公共子表达式”呢它在底层依赖于以下三个维度的严格审查1. 结构等价性Structural Equality两个 Op 必须在结构上完全一比一对应。这意味着它们的Op 类型Name必须相同例如都是arith.mulf。它们的输入操作数Operands必须完全是相同的 SSA 值。它们携带的属性Attributes必须完全等价例如常量 Op 的数值必须一致。它们返回的结果类型Result Types必须完全对齐。2. 支配关系Dominance仅仅结构相同还不够计算发生的先后顺序也至关重要。由于 MLIR 严格遵循 SSA静态单赋值原则第一个 Op提供缓存结果的那个所在的 Block必须在控制流图上支配Dominate第二个 Op 所在的 Block。简单来说无论程序走哪条分支都必须先执行完第一个 Op才能走到第二个 Op。这样第二个 Op 才能安全地直接复用第一个 Op 产生的值。3. 无副作用Side-Effect Free这是最重要的安全底线。如果一个 Op 带有副作用例如改变了物理内存、向屏幕打印了日志、或者包含不确定性的随机数生成那么即使它和前面的某个 Op 长得一模一样也绝对不能被消除。因为消除它会直接改变程序的原有运行逻辑。代码进化论从冗余 IR 到纯净 IR为了让这个过程完全可视化我们来看一段包含冗余计算的真实 MLIR 片段优化前的原始 IR包含重复计算func.func redundant_math(%arg0: f32, %arg1: f32) - (f32, f32) { // 1. 第一次计算 %arg0 * %arg1 %res1 arith.mulf %arg0, %arg1 : f32 // 2. 插入一段不相干的加法计算 %tmp arith.addf %arg0, %arg0 : f32 // 3. 第二次计算完全相同的 %arg0 * %arg1输入没变Op没副作用 %res2 arith.mulf %arg0, %arg1 : f32 // 4. 将两次相乘的结果各自加上临时值 %out1 arith.addf %res1, %tmp : f32 %out2 arith.addf %res2, %tmp : f32 func.return %out1, %out2 : f32, f32 }当我们使用 MLIR 的核心工具链对这段代码运行mlir-opt --cse优化命令后整个中间表示会经历脱胎换骨的化简。优化后的纯净 IRfunc.func redundant_math(%arg0: f32, %arg1: f32) - (f32, f32) { // 1. 保留第一次计算它是公共子表达式的源头 %res1 arith.mulf %arg0, %arg1 : f32 %tmp arith.addf %arg0, %arg0 : f32 // [消亡] 原本的 %res2 消失了 // 2. 随后所有原本需要使用 %res2 的地方被完美重定向Rewritten到了 %res1 上 %out1 arith.addf %res1, %tmp : f32 %out2 ar {-- 核心数据流重定向 --} ith.addf %res1, %tmp : f32 func.return %out1, %out2 : f32, f32 }MLIR 在处理嵌套控制流Region时的黑科技如果只是处理像上面这样平铺直叙的代码传统的编译器后端也能做得很好。但 MLIR 的核心杀手锏在于它支持Region区域的嵌套结构比如scf.for内部包裹着循环体区域。MLIR 的 CSE 优化器拥有深度感知多层作用域的能力。---------------------------------------- | func.func (外层作用域) | | 计算了 %val a b | | | | ---------------------------------- | | | scf.for / scf.if (内层 Region) | | | | | | | | 再次遇到了 a b | | | | [CSE 触发]: 直接复用外层的 %val | | | ---------------------------------- | ----------------------------------------当 CSE Pass 深入扫描到一个嵌套的内层 Region 时它会主动带着外层作用域已经计算好的“已知表达式列表”一同走进去。如果发现在内层的循环或条件分支内部有人在重复计算外层已经算好的东西它会毫不犹豫地隔空切断数据流直接把内层的消费节点连接到外层的输出上。这种跨层级的全局透视在传统的扁平化 IR 中需要付出极其高昂的图分析代价。开发者启示让你的自定义方言完美支持 CSE如果你正在基于 MLIR 开发自己的硬件编译器或者特定的高级 AI 算子方言如何才能让你的自定义 Op 也能享受到框架自带的 CSE 免费红利呢答案非常简单只需要在用 TableGen.td文件定义你的 Operation 时为其挂上一枚特殊的勋章TraitNoMemoryEffect无内存副作用或实现MemoryEffectOpInterface内存副作用接口。def MyCustomAddOp : MyDialect_Opcustom_add, [NoMemoryEffect] { let summary A precision addition op tailored for custom accelerators; let arguments (ins F32:$lhs, F32:$rhs); let results (outs F32:$result); }一旦你声明了这个 TraitMLIR 的 CSE 引擎就会在编译期秒懂“哦这个算子不管执行多少次只要输入一样它就会老老实实吐出一样的结果而且不会乱动物理内存。” 随后框架便会自动接管一切帮你把代码中所有冗余的自定义计算清理得一干二净。