1. 从四处碰壁到真正提速我为什么决定把推理链路整个交给MLIR干推理优化这几年有一个问题几乎每次都会被问到模型部署时的性能瓶颈到底在哪如果你的第一反应是“算子实现不够快”那只能说答对了一小半。我自己的感受是绝大多数性能损失发生在算子与算子之间的缝隙里——中间张量来来去去搬内存调度器在一层层图结构里做局部决策底层硬件能力大量闲置。那时候我手头一个线上项目ONNX导出的卷积网络C侧直接用推理引擎逐算子执行单batch延迟一直压不下去profiling结果里Conv、ReLU、BN各算各的光中间结果写回和重新读入就占掉了整个推理时间将近30%。痛定思痛之后我开始认真看MLIR这条路。最开始就是被它那句“多层级中间表示”的宣传语吸引但真正用起来才发现MLIR模型编译加速的本质是提供了一个把整个模型当成“程序”来编译的框架——把计算图从高层的算子序列逐层降低成循环、向量、内存操作最终落地成机器码而不是在运行时逐算子解释执行。这篇文章我不打算做概念搬运工。我尽量把这条链路拆开讲清楚MLIR到底用怎样的架构设计支撑了模型编译加速我自己是怎么搭出一条从ONNX到动态库的完整编译流水线的以及在实际调优过程中哪些优化手段真正带来了肉眼可见的收益哪些坑让我白白熬了好几个通宵。如果你正在做推理加速或者对TorchInductor、ONNX Runtime的底层思路感兴趣又或者你想为自研芯片写一套编译器工具链这篇文章应该能帮你在MLIR这块少走些弯路。2. 搞懂MLIR的底层逻辑dialect、pass和多级IR是怎么协同工作的2.1 dialect不是玄学它就是一组定义好的算子和规则刚接触MLIR的人十有八九先被术语糊了一脸。一张IR文件里可能同时出现tosa.conv2d、linalg.generic、arith.addf、memref.alloc看得人头皮发麻。实际上dialect翻译成“方言”非常形象——每种方言都有自己的算子和语义规范就像不同地区的语言有自己的词汇和语法一样。关键点在于所有方言都运行在同一个MLIR基础设施之上。一个IR可以同时包含多种方言的算子因为它们本质上都是Operation有统一的属性、类型、区域region结构。这套设计让“跨层级优化”成为可能你在高层还在讲“卷积后面跟着ReLU”到了低层可以讲“这个循环体内先算乘加再算激活”两者都存在于同一棵IR树上只是所处高度不同。我在实际项目中用的最多的一组搭配是输入模型经导入器变成stablehlo或者tosa方言这层算子语义丰富、跟ONNX比较像适合做图级别的优化接下来降到linalg开始显式表达多维循环结构再往下经过scf/affine方言处理循环控制流最后落到llvm方言交给LLVM后端生成机器码。2.2 pass的本质是“一遍遍重写IR”顺序比数量更重要有了dialect定义算子下一步就是怎么把它变快。MLIR里面优化动作被封装成pass一个pass就是一次对整个IR的扫描和重写。--convert-tosa-to-linalg就是把tosa算子改写成等价的linalg算子--canonicalize负责清理冗余操作。这里我想多说一句顺序问题。pass之间的顺序依赖是我见过最多人踩坑的地方。比较典型的两个例子第一你想做算子融合前提是被融合的几个算子已经被降级到同一个dialect层级否则融合pattern根本匹配不上——就好比你让两个说不同语言的人合作起码得先给他们配个翻译第二内存规划必须在算子基本定型之后做做得太早后面一通融合优化把buffer结构全打乱了前面算的就白费了。我自己有一个工作习惯搭流水线时先把每个pass的输入输出IR快照打出来跑完一轮看一眼中间结果确认每个pass都按预期改写出了目标结构。这比到最后发现性能不对再回头排查要省时间得多。顺便说一句mlir-opt给--mlir-print-ir-after-all这个选项能把每个pass之后的IR全部打印出来虽然输出很长但配合grep过滤关注的关键算子名定位问题非常高效。2.3 多级IR的真正价值保留每层最该保留的信息“为什么需要多级IR一个中间表示用到底不行吗”这个问题我当年也纠结过。答案其实很实际没有任何一个单一的IR能同时完美表达“这里是卷积ReLU融合”这种高层语义又能精细控制“向量寄存器怎么分配”这种底层细节。高层IR的优势是语义完整做图级优化省力低层IR的优势是贴近硬件做代码生成和调度精准。MLIR选择的方法论很简单——不找一个“万金油”而是提供一整套阶梯式IR让编译器开发者在不同层级做不同的事每层优化都基于该层最丰富的信息展开。这个设计带来的实际收益我在做算子融合时感受特别明显。如果一开始就把模型降到纯循环层想识别“这是ConvReLUBN三个算子的组合”就非常困难你得从循环体结构里反推语义但如果从tosa层入手模式匹配一目了然。反过来如果一直留在高层循环分块、向量化这些硬件相关的优化步骤就无从下手。MLIR把这两件事拆开每个阶段都做得很专注。3. 手搭一条完整的MLIR加速流水线从ONNX模型到动态链接库理论聊得差不多了直接看看怎么落地。我把自己的那条链路完整走一遍——输入是一个ONNX格式的卷积模型输出是一个可被C动态加载的.so文件。3.1 链路总览一次性讲清五个关键阶段整个流程可以分成五个阶段导入阶段把ONNX模型转成MLIR的stablehlo或tosa方言IR由导入器负责算子映射和张量类型转换。高层优化在图级别做算子融合、常量折叠、死代码消除——这时候语义信息最全改起来成本最低。中低层转换把高层IR降低到linalg和scf显式表达循环和访存结构为后续优化做准备。底层优化循环分块tiling、向量化、内存规划这些动作直接决定生成代码的硬件利用效率。代码生成从linalg一路降到llvm方言再交给LLVM后端生成目标平台机器码封装成动态库。你不需要从零实现全部环节。社区已经有相当成熟的工具链LLVM官方仓库的MLIR子项目自带大量转换passONNX-MLIR项目则提供了完整的ONNX导入和编译工具。自己真正要动脑子的部分是根据目标硬件和模型特点把流水线中的pass组合和顺序调整到最优。3.2 一个能跑的流程pass组合与关键命令我在做实验时流程通常长这样。首先用ONNX-MLIR把模型导入onnx-mi --import-onnx model.onnx -o model.mlir拿到初始IR后做高层到中层的转换和优化。这里要注意具体用哪些pass取决于你导入到的dialect是哪一种我用stablehlo或者tosa的路径在pass名称上略有差异mlir-opt model.mlir \ --convert-tosa-to-linalg \ --canonicalize \ --convert-linalg-to-loops \ -o model_linalg.mlir接下来是底层优化阶段。这一段的pass选择和参数跟目标硬件强相关我实验平台上常见的组合是mlir-opt model_linalg.mlir \ --linalg-tile \ --vectorize \ --convert-vector-to-llvm \ -o model_opt.mlir最后生成LLVM IR并编译成动态库mlir-translate model_opt.mlir --mlir-to-llvmir -o model.ll clang -shared model.ll -o libmodel.so这里有一个我特别想强调的坑不同版本的LLVMpass名称和可用性差异很大。我试过照着网上的旧教程抄命令结果mlir-opt直接报pass不存在。最稳妥的做法是跑一下mlir-opt --help用手头版本的pass列表对照调整。3.3 运行时加载动态库的调用接口怎么设计才不踩坑生成的libmodel.so本质上是机器码加载方式跟普通动态库一致。我一般会让编译侧给每个模型自动生成一段C接口wrapper输入输出统一用一种buffer描述符结构传递而不是为每个模型手工硬编码一堆参数。void* handle dlopen(./libmodel.so, RTLD_LAZY); auto infer (int (*)(void* in, void* out))dlsym(handle, model_infer); infer(inputBuffer, outputBuffer);真正的增量价值在于模型定义和运行时代码彻底解耦了。后续模型迭代只需要重新编译生成新的.so文件替换即可调用方代码一行都不用改。这个设计在模型版本频繁变动的项目里能省掉大量联调时间。4. 真正带来性能收益的三个关键优化手段链路跑通只能算是热身。接下来这部分是三板斧式的核心优化手段——算子融合、循环分块与向量化、常量折叠与内存规划。每一项背后都有明确的性能逻辑不是玄学。4.1 算子融合把Conv、ReLU、BN变成一段连续代码先看最经典的例子。卷积 - ReLU - 批归一化这种组合几乎每个CNN里都有。逐算子执行时Conv的输出要完整写进内存ReLU再把这块数据读回来算完再写一次BN又读一遍。一次推理里这种“写-读-写-读”循环反复发生访存开销高得吓人。算子融合的思路很简单让编译器意识到这三个算子可以合并成一个融合算子中间结果尽量留在寄存器或片上缓存不要反复穿越内存层级。具体到MLIR实现我通过conversion pattern来做定义一个模式匹配规则当匹配到“Conv的后继是ReLUReLU的后继是BN”这样的子图时把三者替换成一个自定的融合op。后续代码生成阶段为这个融合op生成一个单一循环体每个数据点上连续执行三次计算中间不落地。但我也要泼一盆冷水不要逢op就融。盲目把整张图合成一大坨表面上减少了中间张量实际上可能因为单次循环体过大导致寄存器溢出性能反而下降。我现在的习惯是先跑profiling找出真正耗时的热点组合精准打击。4.2 循环分块与向量化把CPU的SIMD单元喂饱现代CPU和GPU都有向量处理单元一条指令能同时处理多个数据。想让生成代码的性能上台阶就得让循环体内的计算尽量替换成向量指令。MLIR在linalg和scf层级提供了对循环结构的完全控制。我一般先做tiling——把大循环切分成小块让内层循环的数据量和向量寄存器的宽度匹配再通过向量化pass把内层标量计算替换成向量操作。tile size的选择是门手艺活。我这边踩过的教训是内层tile大小要按硬件的SIMD宽度定比如AVX-512机器算float数据内层取16个元素外层tile则慢慢试目标是把中间数据留在L1/L2 cache里。分块之后还要关注访存顺序局部性不好分了也白分。4.3 常量折叠与内存规划看似不起眼性价比极高常量折叠是整个编译器领域最经典的优化之一。推理场景下模型权重是固定的很多关于固定张量的计算可以在编译期预先算好。MLIR的canonicalize就能搞定这件事清掉“乘1加0”之类的冗余运算后IR节点数会明显下降后续模式匹配的速度和成功率都会提升。内存规划排在第二序位。所有中间buffer可以在同一个主内存池上复用前提是生命周期分析确认两个buffer不会同时活跃。MLIR在memref层做这个分析比较自然把每个缓冲区按依赖关系求存活区间然后做区间分配。对初学者我的建议是先做保守复用别一上来搞复杂的线性扫描或图着色分配。保守策略同样能拿到可观收益复杂度却低一个量级不容易引入内存冲突问题。5. 实战调试IR打印、pass回退和动态shape那些坑MLIR调试和普通程序调试完全是两个世界。没有断点没有watch窗口你手里全是一摞IR文本。掌握几个排错方法能省掉大量无效时间——这条路我是实打实走到烦才总结出来的。5.1 IR前后对比最朴素也最高效的定位手段MLIR的绝大多数转换都是“IR进IR出”。这意味着任何一次不预期的改写都能通过比对IR快照定位。我在每个pass跑完之后保存一份快照文件名带序号遇到生成代码不对或性能异常就顺着快照找哪一步开始出现异常。实际操作用mlir-opt --mlir-print-ir-after-all就可以让每轮pass后打印IR。虽然输出很长但配合grep过滤目标op名基本能确认是哪一轮pass改出了问题。5.2 索引映射linalg优雅背后的“暗礁”linalg算子用indexing maps描述“输入到输出的坐标映射关系”。这是个优雅的设计但新人十有八九会栽在这里。indexing maps写得不严谨轻则无法融合重则导致代码生成阶段循环展开错误。我自己有一次调融合IR里tosa.conv2d降完应该变成等价的linalg.generic结果输出loop的边界全是?。查了一圈就是indexing map里dimset少写了一个维度导致shape丢失。这个教训让我后来给任何转换pattern都加了一条硬性要求先打印IR跑一遍确认输出的linalg.generic的loop bounds、iterator types、indexing maps全部正确再说下一步。5.3 动态shape一切优化pass的隐形杀手动态shape是MLIR实战中最常见、也最烦人的问题。ONNX模型的batch维度经常被标成动态导入MLIR后一堆op带动态维度而很多优化pass对动态shape直接“透明”处理——匹配规则看到动态shape就放弃导致后续tiling、向量化统统失效性能表现和预期差出一大截。一个可行的处理策略编译时固定一个最大支持batch比如4对外接口做padding和mask换来大量可用的优化机会。这个trade-off在服务端推理场景下通常很划算因为线上请求的batch size其实相对稳定。推理引擎常见的做法本质上就是这么回事。5.4 调试小工具让Python bindings帮你管理pass流水线命令行拼pass串太容易出错了。我后来把常用的流水线封装成一个Python函数输入mlir文件路径内部构造pass manager跑完每个pass后自动把IR写到目录下。这样整个编译过程完全可视化、可回放。from mlir import ir from mlir.passmanager import PassManager with ir.Context() as ctx: module ir.Module.parse(model_linalg.mlir) pm PassManager.parse( builtin.module(canonicalize,convert-linalg-to-loops) ) pm.run(module.operation) print(module)这个思路特别适合做优化实验改一个pass参数重跑一遍看IR的变化——所有决策都基于可见的IR事实而不是凭空猜。6. 从能跑到能用工程化阶段必须做的几件事编译链路跑通后真正的挑战在于怎么把它变成稳定、可维护、可迭代的生产能力。这块的经验总结起来是几个字缓存中间产物、解耦版本沟通、用数据驱动优化决策。6.1 IR缓存与增量编译别让模型迭代吃光你的开发时间模型编译是昂贵的操作。一份ONNX可能只是改了最后一个卷积层的out_channel但每次都要完整重跑一遍编译流水线。我的做法是用输入模型的特征做hash缓存输入模型的hash加编译选项的hash作为key命中就直接返回之前生成的动态库。这个改动在项目初期看不出大收益但到了模型高频迭代阶段就非常关键——省掉的时间可以做更多轮优化实验。缓存粒度我做在了IR层级而不是整个库的层级这样即使某几个pass改了参数也能复用前面没受影响的IR中间结果。6.2 编译期与运行期的“缓冲区接口契约”动态库生成后输入输出buffer的布局信息要同步生成一份描述文件格式输出input个数、shape、dtype、bias等。运行期解析这份描述文件做内存准备。我见过不少团队把输入输出结构硬编码在调用方代码里模型版本一换就崩。更稳妥的做法是让描述文件成为编译期和运行期之间的唯一契约调用方只依赖这份描述不依赖任何硬编码的模型细节。这样模型下线、换版本都是天衣无缝的事。6.3 用benchmark数据反向指导优化别信感觉信数据优化做得对不对最终衡量标准只有一个——运行时的实测数据。我建议任何一次改动都跑同一套benchmark脚本同一份输入数据同样的运行环境统计延迟、吞吐和内存占用。我做过一次“笨”但效果极好的实验把每个优化pass单独开启/关闭跑一遍用全组合实验寻找最优pass配置。MLIR的pass大多可以独立开关这让组合实验变得异常容易。做完一轮沉淀出的“编译配方”可以直接复用到后续其他模型上省掉大量重复试错的时间。注意无论你看到网上的某篇优化经验写得多么肯定都要在自己的硬件和模型上做A/B验证。访存模式、指令集特性、数据尺寸都会改变优化策略的结论没有放之四海而皆准的优化组合。7. 最后说点操作层面的经验MLIR的学习曲线确实陡这点我不否认。我最早啃它的时候光是理顺tosa、linalg、scf几个dialect之间的关系就花了不少时间。但真正用起来之后我的体会是这个框架最强大的地方反而不是某一招具体的优化技术而是它把“编译器工程”变成了一套可组合、可调试、可复用的积木。现在的我面对任何一个新模型推理优化需求第一反应都是“先跑出IR快照看看瓶颈在哪”。第二反应是“这个模型里有没有可以融合的连续访存热点”。然后才是“要不要做tilingSIMD宽度是多少”。整个思维方式和以前看算子库、调图形调度器完全不同。还有个小技巧说说。因为做tiling参数实验很频繁我把tile size的调整也做成了脚本自动化用配置文件驱动。一组tile size跑到基准测试输出延迟然后自动换一组再跑。最后模型交付时用的tile参数全是最优解——而不是我拍脑袋定的值。如果你正准备开始MLIR实战我建议你记住这句话先搭一条最小链路跑通一个简单模型然后盯着IR和性能数据去理解每一步发生了什么。pass不生效、IR降不下去这些坑你迟早会踩但大多数都能通过多看看IR文本、多跑几轮参数实验来解决。这个领域就是这样慢工出细活一旦建立起感觉后面效率高到你自己都惊讶。