1. 编译器领域的乐高积木为什么MLIR值得你花时间第一次接触MLIR是在一个算子融合的项目里当时团队正被TVM的调度原语和手写CUDA之间的割裂感折磨得够呛。一个卷积算子前端框架导出计算图中间经过图优化最后落到硬件上还得靠人肉写模板整个链路像一条用胶带粘起来的流水线每换一个后端就要重新撕一遍。后来有人扔了个MLIR的链接到群里说这玩意儿可能能治我们的病。点进去一看满屏的方言、Pass、重定目标头都大了。但硬着头皮啃了两个月之后我确实回不去了——它把编译器从每个环节各写各的变成了大家说同一种语言按需翻译。这篇博文要聊的就是MLIR最核心的那套玩法可组合的模块化代码生成。说白了它解决的是一个老问题——当你需要为不同硬件、不同算子、不同精度组合生成高效代码时怎么避免每次都从零造轮子。MLIR给出的答案是把编译过程拆成一层层方言每层只管自己的事层与层之间通过标准化的中间表示衔接需要什么就组合什么。这就像乐高积木基础砖块是统一的但你可以拼出飞机、城堡或者机器人。适合读这篇的人如果你正在做AI编译器、算子库开发、或者任何需要把高层计算描述翻译成底层机器码的工作MLIR这套思路能帮你省下大量重复劳动。如果你只是听说过MLIR但没动过手这篇会带你从零跑通一个完整的张量编译器构造流程用的是MLIR 15的稳定版本所有代码都可以直接复现。如果你已经在用MLIR但觉得Pass写得乱七八糟那关于方言设计和Pass流水线的部分应该能给你一些整理思路的启发。我个人的背景是做了五年多的AI框架和编译器从TVM到XLA再到MLIR都踩过坑。这篇的内容基于我实际项目中的经验加上对MLIR官方文档和社区实践的梳理尽量做到看完就能动手。2. 拆解MLIR的核心设计方言、Pass与重定目标2.1 为什么需要方言这层抽象传统编译器比如LLVM它的IR是相对固定的——你拿到的是一个已经比较低层的表示适合做指令选择、寄存器分配这些后端优化。但AI编译器面对的问题不一样输入是TensorFlow/PyTorch的计算图里面全是conv2d、matmul、softmax这种高层算子输出可能是x86指令、CUDA PTX、或者某个NPU的私有指令集。中间这个跨度太大了如果只用一套IR要么高层信息丢失太早要么低层优化做不了。MLIR的方言机制就是来解决这个跨度的。你可以把方言理解成领域特定的IR扩展包。比如tosa方言专门描述张量算子linalg方言描述线性代数运算affine方言描述循环嵌套和内存访问llvm方言直接对应LLVM IR。每个方言有自己的操作定义、类型系统和验证规则。关键点在于这些方言可以在同一个模块里共存并且可以通过Pass互相转换。我刚开始学的时候有个误区以为方言是一层套一层的严格层次结构。实际上MLIR允许方言之间任意转换只要你能定义合法的转换规则。比如你可以从tosa直接降到linalg也可以从linalg降到scf结构化控制流再从scf降到cf控制流最后降到llvm。这种灵活性是MLIR最强大的地方也是最容易让人迷失的地方——因为选择太多了。2.2 可组合性的真正含义标题里的可组合模块化代码生成核心在于可组合三个字。在MLIR的语境下它意味着几件事第一方言可组合。你不需要一个方言包打天下。做张量计算就用linalg做循环优化就用affine做硬件特定优化就自己定义一个方言。每个方言只关心自己领域内的合法性和优化不越界。第二Pass可组合。MLIR的Pass是模块化的每个Pass做一件小事比如把linalg.matmul转成循环、把affine循环展开、把scf转成cf。你可以按任意顺序组合这些Pass形成一条流水线。官方提供的Pass有几百个你还可以自己写。第三目标可组合。同一个高层表示可以通过不同的Pass流水线降到不同的目标。比如同样的linalg算子走x86路线就降到llvm方言然后生成机器码走CUDA路线就降到nvvm方言然后生成PTX。这种一次编写多处编译的能力是MLIR相比传统编译器框架最大的优势。我实际项目里的做法是先定义一套自己的高层方言比如叫mydsl把业务逻辑用这套方言表示然后写几条降级路径分别对应不同的硬件后端。这样新增一个后端时只需要写从mydsl到目标方言的转换前面的优化Pass全部复用。2.3 重定目标的关键接口与模式重写重定目标这个词听起来很学术其实意思很简单让同一套编译器基础设施支持新的硬件目标。MLIR实现重定目标主要靠两个机制接口和模式重写。接口Interface是MLIR里一种类似契约的东西。比如LinalgOp接口定义了所有线性代数算子必须提供的方法比如获取输入输出、获取循环范围等。当你写一个通用的优化Pass时你面向接口编程而不是面向具体算子。这样新增一个算子时只要它实现了LinalgOp接口就能自动享受所有已有的优化。模式重写Pattern Rewrite是MLIR里做方言转换的标准方式。你定义一个模式描述什么样的操作可以替换成什么样的操作然后MLIR的贪婪重写引擎会自动在IR上应用这些模式。比如你可以定义一个模式linalg.matmul在满足某些条件时可以替换成linalg.generic加循环。这种声明式的写法比手写遍历IR树要清晰得多也不容易漏掉边界情况。我踩过的一个坑是模式重写的顺序很重要。如果两个模式都能匹配同一个操作重写引擎会按注册顺序尝试先匹配上的先应用。有时候一个过于激进的模式会阻止后面更优的模式生效。解决办法是用PatternBenefit给模式设置优先级或者用applyPatternsAndFoldGreedily的配置参数控制重写策略。3. 从零搭建一个张量编译器MLIR 15全流程复现3.1 环境准备与项目骨架先说一下环境。我用的MLIR 15.0.7Ubuntu 22.04LLVM/MLIR从源码编译。如果你不想编译也可以用预编译包但自己编译能确保版本一致后面调试也方便。编译命令大概是这样git clone https://github.com/llvm/llvm-project.git cd llvm-project git checkout llvmorg-15.0.7 mkdir build cd build cmake -G Ninja ../llvm \ -DLLVM_ENABLE_PROJECTSmlir \ -DLLVM_BUILD_EXAMPLESON \ -DLLVM_TARGETS_TO_BUILDX86;NVPTX \ -DCMAKE_BUILD_TYPERelease \ -DLLVM_ENABLE_ASSERTIONSON ninja编译大概要一两个小时取决于机器性能。编译完之后build/bin下面会有mlir-opt、mlir-translate这些工具后面会频繁用到。项目骨架我建议这样组织mytensor-compiler/ ├── CMakeLists.txt ├── include/ │ └── MyDialect/ │ ├── MyDialect.td │ ├── MyOps.td │ └── Passes.td ├── lib/ │ ├── MyDialect.cpp │ ├── MyOps.cpp │ └── Passes.cpp ├── tools/ │ └── my-opt.cpp └── test/ └── simple.mlir这个结构是MLIR官方推荐的out-of-tree项目布局方便你把自己的方言和Pass跟MLIR核心分开管理。TableGen文件.td用来声明方言、操作和PassC文件实现具体逻辑。3.2 定义自己的张量方言假设我们要做一个简单的张量编译器支持matmul、add、relu三个算子。先定义方言// MyDialect.td def MyDialect : Dialect { let name my; let cppNamespace ::my; let summary A toy tensor dialect for learning MLIR; let description [{ This dialect contains basic tensor operations for demonstrating MLIRs composable code generation workflow. }]; }然后定义操作。以matmul为例// MyOps.td def My_MatMulOp : My_Opmatmul, [Pure] { let summary Matrix multiplication; let arguments (ins AnyTensor:$lhs, AnyTensor:$rhs); let results (outs AnyTensor:$result); let assemblyFormat $lhs , $rhs attr-dict : type($lhs) , type($rhs) - type($result); }这里Pure表示这个操作没有副作用可以被优化器自由移动或删除。assemblyFormat定义了文本IR的格式这样你在.mlir文件里可以这样写%0 my.matmul %a, %b : tensor4x8xf32, tensor8x16xf32 - tensor4x16xf32定义完方言和操作后需要写C实现来注册它们。核心是MyDialect.cpp里的initialize()方法调用addOperations把操作注册进去。这一步看起来简单但有个细节操作的验证逻辑。MLIR允许你在操作定义里加hasVerifier然后实现一个verify()方法检查输入输出类型是否匹配。比如matmul要求两个输入的维度满足矩阵乘法规则这个检查最好在验证阶段做而不是等到降级时才报错。3.3 降级路径设计从高层方言到Linalg定义完自己的方言后下一步是把它降级到MLIR已有的方言。我选择降到linalg因为linalg是MLIR里做张量计算降级最成熟的中间层后面可以继续降到scf、affine、llvm。降级用模式重写实现。以my.matmul为例struct MatMulToLinalg : public OpRewritePatternmy::MatMulOp { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(my::MatMulOp op, PatternRewriter rewriter) const override { auto lhsType op.lhs().getType().castRankedTensorType(); auto rhsType op.rhs().getType().castRankedTensorType(); auto resultType op.result().getType().castRankedTensorType(); // 创建 linalg.matmul auto linalgOp rewriter.createlinalg::MatmulOp( op.getLoc(), ValueRange{op.lhs(), op.rhs()}, ValueRange{op.result()}); rewriter.replaceOp(op, linalgOp.getResults()); return success(); } };这个模式做的事情很简单把my.matmul替换成linalg.matmul。但实际项目中你可能需要处理更复杂的情况比如输入是动态shape、或者需要插入tensor.cast来对齐类型。我建议一开始就考虑动态shape的支持因为AI模型里动态batch size太常见了如果等到后面再加改动会很大。my.add和my.relu的降级类似分别对应linalg.add和linalg.elemwise_unary用relu的region实现。这里有个技巧linalg的逐元素操作可以用linalg.generic统一表示但用具体的linalg.add更直观也更容易被后续Pass识别。3.4 构建Pass流水线有了降级模式接下来要组织Pass流水线。我的流水线大概是这样void buildPipeline(OpPassManager pm) { // 第一阶段高层方言内部优化 pm.addPass(my::createConstantFoldingPass()); pm.addPass(my::createOperatorFusionPass()); // 第二阶段降级到linalg pm.addPass(my::createConvertMyToLinalgPass()); // 第三阶段linalg层优化 pm.addPass(mlir::createLinalgElementwiseOpFusionPass()); pm.addPass(mlir::createLinalgBufferizePass()); // 第四阶段降级到scf/affine pm.addPass(mlir::createConvertLinalgToLoopsPass()); pm.addPass(mlir::createConvertSCFToCFPass()); // 第五阶段降级到llvm pm.addPass(mlir::createConvertToLLVMPass()); }这个流水线里前两个Pass是我自己写的后面都是MLIR官方提供的。createConvertMyToLinalgPass内部会应用前面定义的模式重写。createLinalgBufferizePass把tensor类型转成memref类型这是从值语义转向内存语义的关键一步。createConvertLinalgToLoopsPass把linalg算子展开成循环嵌套。最后createConvertToLLVMPass把整个IR降到LLVM方言之后就可以用mlir-translate生成LLVM IR再交给LLVM后端生成机器码。实际跑的时候我发现createLinalgBufferizePass需要配置BufferizationOptions指定哪些函数需要bufferize、内存空间怎么分配。默认配置在简单例子上能跑但复杂模型上会出问题。我的经验是尽早引入bufferize并且把内存分配策略显式化。比如用memref.alloc显式分配而不是依赖自动推导这样后面做内存复用优化时更容易控制。3.5 完整复现一个矩阵乘法加ReLU的例子说了这么多跑一个完整例子最直观。假设输入是一个简单的计算C relu(A * B)其中A是4x8B是8x16都是f32。首先写高层IRfunc.func matmul_relu(%a: tensor4x8xf32, %b: tensor8x16xf32) - tensor4x16xf32 { %0 my.matmul %a, %b : tensor4x8xf32, tensor8x16xf32 - tensor4x16xf32 %1 my.relu %0 : tensor4x16xf32 - tensor4x16xf32 return %1 : tensor4x16xf32 }用my-opt跑流水线./build/bin/my-opt --my-pipeline test/simple.mlir -o output.mlir跑完之后output.mlir里应该能看到llvm方言的操作。如果中间某一步报错可以用--mlir-print-ir-after-all打印每一步之后的IR方便定位问题。我实际跑的时候遇到过一个错误linalg.matmul要求输入是memref类型但我的IR里还是tensor类型。原因是createLinalgBufferizePass没有正确配置。解决办法是在bufferize之前先跑createConvertTensorToMemrefPass或者直接在bufferize配置里指定tensor到memref的转换规则。另一个常见问题是动态shape。如果A的维度是?x8linalg.matmul需要知道具体的循环边界。MLIR 15里可以用tensor.dim操作获取动态维度然后在降级时生成相应的边界检查。这部分比较复杂建议先跑通静态shape再逐步加入动态支持。4. 实战中踩过的坑与排查技巧4.1 方言转换中的类型不匹配这是最常见的问题。比如my.matmul的输出类型是tensor4x16xf32但linalg.matmul期望的输出是memref4x16xf32。如果直接替换MLIR的验证器会报类型错误。解决办法有两种一是在模式重写里插入tensor.cast或memref.cast做类型转换二是调整流水线顺序先做bufferize再做方言转换。我倾向于第二种因为bufferize是一个全局性的决策放在流水线早期更可控。排查这类问题的技巧用mlir-opt --verify-each让每个Pass之后都验证IR合法性这样能精确定位到哪个Pass产生了非法IR。另外--mlir-print-op-generic可以打印操作的通用形式方便看类型信息。4.2 Pass顺序导致的优化失效MLIR的Pass顺序对最终代码质量影响很大。我遇到过一个问题先跑linalg的融合Pass再跑降级Pass结果融合没生效。原因是降级Pass把linalg算子转成了循环融合Pass就认不出来了。正确的顺序应该是先做高层优化再做降级。具体来说linalg层的融合、tiling、向量化都应该在convert-linalg-to-loops之前完成。一旦降到scf或affine优化空间就小很多了。另一个经验是不要一次性把所有Pass都加上。我刚开始图省事把官方推荐的Pass全塞进流水线结果编译时间暴涨而且有些Pass互相干扰。后来改成增量式先跑最小流水线确认能生成正确代码再逐个加优化Pass每加一个就测一下性能和正确性。4.3 内存分配与生命周期管理从tensor降到memref之后内存分配就成了必须显式处理的问题。MLIR提供了memref.alloc和memref.dealloc但手动管理容易漏掉dealloc导致内存泄漏。我的做法是用bufferization的自动分析功能让它推导出内存生命周期然后插入alloc/dealloc。MLIR 15的OneShotBufferize可以处理大部分情况但对于有控制流的函数可能需要手动干预。一个实用的技巧在bufferize之后跑--memref-optimize它会做一些内存复用和提升。但要注意这个Pass可能会改变内存布局如果后面有依赖特定布局的硬件指令需要谨慎使用。4.4 常见问题速查表问题现象可能原因排查方法解决方案验证器报类型错误方言转换时类型不匹配--verify-each定位到具体Pass插入cast或调整流水线顺序优化Pass不生效Pass顺序错误--mlir-print-ir-after-all对比IR把优化Pass移到降级Pass之前编译时间过长Pass过多或IR过大用--mlir-timing统计各Pass耗时精简流水线增量添加生成代码性能差缺少tiling/向量化对比有无优化Pass的IR加入linalg的tiling和vectorization动态shape报错循环边界未处理检查tensor.dim的使用显式生成边界检查或特化4.5 调试工具与技巧MLIR自带了不少调试工具我常用的有--mlir-print-ir-after-all打印每个Pass之后的IR最常用的调试手段。--mlir-print-ir-before-all打印每个Pass之前的IR配合上面的一起用。--mlir-timing统计各Pass耗时找性能瓶颈。--mlir-disable-threading单线程跑方便调试多线程相关的问题。--debug-onlydialect-conversion打印方言转换的详细日志。另外我强烈建议在项目里加单元测试。MLIR的lit测试框架很好用写一个.mlir文件用RUN行指定命令和期望输出就能自动化验证。比如// RUN: my-opt --my-pipeline %s | FileCheck %s // CHECK-LABEL: func matmul_relu // CHECK: llvm.call这样每次改Pass之后跑一遍测试能快速发现回归。5. 从玩具到生产MLIR工程化的几点体会5.1 方言设计的粒度控制自己定义方言时最容易犯的错是什么都往里塞。我见过一个项目把卷积、池化、归一化、激活全塞进一个方言结果这个方言变得无比庞大每个操作都要写一堆验证和降级逻辑。我的建议是按抽象层次划分方言。高层方言只描述做什么比如my.conv2d中层方言描述怎么做比如linalg.conv_2d低层方言描述用什么指令做比如llvm或硬件特定方言。每个方言只关心自己层次的事跨层转换用模式重写。另一个经验是尽量复用MLIR已有的方言。linalg、affine、scf、arith这些方言已经非常成熟没必要重新造。你的自定义方言应该只包含那些MLIR没有的、业务特定的操作。5.2 Pass的可测试性与可维护性Pass写多了之后维护是个大问题。我的做法是第一每个Pass只做一件事。比如把my.matmul转成linalg.matmul是一个Pass把my.add转成linalg.add是另一个Pass。这样测试和调试都简单。第二Pass的配置参数化。比如tiling的块大小、向量化的宽度都做成Pass选项而不是硬编码。这样同一套Pass可以适配不同的硬件。第三用PassPipeline注册流水线。MLIR允许你把一组Pass注册成一个命名的流水线比如--my-pipeline。这样用户不需要知道内部有哪些Pass只需要调用流水线名字。也方便你做A/B测试对比不同流水线的效果。5.3 性能调优的实际经验最后聊一下性能。MLIR生成的代码质量很大程度上取决于你用了哪些优化Pass。我实测下来对矩阵乘法这类计算密集型算子以下几个Pass效果最明显linalg-tile把大矩阵切成小块提高缓存命中率。块大小需要根据目标硬件的缓存大小调我一般从32x32开始试。linalg-vectorize把标量循环转成向量指令。在x86上配合AVX2/AVX-512效果很好。affine-loop-unroll循环展开减少分支开销。展开因子一般设2或4。canonicalize和cse清理冗余代码几乎每个阶段之后都应该跑。但要注意这些优化不是越多越好。我试过把所有优化都打开结果编译时间从几秒涨到几分钟而性能只提升了不到5%。后来改成按需开启先跑基准测试找到瓶颈再针对性地加优化Pass。还有一个容易忽略的点数据布局。MLIR默认的memref布局是行优先但某些硬件可能更适合列优先或分块布局。如果目标硬件有特殊要求需要在bufferize阶段就指定布局而不是等到最后再转换。5.4 后续扩展方向这套框架跑通之后可以往几个方向扩展。一是加入更多算子比如卷积、池化、归一化把常见CNN和Transformer的算子覆盖全。二是支持更多后端比如通过nvvm方言生成CUDA代码或者通过自定义方言对接特定加速器。三是加入自动调优用MLIR的transform方言描述优化策略然后搜索最优的tiling和向量化参数。我目前在做的是把transform方言集成进来用脚本描述优化流程而不是硬编码Pass顺序。这样调优的时候只需要改脚本不用重新编译C代码。这个方向MLIR社区也在推应该是未来的主流做法。踩了这么多坑最大的体会是MLIR的学习曲线确实陡但一旦理解了方言Pass模式重写这套组合拳后面的事情就顺了。它不像TVM那样给你一个端到端的黑盒而是给你一套积木让你自己搭。自由度高了责任也大了——你需要自己想清楚每一层做什么、怎么衔接。但正是这种自由度让MLIR能适配从手机到数据中心的各类硬件也让它在AI编译器领域越来越受欢迎。