MLIR 之旅：人人都依赖的方言栈

# MLIR 漫游：人人都依赖的方言栈 来源：https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/ 无论你是否写过一行 MLIR 代码，只要你训练或部署模型，你就已经依赖了它。XLA 通过它进行降层，Triton 构建在它之上，Mojo 原生就是 MLIR，而 Torch-MLIR、IREE 和 ONNX-MLIR 的存在则是为了让各自的前端落入其中。一个单一的基础设施能够栖身于如此多本应相互竞争的栈之下，个中原因值得理解——因为它很大程度上解释了现代 ML 编译器究竟是如何构建的，以及它们的接缝在哪里。 这篇文章是一次 MLIR 漫游：它是什么，使其与众不同的方言（dialect）理念是什么，一个张量操作如何逐步降层为机器码，以及该基础设施提供了什么、又不提供什么。 ## MLIR 究竟是什么 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#what-mlir-actually-is 常见的误解是，MLIR 是“另一个像 LLVM IR 一样的 IR”。更准确地说，它是一个 IR **构造工具包**（construction kit）。LLVM IR 是一个单一的、固定的、低层表示：大致是一种带类型的汇编语言，配有 SSA 值。对于通往机器码的最后一公里来说，这是正确的抽象，但对于张量上的矩阵乘法来说，这却是错误的抽象。 历史上，每个需要高层表示的领域编译器都从头发明了自己的表示：XLA 有 HLO，Halide 有自己的 IR，TensorFlow 有它的图，每个都带着独立的 pass 管理器、序列化格式、验证器和一堆 bug。[^1] MLIR 的前提是，这些表示的共同点远多于不同点，并且这些共同部分可以一次性构建并共享：SSA、区域和基本块构成的 CFG、pass 基础设施、模式重写器、位置追踪和验证。领域之间的差异则通过**方言**（dialect）来表达。 MLIR 中一切的单位是**操作**（Operation）。一个操作有操作数和结果（SSA 值）、一组有类型的**属性**（编译时常量，如形状或步长），以及零个或多个**区域**（region），区域本身包含进一步操作的基本块。这最后一个属性使得 IR 真正成为多层次的：单个操作可以携带一个完整的嵌套计算，因此一个高层的 `linalg.generic` 和一个低层的 `llvm.add` 是相同种类的对象，只是处于不同海拔。 每个操作都属于一个方言，方言只是一个相关操作、类型和属性的名称空间。用语法来描述的话，这种关系既简短又递归。方言提供**词汇表**（vocabulary）：操作名、类型和属性，而操作的**形状**（shape）是通用的。下面的形式是从 MLIR 的文本语法中简化而来，以展示基本结构；权威产生式请参考语言参考。[^2] ```ebnf ; 方言是一个名称空间，贡献一系列操作、类型和属性。 dialect ::= (operation-def | type-def | attribute-def)* ; 操作的语法在所有方言中都是一致的。 operation ::= [ result ("," result)* "=" ] op-name "(" [ operand ("," operand)* ] ")" attr-dict? region* ":" type-signature op-name ::= dialect-name "." mnemonic ; 例如 linalg.matmul；通用形式中以字符串形式引用 attr-dict ::= "{" attr-entry ("," attr-entry)* "}" ; 编译时常量 region ::= "{" block+ "}" block ::= operation+ ; 操作包含区域，区域包含操作 -> 递归 result ::= ssa-value ; %C operand ::= ssa-value ; %A, %B type ::= "!" dialect-name "." mnemonic | builtin-type ; 例如 !llvm.ptr；内建类型：tensor<...>, memref<...> attr-entry ::= name "=" attribute-value ; 例如 1 : i64, "foo", #dialect.attr<...> ``` 这带来了两个后果。首先，操作名、类型和属性都由方言命名空间限定（`linalg.matmul`、`tensor<128x256xf32>`），因此“添加一个方言”就是扩展词汇表，而不触及语法——这正是为什么周围的基础设施可以独立于方言。其次，由于一个 `operation` 可能包含一个 `region`，而一个 `region` 又包含进一步 `operation` 的 `block`，这个结构可以无限嵌套，这种递归正是让一个操作能够携带整个计算（而非一条指令）的原因。 ## 方言：一个 IR，多种海拔 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#dialects-one-ir-many-altitudes 其定义特征是方言可以**共存**。一个处于编译中期的模块通常同时包含来自多个方言的操作，而降层（lowering）就是逐步用更低层的操作替换高层操作，直到只剩下目标方言。 一个 ML 流水线经过的方言，从高到低排列： - **高层（你意图）**：`stablehlo` 和 `tosa`（全张量算子集）、`linalg`（张量/缓冲上的结构化操作）、`tensor`（值语义的张量操作）。[^3] - **中层（如何结构化）**：`memref`（带有布局/步长的缓冲区）、`affine` 和 `scf`（循环嵌套和结构化控制流）、`vector`（SIMD）、`arith`（标量数学）。[^4] - **低层（在哪里运行）**：`llvm`（翻译为 LLVM IR 以面向 CPU）、`gpu` 加上 `nvvm`/`rocdl`（GPU 目标）、`spirv`（Vulkan/计算）。 MLIR 编码的技巧是选择**何时**从一个海拔下降到下一个。在高层停留太久，你将无法表达硬件特定的调度；过早下降到低层，你就丢弃了优化器所需的结构。 “共存”一词很容易被轻描淡写，所以这里用一个同时使用四个方言的函数来说明。此时还没有任何降层；这些操作只是共存在同一个 SSA 区域中，验证器会一起检查它们： ``` // 一个区域，四个方言；每一行标注了来自哪个方言。 func.func @scale_in_place(%buf: memref<1024xf32>, %a: f32) { %c0 = arith.constant 0 : index // 1. arith：循环边界是标量 index 常量 %c1 = arith.constant 1 : index // 2. arith %n = arith.constant 1024 : index // 3. arith scf.for %i = %c0 to %n step %c1 { // 4. scf：一个携带区域的结构化循环 %x = memref.load %buf[%i] : memref<1024xf32> // 5. memref：从显式缓冲区读取 %y = arith.mulf %x, %a : f32 // 6. arith：标量乘法 memref.store %y, %buf[%i] : memref<1024xf32> // 7. memref：将结果写回 } return // 8. func：终结符 } ``` 四个方言——`func`、`arith`、`scf` 和 `memref`——出现在同一个区域中，彼此之间没有任何阻抗不匹配。后续的 pass 可能会将 `scf.for` 重写为 `cf` 分支，或者将循环体向量化到 `vector` 方言，但在这个海拔上，它们只是简单地组合在一起。这种可组合性正是关键：方言不是你要在之间翻译的独立 IR，而是你在一个程序中混合使用的词汇表。 ```mermaid graph TD subgraph FE["前端"] SH["StableHLO / TOSA"] TM["Torch-MLIR"] end SH --> LIN TM --> LIN LIN["linalg + tensor(值语义)"] -->|bufferize| MR["memref(显式缓冲区)"] MR --> LOOPS["scf / afffine(循环嵌套)"] LOOPS --> VEC["vector(SIMD)"] VEC --> TGT{目标} TGT -->|CPU| LLVMD["llvm dialect → LLVM IR"] TGT -->|NVIDIA GPU| NVVM["gpu + nvvm → PTX"] TGT -->|Vulkan| SPV["spirv"] classDef hi fill:#e1f5ff,stroke:#4a6572,color:#1a1a1a; classDef mid fill:#fff3cd,stroke:#7a6a00,color:#1a1a1a; classDef dec fill:#ffe0b2,stroke:#7a5200,color:#1a1a1a; classDef tgt fill:#f3e5f5,stroke:#6a4a6a,color:#1a1a1a; class SH,TM,LIN hi; class MR,LOOPS,VEC mid; class TGT dec; class LLVMD,NVVM,SPV tgt; ``` ## 渐进降层的具体示例 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#progressive-lowering-concretely 以一个矩阵乘法为例。在最顶层，它是一个结构化的**操作**，作用于（值语义的）张量。形状存在于类型中，关于内存或循环的任何信息都还未提及： ``` // 形状是类型的一部分；这是纯数据流，基于张量。 func.func @matmul(%A: tensor<128x256xf32>, %B: tensor<256x64xf32>) -> tensor<128x64xf32> { %zero = arith.constant 0.0 : f32 // 1. 累加器的标量单位元 %init = tensor.empty() : tensor<128x64xf32> // 2. 未初始化的结果值 %acc = linalg.fill ins(%zero : f32) outs(%init : tensor<128x64xf32>) -> tensor<128x64xf32> // 3. 将累加器初始化为零 %C = linalg.matmul // 4. 整个乘法，作为一个操作 ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>) outs(%acc : tensor<128x64xf32>) -> tensor<128x64xf32> return %C : tensor<128x64xf32> } ``` 两个变换将其转化为接近机器码的形式。首先，**缓冲化**（bufferization）将值语义的 `tensor` 转换为 `memref`，即带有地址、布局和生命周期的显式缓冲区。[^5] 这是程序不再纯数据流、开始具有别名、分配以及需要将结果写入调用者提供的输出缓冲区（而非返回新值）的时刻。 其次，结构化的 `linalg.matmul` 被降层为显式的循环嵌套： ``` // 张量现在变成了 memref；结果就地写入 %C。 func.func @matmul(%A: memref<128x256xf32>, %B: memref<256x64xf32>, %C: memref<128x64xf32>) { affine.for %i = 0 to 128 { // 1. matmul 隐含的循环嵌套 affine.for %j = 0 to 64 { affine.for %k = 0 to 256 { %a = affine.load %A[%i, %k] : memref<128x256xf32> // 2. 从缓冲区显式加载 %b = affine.load %B[%k, %j] : memref<256x64xf32> %c = affine.load %C[%i, %j] : memref<128x64xf32> %p = arith.mulf %a, %b : f32 // 3. 算术运算，现在是标量 %s = arith.addf %c, %p : f32 affine.store %s, %C[%i, %j] : memref<128x64xf32> // 4. 累加回缓冲区 } } } return } ``` 从这里开始，`affine`/`scf` 循环可以被分块（tiling）、向量化到 `vector` 方言，最后降层到 `llvm` 方言（面向 CPU）或 `gpu`/`nvvm` 路径。本质上，**相同**的高层出发点可以通过选择不同的降层 pass 被驱动到任意这些目标路径上。这种重用——一个前端表示、多个目标降层——就是 MLIR 的全部经济论据。 ## 真正的好处：可重用的降层机制 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#the-real-payoff-reusable-lowering-machinery 使这一切变得实用的不是方言本身，而是它们下方的共享基础设施。MLIR 提供了一个**方言转换**（dialect conversion）框架：你声明一个转换目标（哪些方言/操作在最终是“合法的”），提供将非法操作转换为合法操作的重写模式，然后框架将它们应用到不动点，同时处理类型转换和操作数重映射。[^2] 模式重写、验证器、位置/调试信息传播以及 pass 管理器都是独立于方言的。大多数变换被编写为局部重写模式，而不是手写 IR 遍历，包括每个方言注册的规范化模式（用来折叠和简化自身的操作）；这些模式可以通过声明式方式（通过 PDL/DRR）定义，并由共享驱动程序应用。[^6] 一个新的抽象就是一组操作定义加上一组模式；其他一切都是继承来的。这就是为什么生态系统得以整合。每个项目不再维护一个定制的编译器中间端，而是共享一个： ```mermaid graph LR MLIR["MLIR 核心\n操作 · 方言 · pass · 重写器"] MLIR --- XLA["XLA / StableHLO"] MLIR --- TRITON["Triton"] MLIR --- MOJO["Mojo"] MLIR --- IREE["IREE"] MLIR --- TORCH["Torch-MLIR"] MLIR --- CIRCT["CIRCT (硬件设计)"] classDef core fill:#e1f5ff,stroke:#4a6572,color:#1a1a1a; classDef leaf fill:#f3e5f5,stroke:#6a4a6a,color:#1a1a1a; class MLIR core; class XLA,TRITON,MOJO,IREE,TORCH,CIRCT leaf; ``` 这也是新硬件获得支持的方式。厂商添加一个建模其设备操作和内存的方言，再加上将标准中层方言（`linalg`、`memref`、`vector`）降层到该方言的 pass。前端、优化器和工具链都是免费的；厂商只编写真正针对其硅片特性的部分。[^7] ## 成本：验证、调试和脆弱的流水线 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#the-costs-verification-debugging-and-fragile-pipelines 使得 MLIR 高效工作的机制也有明显的尖锐边缘，值得坦诚说明。 **定义正确的验证是真正的工作。** 每个方言都提供一个验证器，用以强制执行其操作和类型的不变式，共享框架会将它们全部运行起来。这能及早捕获格式错误的 IR，但编写*正确*的验证规则的责任落到了定义方言的人身上。如果一个类型不变式写错了，或者遗漏了一个，那么格式错误的 IR 就会通过验证器，并在后续的 pass 中作为崩溃重新浮现，距离其真正引入的位置很远。[^8] **调试跨越海拔。** 由于一个模块中有多个方言，并有一长串 pass，一个错误编译或性能回归必须从高层张量操作一直追溯到特定的 LLVM 或 PTX 指令。位置追踪会在降层过程中传播源代码位置并有所帮助，但 bug 常常在导致它的 pass 之后的多个 pass 后才显现。实用的工具包括在 pass 之间转储 IR（`mlir-opt --mlir-print-ir-after-all`）和自动崩溃复现器；即使有了这些，多级调试本身也是一项技能，并且在 MLIR 社区中是一个反复出现的话题。[^9] **方言是稳定的；流水线不是。** StableHLO 是一个版本化的、可移植的中间交换格式，内建方言也相当稳定。但将变换串接在一起的 **pass 流水线** 并不是。为 Torch-MLIR 调优的降层流水线一般不能在 IREE 中未经修改地运行，而一个在某个版本上能工作的序列可能在下一个版本上就坏了。“构建在 MLIR 之上”意味着方言可以组合；这并不意味着两个基于 MLIR 的栈可以端到端地互操作。[^10] ## 关键技术要点 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#key-technical-takeaways | 概念 | 是什么 | 为什么重要 | |------|--------|------------| | **方言** | 操作和类型的名称空间（例如 `linalg`、`arith`）。 | 允许高层和低层操作在一个模块中混合，无需阻抗不匹配。 | | **区域** | 操作可以包含嵌套的进一步操作的基本块。 | 实现多级 IR，其中一个操作可以持有整个子计算（结构化操作内部的循环体）。 | | **缓冲化** | 将 `tensor`（值语义）转换为 `memref`（缓冲区）。 | 程序获得别名、分配和副作用的关键转换。 | | **方言转换** | 将非法操作重写为合法操作直至不动点的框架。 | 驱动降层的引擎，例如 `linalg` → `vector` → `llvm`。 | ## 基础设施的边界 https://hiraditya.github.io/posts/mlir-dialect-stack-for-ml/#where-the-infrastructure-stops MLIR 有意只提供管道。它为你提供一个*表示*硬件的地方，以及一种*降层*到硬件的规范方式，这已经相当不错了。但它不是源语言，它本身也不决定计算应该在何处运行，也不保证缓冲区存在于操作所期望的内存空间中。这些选择由位于其上层的部分决定：前端语言、pass 流水线、代价模型。如今，这些大多仍然在类型系统外部表达。