MLIR 这套东西我最早是在做算子融合的时候被同事安利的。当时我们的推理引擎要同时支持 CPU、GPU 和一颗自研 NPU三套后端各写各的 IR 和 Pass改一个算子要动三个地方维护成本高得离谱。后来把中间表示统一到 MLIR 上用一套方言描述计算再分别 lower 到不同目标工作量直接砍掉一大半。这篇就围绕 MLIR 的可组合模块化代码生成思路把结构化、可重定目标的张量编译器构造法讲透并附上 MLIR 15 的全流程复现步骤。不管你是刚接触编译器的新手还是已经在做算子编译的老手都能从里面拿到能直接上手的东西。1. 为什么张量编译器需要 MLIR 这层中间表示1.1 传统张量编译器的三座大山先说清楚问题在哪。一个张量编译器本质上要干三件事把高层计算图翻译成中间表示、在中间表示上做优化、再把优化后的表示翻译成目标硬件的机器码或库调用。听起来简单但真做起来麻烦全在细节里。第一座山是多后端适配。你有一个卷积算子CPU 上想调 oneDNNGPU 上想生成 CUDA kernelNPU 上要生成自定义指令。如果每个后端都从计算图直接翻译那你有 N 个后端就要写 N 套翻译逻辑而且每套逻辑里都要重复处理广播、类型提升、边界填充这些公共问题。改一个语义N 个地方都要跟着改。第二座山是优化 Pass 的复用。常量折叠、死代码消除、算子融合这些优化逻辑上跟目标硬件没关系但如果你把 IR 设计得跟某个后端绑死这些 Pass 就没法跨后端复用。结果就是每个后端各写一套优化重复造轮子。第三座山是渐进式 lowering 的缺失。从高层算子到机器码中间跨度太大。如果只有高层 IR和机器码两级那 lowering 就是一个巨大的、难以调试的黑盒。出了问题你根本不知道是哪一步引入的。MLIR 的设计恰好对着这三座山来的。它的核心思路是用方言Dialect把不同抽象层级隔开用渐进式 lowering 把大跳跃拆成小步用统一的 Pass 基础设施让优化跨方言复用。这就是标题里可组合模块化和结构化、可重定目标的实际含义。1.2 方言机制到底解决了什么方言是 MLIR 里最核心的概念。你可以把它理解成一套操作、类型和属性的命名空间。比如tensor方言管张量类型和基础操作linalg方言管结构化线性代数运算arith方言管算术scf方言管结构化控制流llvm方言直接对应 LLVM IR。关键在于这些方言不是孤立的它们可以互相嵌套、互相转换。一个linalg.matmul可以 lower 成scf.for循环加arith.mulf和arith.addf再 lower 成llvm方言最后导出成 LLVM IR。每一步 lowering 都是一个独立的、可测试的 Pass。这种设计带来的直接好处是你新增一个后端只需要写从某个中间方言到目标方言的 lowering前面的公共优化全部白嫖。比如你想支持一个新的加速器只要它能在linalg或vector层面接住那从高层到linalg的所有优化你都不用重写。我实测下来一个中等复杂度的张量编译器用 MLIR 重构后后端相关代码量能压到原来的三分之一左右而且新增后端的周期从几周缩短到几天。这不是夸张是因为公共部分真的被复用掉了。1.3 可重定目标的关键接口与契约可重定目标这个词听起来玄乎说白了就是同一套上层逻辑能生成面向不同硬件的代码。MLIR 实现这一点靠的是接口Interface和契约Contract。比如BufferizableOpInterface它定义了一个操作如何被 bufferize的契约。任何方言只要实现了这个接口就能接入统一的 bufferization 流程。再比如TargetInterface它定义了目标硬件的查询接口你问它支持不支持某个类型某个操作的代价是多少它来回答。这种接口化的设计让上层 Pass 不需要知道具体后端是谁只需要按接口调用。后端只需要实现接口就能被上层流程接纳。这就是可组合的技术底座——组件之间通过接口通信而不是硬编码依赖。理解了这一层你再看 MLIR 的代码生成流程就不会觉得它是一堆零散 Pass 的堆砌而是一个有明确契约的分层系统。2. 结构化代码生成的核心链路拆解2.1 从计算图到高层方言入口怎么选张量编译器的入口通常是某种计算图描述可能是 ONNX、可能是自定义的图格式也可能是 Python 里直接 trace 出来的。第一步是把它转成 MLIR 的高层方言。这里有个选型问题入口方言用哪个常见的有几个选择。linalg方言适合结构化、规则的计算比如矩阵乘、卷积、逐元素运算。它的抽象层级适中既有足够的信息做优化又不会太高层导致 lowering 困难。tosa方言适合从 TensorFlow Lite 这类框架导入算子定义比较固定。mhlo方言适合从 JAX、TensorFlow 这类框架导入算子粒度更细。我的经验是如果你是自己定义计算图优先选linalg。因为linalg的structured op概念非常清晰一个操作由迭代空间、输入输出和标量计算三部分组成这种结构化的描述让后续的 tiling、fusion、vectorization 都有统一的处理框架。具体转换时你需要为每个高层算子写一个 pattern把它映射成linalg操作。比如一个带 bias 的卷积可以拆成linalg.conv_2d加linalg.generic做 bias 加法。这一步的关键是保持语义等价不要在这一层做激进的优化把优化留给后面的 Pass。2.2 渐进式 lowering 的层级划分从linalg到最终代码中间要经过好几层。我把常见的层级列一下方便你建立整体认知。层级代表方言主要职责高层计算linalg, tosa结构化算子描述循环与内存scf, memref循环结构、内存分配与访问向量与并行vector, gpu向量化、线程映射底层算术arith, math标量算术、数学函数目标相关llvm, nvvm, rocdl目标指令、导出这个划分不是绝对的但大体遵循从抽象到具体、从结构化到指令化的顺序。每一层 lowering 只负责一件事比如linalg到scf只负责把结构化算子展开成循环scf到cf只负责把结构化控制流转成基本块跳转。这种渐进式的好处是可调试。如果最终代码有问题你可以逐层 dump IR看是哪一层引入的。MLIR 提供了--mlir-print-ir-after-all这样的选项能把每个 Pass 之后的 IR 都打出来排查问题非常方便。2.3 张量到缓冲的转换时机有一个坑我必须提前说tensor 和 memref 的转换时机。MLIR 里tensor类型是不可变的、值语义的memref类型是可变的、引用语义的。高层优化在tensor层面做因为值语义让数据流分析更简单底层代码生成在memref层面做因为要精确控制内存。转换的时机很关键。转太早你会失去很多基于值语义的优化机会转太晚很多需要内存信息的优化又做不了。我的经验是在 tiling 和 fusion 之后、vectorization 之前做 bufferization。这时候计算结构已经确定内存访问模式也清晰了转成memref正好能支撑后续的向量化和循环优化。MLIR 的One-Shot Bufferize是现在推荐的方案它通过分析每个 tensor 的读写冲突来决定是否要拷贝比早期的linalg-comprehensive-bufferize更高效。用的时候注意配置allow-return-allocs和bufferize-function-boundaries这些选项具体含义后面复现部分会讲。3. MLIR 15 全流程复现从环境到可执行3.1 环境准备与构建方式选择复现 MLIR 流程第一步是拿到可用的工具链。有两条路装预编译包或者从源码构建。预编译包最省事很多 Linux 发行版的仓库里就有mlir相关的包或者用 LLVM 官方发布的预编译版本。但预编译版本通常是 release 构建不带 assertions调试的时候看不到内部检查而且版本可能和你需要的对不上。从源码构建更可控尤其是你要做二次开发的时候。MLIR 15 对应的 LLVM 15构建命令大致是这样git clone --branch llvmorg-15.0.0 https://github.com/llvm/llvm-project.git cd llvm-project cmake -S llvm -B build -G Ninja \ -DCMAKE_BUILD_TYPERelease \ -DLLVM_ENABLE_PROJECTSmlir \ -DLLVM_TARGETS_TO_BUILDX86;NVPTX \ -DLLVM_ENABLE_ASSERTIONSON cmake --build build --target mlir-opt mlir-translate这里有几个点要注意。LLVM_ENABLE_PROJECTS只开mlir就够了除非你还要用 Clang。LLVM_TARGETS_TO_BUILD按需开做 CPU 代码生成开X86做 GPU 开NVPTX。LLVM_ENABLE_ASSERTIONSON强烈建议打开虽然会慢一点但能在 Pass 出错时第一时间报出来省去大量排查时间。构建完成后build/bin/下会有mlir-opt、mlir-translate这些工具。mlir-opt是跑 Pass 的主力mlir-translate负责 MLIR 和 LLVM IR 之间的转换。3.2 一个最小可跑的 linalg 例子环境好了先跑一个最小例子建立信心。写一个matmul.mlirfunc.func matmul(%A: tensor4x8xf32, %B: tensor8x16xf32) - tensor4x16xf32 { %init arith.constant dense0.0 : tensor4x16xf32 %0 linalg.matmul ins(%A, %B : tensor4x8xf32, tensor8x16xf32) outs(%init : tensor4x16xf32) - tensor4x16xf32 return %0 : tensor4x16xf32 }用mlir-opt跑一下 tilingmlir-opt matmul.mlir \ --linalg-tiletile-sizes2,4 \ --convert-linalg-to-loops \ -o matmul_tiled.mlir--linalg-tile把矩阵乘按 2x4 分块--convert-linalg-to-loops把linalg.matmul展开成scf.for循环。跑完你 dump 一下 IR能看到循环结构已经出来了。这一步的意图是验证工具链通了同时让你直观看到 lowering 的效果。很多人一上来就搞复杂例子结果卡在环境问题上很打击信心。先跑通最小例子再逐步加复杂度这是我一贯的做法。3.3 完整 pipeline 的 Pass 编排真正做代码生成需要把 Pass 串成 pipeline。下面是一个面向 CPU 的典型 pipeline我按顺序列出来并解释每一步的意图。mlir-opt input.mlir \ --linalg-fuse-elementwise-ops \ --linalg-tiletile-sizes8,8 \ --linalg-bufferize \ --convert-linalg-to-loops \ --convert-scf-to-cf \ --convert-vector-to-llvm \ --convert-memref-to-llvm \ --convert-func-to-llvm \ --reconcile-unrealized-casts \ -o output.mlir逐条说--linalg-fuse-elementwise-ops把逐元素操作融合进生产者减少中间张量。这一步在tensor层面做效果最好。--linalg-tile分块为后续的 cache 友好访问和向量化做准备。tile size 的选择要看目标硬件的 cache 大小8x8 是个保守的起点。--linalg-bufferizetensor 转 memref。注意这一步之后就不能再做基于值语义的优化了。--convert-linalg-to-loops结构化算子展开成循环。--convert-scf-to-cf结构化控制流转成基本块跳转为导出 LLVM IR 做准备。--convert-vector-to-llvm、--convert-memref-to-llvm、--convert-func-to-llvm把各个方言转成 LLVM 方言。--reconcile-unrealized-casts清理类型转换过程中产生的 unrealized cast这是 MLIR 里常见的一步收尾。这套 pipeline 跑通后用mlir-translate --mlir-to-llvmir就能导出 LLVM IR再交给 LLVM 后端生成机器码。注意Pass 的顺序不能随便调。比如 bufferize 必须在 tiling 之后否则 tiling 会失去结构化信息convert-scf-to-cf 必须在所有依赖 scf 的 Pass 之后。顺序错了轻则报错重则生成错误代码。4. 踩坑实录复现过程中最容易翻车的几个点4.1 bufferization 的读写冲突与拷贝第一个大坑是 bufferization。One-Shot Bufferize会分析每个 tensor 的读写关系如果发现某个 tensor 在被读取的同时又被写入就会插入一个拷贝来保证语义正确。这个拷贝有时候是必要的有时候是多余的取决于你的 IR 结构。我遇到过一个情况一个逐元素操作链因为中间某个 tensor 被两个下游同时消费bufferize 时插入了大量拷贝性能直接掉了一半。排查的时候用--one-shot-bufferizetest-analysis-only把分析结果打出来发现是某个linalg.generic的输出被复用导致的。解决办法有两个一是调整 IR 结构让每个 tensor 只有一个消费者二是用--linalg-fuse-elementwise-ops提前把链融合掉减少中间 tensor。我一般两个都用先融合再 bufferize拷贝能少很多。4.2 类型不匹配导致的 lowering 失败第二个坑是类型不匹配。MLIR 是强类型的每个操作对操作数类型有严格要求。比如linalg.matmul要求输入是tensor或memref如果你传了个vector直接报错。更隐蔽的是索引类型。MLIR 里索引默认是index类型但有些操作要求i32或i64。在 lowering 过程中如果类型转换没处理好会出现arith.index_cast满天飞的情况最后导出 LLVM IR 时一堆冗余转换。我的做法是在高层就统一索引类型尽量用index需要具体整数类型的地方显式转换并且用--canonicalize把冗余转换消掉。canonicalize 这个 Pass 看着不起眼但实际用起来能清掉大量噪音建议每个阶段之后都跑一遍。4.3 Pass 顺序引发的隐蔽错误第三个坑最阴险Pass 顺序错了但不报错只是生成的结果不对。比如你先跑了--convert-scf-to-cf再跑--linalg-tiletiling 需要 scf 结构结果发现没有要么报错要么静默跳过最后生成的代码根本没分块。这种问题的排查方法是逐 Pass dump IR。用--mlir-print-ir-after-all把每个 Pass 之后的 IR 打出来对比预期和实际。虽然输出很长但定位问题非常有效。我一般会把输出重定向到文件然后搜索关键操作名看它在哪一步消失或变形了。还有一个经验是把 pipeline 拆成小段测试。不要一次性跑完整 pipeline先跑前三个 Passdump 看结果对不对再加后面的。这样出问题能快速定位到具体是哪个 Pass 引入的。5. 可重定目标实战同一 IR 生成多后端代码5.1 从 linalg 到 GPU 的 lowering 路径可重定目标的价值在多后端场景下体现得最明显。同一份linalgIR面向 CPU 走一套 lowering面向 GPU 走另一套。GPU 路径大致是这样mlir-opt input.mlir \ --linalg-tiletile-sizes16,16 \ --convert-linalg-to-parallel-loops \ --gpu-map-parallel-loops \ --convert-parallel-loops-to-gpu \ --gpu-kernel-outlining \ --convert-gpu-to-nvvm \ --convert-nvvm-to-llvm \ -o output_gpu.mlir关键区别在于--convert-linalg-to-parallel-loops和--gpu-map-parallel-loops。前者把结构化算子展开成并行循环后者把并行循环映射到 GPU 的 block 和 thread 维度。--gpu-kernel-outlining把 kernel 代码提取成独立的函数--convert-gpu-to-nvvm转成 NVVM 方言最后导出 LLVM IR。对比 CPU 路径你会发现前面的 tiling 和 fusion 是共用的只有从并行循环开始才分叉。这就是可重定目标的实际收益——公共部分复用差异部分隔离。5.2 目标相关信息的注入方式不同后端需要不同的目标信息比如 GPU 需要知道 block 大小、shared memory 容量CPU 需要知道 cache 大小、SIMD 宽度。这些信息怎么注入MLIR 提供了几种机制。一是属性Attribute你可以在操作上挂属性来传递信息比如linalg.tile的 tile size 就是个属性。二是目标接口通过TargetInterface查询目标能力。三是Pass 选项在跑 Pass 时通过命令行参数传入。我的建议是能用属性就用属性因为属性跟着 IR 走dump 出来一目了然调试方便。目标接口适合查询类的信息Pass 选项适合全局配置。三者结合使用但不要滥用否则信息散落各处维护起来头疼。5.3 多后端共用的优化 Pass 清单哪些 Pass 是跨后端共用的我整理了一份清单这些 Pass 在 CPU 和 GPU 路径里都会跑。Pass作用是否跨后端canonicalize规范化 IR是cse公共子表达式消除是linalg-fuse-elementwise-ops逐元素融合是linalg-tile分块是tile size 不同constant-fold常量折叠是symbol-dce死符号消除是可以看到优化类 Pass 基本都是共用的只有 lowering 类 Pass 才分后端。这正是 MLIR 分层设计的威力——把跟目标无关的优化和跟目标相关的 lowering 彻底分开。实际项目里我会把这些共用 Pass 封装成一个函数CPU 和 GPU 路径都调用它只是传入不同的 tile size 等参数。这样既保证了优化一致性又保留了后端灵活性。6. 从 MLIR 到可执行文件的最后一公里6.1 导出 LLVM IR 与链接MLIR 本身不生成机器码它导出 LLVM IR剩下的交给 LLVM。导出命令mlir-translate --mlir-to-llvmir output.mlir -o output.ll然后用llc编译成目标文件再用clang链接成可执行文件llc -filetypeobj output.ll -o output.o clang output.o -o output -lm这一步的坑主要在运行时库。MLIR 生成的代码可能依赖一些运行时函数比如内存分配、打印等。这些函数需要链接对应的库否则会出现未定义符号。CPU 路径一般依赖libmlir_runner_utils和libmlir_c_runner_utilsGPU 路径还要加上 CUDA 运行时。我建议先用mlir-cpu-runner跑通它内置了这些运行时库能直接执行 MLIR 文件。跑通之后再走完整的导出链接流程这样能把运行时问题和代码生成问题分开排查。6.2 性能验证与对比方法代码生成出来得验证性能。我的做法是先正确性后性能。正确性用mlir-cpu-runner跑对比参考实现的结果。性能用perf或nsight这类工具测看热点在哪。对比的时候要注意公平性。跟手写代码比要保证算法一致、编译选项一致跟库调用比要保证输入规模一致、预热充分。我见过太多对比不公平导致的错误结论比如拿没优化的 MLIR 代码跟高度优化的库比然后得出MLIR 性能不行的结论这是不客观的。一个实用的方法是逐层对比。先看 tiling 前后的性能差异再看 vectorization 前后的差异这样能定位到具体哪个优化起了作用哪个优化没效果。MLIR 的 Pass 粒度细正好支持这种细粒度的性能分析。6.3 常见运行时错误与定位最后说几个常见的运行时错误和定位方法。段错误多半是内存访问越界。用mlir-cpu-runner加--entry-point-resultvoid跑配合gdb定位。检查 bufferization 后的 memref 形状和索引是否正确。结果错误多半是 lowering 语义不对。逐层 dump IR对比每一层的语义。重点检查循环边界、索引计算、类型转换。性能异常多半是优化没生效。用--mlir-print-ir-after-all看优化 Pass 是否真的改了 IR还是静默跳过了。我踩过最深的一个坑是索引类型不一致导致的静默错误。某个循环的索引是i32但另一个地方期望index隐式转换后数值溢出结果在小规模输入下正常大规模输入下就错了。这种问题只能靠严格的类型检查和充分的测试覆盖来避免。7. 我在实际项目中的几点体会做了一段时间 MLIR 之后有几个体会比较深分享出来供参考。第一不要试图一步到位。MLIR 的方言和 Pass 很多一开始容易贪多想把所有优化都用上。实际上先把最小 pipeline 跑通再逐步加优化效果更好。我见过有人一上来就配了二十几个 Pass结果 IR 被改得面目全非根本没法调试。第二IR 的可读性很重要。MLIR 的 IR 是给人看的不是只给机器看的。写 pattern 的时候尽量生成结构清晰、命名规范的 IR这样调试的时候一眼就能看懂。我习惯给关键的 tensor 和循环起有意义的名字虽然多花点时间但排查问题时省的时间更多。第三测试要覆盖边界情况。张量编译器最容易在边界上出错比如空张量、单元素张量、非对齐形状。这些情况在正常输入里不常见但一旦出现就是崩溃。我的做法是专门写一组边界测试每次改 Pass 都跑一遍。第四善用社区资源。MLIR 的文档和示例质量很高mlir/test目录下有大量可运行的测试用例遇到问题先搜测试用例往往能找到答案。另外 MLIR 的 discourse 论坛也很活跃提问一般都能得到回复。最后再分享一个小技巧用--debug-only打开 Pass 的内部日志。很多 Pass 支持--debug-onlypass-name来输出内部调试信息比如 tiling 的决策过程、bufferize 的冲突分析结果。这些信息在排查复杂问题时非常有用比单纯看 IR 直观得多。用之前记得构建时打开 assertions 和 debug 支持。