声明式图优化的威力深入理解 MLIR 中的 DRR Pattern声明式重写规则在上一篇文章中我们见证了MLIRGen是如何把前端的语法树或计算图一步步引渡进 MLIR 的多级方言世界。当初始的高层方言 IR如StableHLO或TOSA生成完毕后编译器接下来最核心的任务就是图级优化Graph Optimization——比如算子融合、代数化简如x×1→xx \times 1 \to xx×1→x以及消除冗余操作。在传统的编译器设计中编写这些图优化 Pass 是一件极其痛苦的体力活。你需要用 C 手写复杂的、嵌套多层的if-else状态机去探测图的拓扑结构。为了彻底解放编译器工程师MLIR 引入了一套惊艳的元编程图重写基础设施——DRRDeclarative Rewrite Rules声明式重写规则。本文将带你深度拆解 DRR Pattern 的设计哲学、核心语法与工作机制看它如何用区区几行声明取代上百行繁琐的 C 模式匹配。1. 什么是 DRR PatternDRRDeclarative Rewrite Rules是 MLIR 体系下基于TableGen语法的声明式图重写规范。它的核心理念极其纯粹用一种直观的、高度接近数学公式的声明语法定义“源图结构Source Pattern”与“目标图结构Target Pattern”的映射关系。[ 开发者编写的优化规则 .td ] │ ▼ (编译期) [ mlir-tblgen ] (带 -gen-rewriters 后端) │ ▼ (自动喷涌) [ 高效、高内联的 C 模式匹配类 ] (自动处理拓扑遍历与类型校验)通过 DRR你不再需要关心如何用 C 去一步步检查一个算子的输入是不是另一个算子也不用关心如何手动呼叫OpRewriter去删除或替换节点。这一切通通由 TableGen 在编译期自动帮你生成。2. DRR 核心概念Pat的黄金模板在.td文件中定义一个 DRR 规则的核心是使用官方提供的Pat类。其基本语法结构如下def : Pat(SourcePattern), (TargetPattern), [Constraints], (Benefit);SourcePattern源模式描述你想在 IR 图中捕捉的、带有冗余或能被优化的算子结构。TargetPattern目标模式描述你想把源结构替换成的新算子结构。Constraints约束列表可选用来执行额外的动态检查例如要求某个张量的维度必须大于 0或者数据类型必须是 FP32。Benefit优化收益可选一个整数用来告诉优化驱动器PatternApplicator这条规则的优先级。数值越高编译器越优先匹配它。3. 经典案例拆解大模型推理中的常见 DRR 优化为了让你切实感受到 DRR 的强大我们来看三个大模型推理与算子融合中非常高频的真实场景。① 代数化简消除乘以 1 的冗余x×1.0→xx \times 1.0 \to xx×1.0→x在模型量化或多图拼接后经常会产生一些由于算子退化而导致的乘以 1 的操作。// 引入 StableHLO 算子定义 include stablehlo/dialect/StableHloOps.td // 匹配模式 stablehlo.multiply( $input, (stablehlo.constant 1.0) ) // 重写模式 直接返回 $input def SimplifyMultiplyOne : Pat (StableHLO_MulOp $input, (StableHLO_ConstantOp $attr)), (string_to_op return $input;) // 伪代码逻辑指示直接返回该输入值 // 实际工程中通常配合内联约束如判断 $attr 的值是否为 1 ;② 算子融合Op Fusion重塑激活函数TanhMul→\to→GELU逼近在大模型如 BERT, GPT的 MLP 层中经常会遇到特定的算子组合我们可以直接将其融合成一个芯片后端支持得更好的原子高性能算子。// 假设自研 NPU 方言里有一个高度优化的高性能 FusedGeluOp def FuseGeluPattern : Pat (MyNPU_MulOp $input, (MyNPU_TanhOp (MyNPU_ScaleOp $input, $coeff))), (MyNPU_FusedGeluOp $input, $coeff) ;通过上面这段简单的多重嵌套表达DRR 编译器会自动在底层生成一套深度优先搜索DFS状态机精准在图里捕捉到先Scale再Tanh最后与原输入Mul的复杂拓扑结构并完成一击替换。③ 动态约束带条件的类型安全重写有时候只有当算子满足特定条件时优化才能安全进行。此时我们需要引入HasStaticShape或自定义的 C 约束。// 定义一个自定义的 C 约束检查两个操作数是否具有相同的元素类型 def HasSameElementType : ConstraintCPred$0.getType().getContext() // 实际为 C 校验表达式; def SafeOpFusion : Pat (MyNPU_AddOp $arg0, $arg1), (MyNPU_FastAddOp $arg0, $arg1), [(HasSameElementType $arg0, $arg1)] // 只有满足该约束优化才会发生 ;4. DRR 深度进阶如何表达多输出与高级变换NativeCodeCall当你深入大模型的张量切分Tensor Parallelism或者图层面的Split/Concat优化时你会发现有些高层算子的返回结果不止一个多输出或者重写时需要进行复杂的数学计算。DRR 为此提供了极其强悍的扩展手段① 绑定算子结果Binding Result如果一个源算子产出多个结果你可以用$符号给特定输出槽位命名// 假设某个 TopK 算子输出两个值Values 和 Indices def : Pat (MyNPU_TopKOp:$results $input, $k), (MyNPU_NewTopKOp $input, $k) // 替换整个算子自动对齐两个输出槽位 ;② 逃生舱口NativeCodeCall当你在目标重写模式中需要计算一些极为繁琐的动态属性比如计算融合后的卷积新的 Padding 大小而这又是声明式语法无法表达的DRR 允许你直接呼叫纯手写的 C 函数// 声明一个 C 辅助函数用来动态计算融合后的 Stride 属性 def ComputeFusedStride : NativeCodeCallcomputeStride($0, $1); def MergeStrides : Pat (MyNPU_ConvOp (MyNPU_PoolOp $in, $stride1), $stride2), (MyNPU_FusedConvOp $in, (ComputeFusedStride $stride1, $stride2)) // 在重写发生的瞬间TableGen 自动调用 C 的 computeStride 函数注入属性 ;5. DRR Pattern 的运行时激活在.td文件里优雅地写完这些规则后如何让它们真正滚动在编译器的 Pass 链路中呢TableGen 会为你生成一个名为GeneratedConvert.inc的文件。你只需要在一个标准的 C 图重写 Passmlir::RewritePatternSet中用一行宏控代码将它们全部收拢voidpopulateGraphRewritePatterns(mlir::RewritePatternSetpatterns){// 这一行宏会瞬间将 .td 里定义的几十个声明式优化规则实例化并塞入 Pass 集合#includeMyNPU/TransformPatterns.inc}随后将这个patterns集合丢给 MLIR 官方的贪心图重写驱动器applyPatternsGreedily。驱动器就会以极致的效率在你的计算图上反复横扫Iterative Scan直到没有任何一个节点能被 DRR 规则匹配到为止。总结一句话概括如果说 C 图重写OpRewritePattern是用笨重的乐高积木拼雕塑那么 DRR Pattern 就是直接用 3D 打印机的设计图纸批量产出高性能优化通道。在大模型编译和专用芯片编译器的高频更迭时代DRR 用声明式的单源真理把极其容易出错的图拓扑校验逻辑彻底托管给元编译器。掌握 DRR 的模式表达艺术是让 AI 编译器能够在短短几行代码之间将顶层数学期望完美转化为底层硬件暴击的极致生产力密码。