前阵子为了把手里的稀疏张量内核搬到不同硬件上我把 MLIR 相关的几篇经典论文翻出来重读了一遍然后顺手用 MLIR 15 走了一遍完整的代码生成流程。读完之后最大的感受是过去我们写张量编译器总觉得目标平台绑定和优化策略是一团浆糊你不动底层就做不了局部优化一改底层又全盘重写。MLIR 的答案是把编译器本身拆成一层一层的方言让每一层只做一件事然后用 pass 把层级串起来——所谓可组合模块化代码生成就是这个意思。这篇精读我不打算逐句翻译论文而是想回答三个问题第一大家天天挂在嘴边的结构化和可重定目标在 MLIR 工程实践里到底指什么第二一个张量算子从数学描述到机器指令中间到底经历了哪几层变化第三如果你手上只有 MLIR 15 的二进制怎么把一条linalg.matmul完整落到 LLVM IR 和汇编。后面还会专门列一节复现时踩过的坑那些坑基本都在官方文档的角落里不自己跑一遍很难发现。1. 先把话说透张量编译器的老问题出在哪儿1.1 手写代码生成为什么是一锤子买卖很多做算法的人在接触编译底层之前以为代码生成就是后端加几个指令模板。直到自己写过一个加速器内核才发现张量编译器最头疼的不是生成指令而是生成之前那些结构信息早就在前几轮优化里丢光了。传统的编译器前端把循环、数组、张量算子一路降到基本块和 phi 节点循环结构变成了跳转图张量维度和数据布局信息变成了地址计算——到了后端你只看到一堆标量操作想再做分块、融合、向量化只能靠非常脆弱的重建算法去猜原来的语义。这种情况在 Halide、TVM 时代有所缓解因为它们把 schedule 和 compute 分开但调度描述和后端目标仍然绑定得比较死。你为 CUDA 写好的 tile thread binding 规则换到 AVX2 就要重写一层你为 AVX2 设计好的 vector width换到 Mali GPU 又要改一遍。说白了这就是一锤子买卖每个目标平台一条专用代码路径每加一个后端就多一批维护负担。1.2 可重定目标意味着什么MLIR 论文里反复出现可重定目标这个词。它并不是说同一份 IR 能自动在所有硬件上跑出最好性能而是指把描述计算和映射到硬件这两件事彻底分开并且让映射过程本身也分成多级、可复用。所谓一次描述多个出口一个linalg.matmul在 CPU 出口可以走 scf.for vector ops AVX2 指令在 GPU 出口可以映射到 gpu.thread/gpu.block再降到 NVVM在自研加速器出口甚至可以保留 tiling 结构直接生成定制的循环流水。关键在于这些不同出口共享了前面一大段上层 IR 和优化 pass。你不需要为每个硬件重写如何表达 matmul这件事只需要针对目标平台选择不同的 lower 路径和参数。从工程角度看这比一个巨型后端要划算得多。下面这个表大致描述了不同目标的典型出口目标平台较高层优化主要下降路径最终出口通用 CPUlinalg tile/fuse、vectorizescf.for → affine/vectorLLVM dialect → X86/AArch64GPUlinalg tile distributegpu dialect → NVVM/ROCDLLLVM dialect → PTX/AMD自研加速器linalg tile、bufferization 定制自定义 dialect 或直接 LLVMLLVM dialect / 自定义后端1.3 结构化是这次设计里最关键的字眼标题里的结构化不是形容词而是名词在 MLIR 里它是一个具体的设计选择让高层次的算子保持数学语义不要过早摊平。举个例子linalg.matmul的输入输出是张量它自带循环域和索引映射关系当它出现在 IR 中时编译器知道这是一个矩阵乘可以做融合、分块、矢量化。可一旦被降成scf.for加memref.load这些信息就没了只能靠循环结构去猜。打个比方结构化 IR 像一份带科目层级的财务报表收入和成本各自归位审计师能直接看懂非结构化 IR 把所有数字摊平成流水账每一笔都在但想按业务口径重新聚合就非常痛苦。MLIR 的核心思想之一就是尽量延长带科目层级的阶段把摊平成流水账这一步推迟到最后。这也是为什么 linalg、tensor、scf 这些 dialect 会被反复拿出来讲——它们是同一份计算的三种不同精细度描述。2. 论文核心思路把编译过程拆成三个可组合层级2.1 第一层tensor / linalg数学结构还在最上层是 tensor dialect 和 linalg dialect。tensor 负责表达张量值语义比如一个tensor64x64xf32是一个不可变的值不存在内存地址这种概念。linalg 则负责表达带明确计算结构的算子包括 matmul、conv、pooling、pointwise 等。这一层的 IR 长这样func.func matmul(%A: tensor64x64xf32, %B: tensor64x64xf32) - tensor64x64xf32 { %cst arith.constant 0.0 : f32 %init linalg.init_tensor : tensor64x64xf32 %filled linalg.fill ins(%cst : f32) outs(%init : tensor64x64xf32) - tensor64x64xf32 %res linalg.matmul ins(%A, %B : tensor64x64xf32, tensor64x64xf32) outs(%filled : tensor64x64xf32) - tensor64x64xf32 return %res : tensor64x64xf32 }这一段 IR 里完全看不到循环变量、内存地址、维度计算它只描述了一件事C A * B 0。这就是结构化的好处tiling、fusion、distribution 这些优化都可以在 linalg 层级直接做而不需要先恢复循环结构。2.2 第二层scf / affine / vector把顺序和并行显性化当优化做完需要把结构化的算子变成可执行的控制流就进入第二层。convert-linalg-to-loops会把上面的 matmul 展开成三层循环scf.for %i 0 to 64 step 1 { scf.for %j 0 to 64 step 1 { scf.for %k 0 to 64 step 1 { %a memref.load %A[%i, %k] : memref64x64xf32 %b memref.load %B[%k, %j] : memref64x64xf32 %c memref.load %C[%i, %j] : memref64x64xf32 %prod arith.mulf %a, %b : f32 %sum arith.addf %c, %prod : f32 memref.store %sum, %C[%i, %j] : memref64x64xf32 } } }到这里计算的顺序显式化了内存的读写作也显式化了但循环仍然是高级结构不是裸跳转。这个层级是优化的重要阵地可以调整循环分块、交换顺序、向量化、并行化。如果目标是 CPU往往还会经过 vector dialect 把标量运算打包成向量宽度对应的运算如果目标是 GPU则在这个阶段把循环分布到 thread/block。2.3 第三层LLVM 方言与机器边界最后所有高级结构都要降成 LLVM 方言离机器只剩一步。这一层里memref 变成裸指针加上必要的维度元数据arith 变成llvm.add、llvm.mulscf/affine 变成llvm.br、llvm.cond_br。张量语义彻底消失内存布局、指针运算、调用约定这些底层细节成为主角。这个分层序列最关键的地方在于每一步都是显式、可观测、可中断的。你不是从一个黑盒跳到另一个黑盒而是每一条 pass 命令的输入和输出都在你眼前每一层 IR 都可以打印出来检查。这种可观测性对做编译器的人有多重要只有被老式编译器的 debug log 折磨过的人才懂。3. 组合的机制可插拔管线与看得见的中间表示3.1 pass 不是流水线工人而是处理 IR 的函数MLIR 把一个编译过程组织成若干个 pass每个 pass 的职责非常单一。比如convert-linalg-to-loops只负责把 linalg 算子变成 scf 循环convert-scf-to-cf只负责把结构化控制流变成基本块跳转convert-arith-to-llvm只负责把算术运算降到 LLVM 方言。它们之间的依赖关系取决于你当前 IR 里有哪些 dialect而不存在一个标准的编译流程写死在框架里。这一点跟传统印象很不一样。很多人以为编译器是一个大流水线前端喂进去后端吐出来中间不可打断。MLIR 的可组合体现在你可以自由地挑选 pass、自由地排列顺序、甚至在中间插入自己的自定义方言和自定义 pass。整个管线就是一个可复用的函数组合不是只有一条固定赛道。3.2 中间表示的可打印性是调试神器MLIR 几乎所有的 IR 都能以文本形式导出并重新读入这种 round-trip 能力是我最喜欢的特性。跑完一个 passmlir-opt会把整个模块打印出来你不满意直接改文本再进去。这在做自定义加速器后端时价值极高你可以在任意阶段截断把 IR 保存成文件作为后续调试的输入而不需要构造一套完整的假程序。论文里强调模块化代码生成本质上就是强调每一层的输入输出都是合法、可读、可替换的。如果中间表示不可打印再好的模块化设计也只是纸上谈兵。3.3 transform dialect把组合也变成数据MLIR 15 里已经有 transform dialect 的雏形了。它干的事是把 pass pipeline 的部分决策写成 IR 中的一组 transform op让编译器在执行时按这个描述来 tiling、fusion、lower。换句话说不只是中间表示是数据连如何转换中间表示这一策略本身也是数据。虽然 15 里的 transform dialect 还很年轻API 变动也大但它背后的思路对理解可组合很重要不要把策略硬编码在一个 pass 里而是让它成为可以按输入 IR 和目标任务动态选择的脚本。到后续版本里transform 能做的事越来越多已经成为写 MLIR 专项编译器时很关键的一层。4. MLIR 15 全流程复现从 linalg.matmul 到汇编4.1 环境准备只要三个工具完整复现的第一步是得到可用的工具链。如果你不想从源码编译可以尝试找官方或社区打包的 LLVM 15 预编译版本但想可控地复现我建议还是自己编一次顺便把 MLIR 的源码熟悉起来。git clone --depth 1 --branch llvmorg-15.0.0 https://github.com/llvm/llvm-project.git cd llvm-project cmake -G Ninja -S llvm -B build \ -DLLVM_ENABLE_PROJECTSmlir \ -DLLVM_TARGETS_TO_BUILDhost \ -DCMAKE_BUILD_TYPERelease \ -DLLVM_ENABLE_ASSERTIONSON \ -DCMAKE_C_COMPILERclang \ -DCMAKE_CXX_COMPILERclang cmake --build build --target mlir-opt mlir-translate mlir-cpu-runner -j$(nproc)这里有几个选择要解释一下。LLVM_TARGETS_TO_BUILD只保留 host 目标可以大幅缩短编译时间如果你后面要交叉到 ARM 或者做 GPU需要把它换成AArch64;X86;NVPTX这类列表。LLVM_ENABLE_ASSERTIONS在开发期建议打开很多 pass 的检查逻辑只在断言开启时生效报错信息会友好很多。编译耗时取决于机器但单独 build 那三个 target 通常不会太离谱建议配好 ccache后面改代码重编会舒服很多。4.2 准备一个最小但真实的张量计算用例我们用一个 64x64 的矩阵乘作为例子它简单到不会引入多余噪音又足够展现 MLIR 多级下降的本质。把下面内容保存为matmul.mlirfunc.func matmul(%A: tensor64x64xf32, %B: tensor64x64xf32) - tensor64x64xf32 { %cst arith.constant 0.0 : f32 %init linalg.init_tensor : tensor64x64xf32 %filled linalg.fill ins(%cst : f32) outs(%init : tensor64x64xf32) - tensor64x64xf32 %res linalg.matmul ins(%A, %B : tensor64x64xf32, tensor64x64xf32) outs(%filled : tensor64x64xf32) - tensor64x64xf32 return %res : tensor64x64xf32 }这个用例里包含了arith.constant、linalg.init_tensor、linalg.fill、linalg.matmul四个关键算子已经足够覆盖从 tensor 世界到 memref 世界的完整变化。4.3 四段管线逐级观察第一步是 bufferization把 tensor 语义变成 memref 内存语义。MLIR 15 里推荐用 one-shot bufferizemlir-opt --one-shot-bufferizebufferize-function-boundariestrue matmul.mlir -o matmul.buf.mlir这一步的变化非常明显函数签名从 tensor 参数和返回值变成了纯 memref 参数原来的linalg.matmul也从 tensor 版本变成 memref 版本。如果你看到func.func matmul(%arg0: memref64x64xf32, %arg1: memref64x64xf32, %arg2: memref64x64xf32)那就说明 bufferization 成功了。第二步是把 linalg 算子降成循环mlir-opt --convert-linalg-to-loops matmul.buf.mlir -o matmul.loops.mlir此时 IR 里出现scf.for三重循环循环体内部是memref.load、arith.mulf、arith.addf、memref.store和前面展示的片段一致。到这一步编译器已经理解了计算顺序可以做循环层面的大多数优化。第三步是把结构化控制流降成 CFGmlir-opt --convert-scf-to-cf matmul.loops.mlir -o matmul.cf.mlir输出里会出现cf.br、cf.cond_br这样的基本块跳转。注意 scf 的循环变量和边界信息开始消失剩下来的是如何判断、如何跳转。第四步是降到 LLVM 方言mlir-opt \ --convert-memref-to-llvm \ --convert-arith-to-llvm \ --convert-cf-to-llvm \ --convert-func-to-llvm \ --reconcile-unrealized-casts \ matmul.cf.mlir -o matmul.llvm.mlir最终再用mlir-translate导出 LLVM IR并用llc得到目标汇编mlir-translate --mlir-to-llvmir matmul.llvm.mlir -o matmul.ll llc matmul.ll -o matmul.s本地跑完可以看到生成的汇编里是一串标量mulss/addss指令夹在循环控制指令之间。这是非常朴素的结果我们没有做任何 tiling 和 vectorization编译器只是忠实地把循环代码映射到了指令集。想要看到 AVX2 向量指令需要在链路里加入 vectorization 或直接对linalg.matmul做向量化 tiling这一步留给你自己用--linalg-vectorize做实验。4.4 运行验证给矩阵乘加一个入口如果只想看到 LLVM IR 和汇编到 4.3 就足够了。但全流程复现最好还是确认生成的代码真的可执行。最简单的方式是再写一个main函数在matmul.mlir里加一个入口之后用mlir-cpu-runner直接执行 JIT 好的代码mlir-cpu-runner matmul.llvm.mlir -e main -entry-point-resultvoid为了让输出有意义可以在main里初始化 A、B 并调用 matmul再用vector.print把结果的前几个元素打出来。这一步会额外涉及向量 dailect 的 lower建议单独作为练习不要混在纯代码生成流程里否则出现问题时很难定位是哪一层的问题。5. 复现时踩过的坑15 版特有5.1 别再用 linalg-bufferize 了MLIR 15 里-linalg-bufferize还能用但已经是旧时代方案。它只 bufferize linalg 算子对 tensor.cast、tensor.extract_slice、tensor.insert_slice 这些算子处理得不彻底经常留下一堆tensor.cast和bufferization.to_memref到后续 LLVM lowering 时爆出一堆类型不匹配错误。one-shot-bufferize 是更完整的方案它一次性把整个函数体的 tensor 语义转成 memref并统一处理函数边界。我在最初复现时就偷懒照抄了旧博客的-linalg-bufferize结果卡在unrealized_conversion_cast上整整半天。换成 one-shot 后干净利落地一次通过。5.2 函数边界tensor 返回值会卡住后续 passone-shot-bufferize如果不开bufferize-function-boundariestrue默认只处理函数内部的张量函数签名里的 tensor 参数和返回值会保留原样。结果convert-linalg-to-loops之后函数边界仍然有 tensor 类型后续 pass 根本不知道该怎么处理这种半个函数已 bufferize的状态。我建议从第一步就把函数边界打开除非你有特殊理由要保留函数外部的 tensor 接口。打开之后原本的返回值会变成最后一个 memref 参数调用方的语义要相应调整。5.3 pass 顺序与 unrealized_conversion_castMLIR 在类型不匹配时会插入unrealized_conversion_cast作为占位符这些 cast 在后续 pass 中如果一直没被解析掉会在 LLVM lowering 阶段变成一堆难懂的报错。常见的处理方式是在降级链末尾加一个--reconcile-unrealized-casts它会尝试清理剩余的不一致类型匹配。顺序也很重要不要想着先convert-scf-to-cf再convert-linalg-to-loops。linalg 算子必须先在结构化循环层面被展开再处理 scf 控制流顺序反过来pass 会直接跳过不认识的 linalg op留到后面报错。5.4 版本漂移15 和 17/18 之间的差异如果你用的是 MLIR 16 之后的版本这套命令大概率不能直接跑。改名是常态部分 pass 名从--convert-func-to-llvm变成带新命名空间的版本memref 的 LLVM lowering 相关 pass 也经历过拆分和改名。最稳妥的做法是锁死版本如果一定要用新版本每跑一步都用mlir-opt --help确认 pass 名存在。问题MLIR 15 表现新版本变化bufferize用one-shot-bufferize即可15 之后整体移入 bufferization 包命名空间调整linalg 工具 passconvert-linalg-to-loops可用部分版本拆成-expand-strided-metadata前后分开func loweringconvert-func-to-llvm可用新版本可能要求带finalize阶段cast 清理reconcile-unrealized-casts可用新版本需要在 pass 链里显式指定版本漂移是 MLIR 项目最现实的问题。GitHub 上不少 issue 其实不是 bug而是提问者用了 A 版本的命令去跑 B 版本的工具链。复现任何 MLIR 相关实验第一件事就是查看--help。6. 跳出论文这套方法论能用到哪儿6.1 张量编译器之外工具链式代码生成的老问题依然存在如果你用过 Simulink 做模型开发并生成嵌入式 C 代码会发现它面临的问题和 MLIR 论文完全同构模型是结构化描述目标设备多种多样生成代码的效率和可验证性都取决于能否在较高抽象层保留语义信息。老牌的 Embedded Coder 是封闭工具链模型一层、代码一层、硬件一层中间每一层都是黑盒。MLIR 的多级可组合思路恰好可以给这类工具链提供一个开放的设计参考把模型、优化、目标映射拆开让每一层都可见、可替换。6.2 PLC 与结构化程序生成结构化不是学术黑话PLC 领域里IEC 61131-3 的梯形图、功能块图、结构化文本本质上都是结构化编程模型。把 PLC 程序自动生成 C/C 甚至原生二进制同样面临结构如何保留、如何按目标控制器配置做代码生成的问题。MLIR 的分层思想在这里一样适用高层的功能块图保留控制流与数据流语义中间层做寄存器分配风格优化底层再映射到具体控制器的函数 ABI。很多工业软件公司已经在往这个方向靠。6.3 AI 结构化输出与结构化建模约束即代码生成最近 AI 领域经常讨论结构化输出比如让大模型严格按 JSON Schema 生成内容。你会发现这跟编译器做的事情本质上是同一件事在生成过程中保留约束不要让不合法的中间产物出现。MLIR 的 dialect 就是一套类型化约束系统非法操作和非法类型在早期就会被拒绝而不是等到生成随机字节码之后再纠错。工具链之间互相借鉴的潜力其实比想象中要大。从一个具体张量算子出发最终能看到模型推理、嵌入式工具链、工业控制代码生成这几条线的交汇点这种展开方式也是我读论文时最享受的部分好的抽象一旦建立它不会只在一个领域有用。最后再说点个人体会第一次完整跑通这条链路的时候我对 MLIR 复杂方言生态的抱怨少了很多。它真正解决了我在做加速器代码生成时最痛苦的问题每加一个硬件不用把上层代码重构一遍。MLIR 15 的很多接口现在回头看已经有些旧了但这套分层的思维没有过时。如果你也想踩一遍这条路我的建议是从最小的 matmul 开始手动跑完 bufferize、loop、CFG、LLVM 方言四段每段都停下来打印 IR 看变化。这个过程比读十篇概述文章都管用跑通了之后你对结构化可重定目标可组合这几个词的感受会完全不一样。
