MLIR可组合模块化代码生成：从零构建张量编译器实战

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.7，Ubuntu 22.04，LLVM/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_PROJECTS=mlir \ -DLLVM_BUILD_EXAMPLES=ON \ -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" \ -DCMAKE_BUILD_TYPE=Release \ -DLLVM_ENABLE_ASSERTIONS=ON 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）用来声明方言、操作和Pass，C++文件实现具体逻辑。

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 MLIR's composable code generation workflow. }]; }

然后定义操作。以matmul为例：

// MyOps.td def My_MatMulOp : My_Op<"matmul", [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 : tensor<4x8xf32>, tensor<8x16xf32> -> tensor<4x16xf32>

定义完方言和操作后，需要写C++实现来注册它们。核心是MyDialect.cpp里的initialize()方法，调用addOperations把操作注册进去。这一步看起来简单，但有个细节：操作的验证逻辑。MLIR允许你在操作定义里加hasVerifier，然后实现一个verify()方法检查输入输出类型是否匹配。比如matmul要求两个输入的维度满足矩阵乘法规则，这个检查最好在验证阶段做，而不是等到降级时才报错。

3.3 降级路径设计：从高层方言到Linalg

定义完自己的方言后，下一步是把它降级到MLIR已有的方言。我选择降到linalg，因为linalg是MLIR里做张量计算降级最成熟的中间层，后面可以继续降到scf、affine、llvm。

降级用模式重写实现。以my.matmul为例：

struct MatMulToLinalg : public OpRewritePattern<my::MatMulOp> { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(my::MatMulOp op, PatternRewriter &rewriter) const override { auto lhsType = op.lhs().getType().cast<RankedTensorType>(); auto rhsType = op.rhs().getType().cast<RankedTensorType>(); auto resultType = op.result().getType().cast<RankedTensorType>(); // 创建 linalg.matmul auto linalgOp = rewriter.create<linalg::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是4x8，B是8x16，都是f32。

首先写高层IR：

func.func @matmul_relu(%a: tensor<4x8xf32>, %b: tensor<8x16xf32>) -> tensor<4x16xf32> { %0 = my.matmul %a, %b : tensor<4x8xf32>, tensor<8x16xf32> -> tensor<4x16xf32> %1 = my.relu %0 : tensor<4x16xf32> -> tensor<4x16xf32> return %1 : tensor<4x16xf32> }

用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的维度是?x8，linalg.matmul需要知道具体的循环边界。MLIR 15里可以用tensor.dim操作获取动态维度，然后在降级时生成相应的边界检查。这部分比较复杂，建议先跑通静态shape，再逐步加入动态支持。

4. 实战中踩过的坑与排查技巧

4.1 方言转换中的类型不匹配

这是最常见的问题。比如my.matmul的输出类型是tensor<4x16xf32>，但linalg.matmul期望的输出是memref<4x16xf32>。如果直接替换，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 常见问题速查表

4.5 调试工具与技巧

MLIR自带了不少调试工具，我常用的有：

• --mlir-print-ir-after-all：打印每个Pass之后的IR，最常用的调试手段。
• --mlir-print-ir-before-all：打印每个Pass之前的IR，配合上面的一起用。
• --mlir-timing：统计各Pass耗时，找性能瓶颈。
• --mlir-disable-threading：单线程跑，方便调试多线程相关的问题。
• --debug-only=dialect-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编译器领域越来越受欢迎。