9.5 KiB
阶段 1:MLIR 速查(对应规划阶段 1)
2026-08-14。面向「已有 TVM 经验、看过 IREE pipeline」的 MLIR 速查。 重点是 dialect 概念 + 各 dialect 职责,配合阶段 2 的 pipeline dump 一起看。
1. MLIR 的核心思想
MLIR 不是一个固定 IR,而是一个 IR 框架:它定义了一套通用的 IR 结构(Operation/Block/Region/Type), 然后让你自己定义 dialect。每个 dialect 定义自己的 op、类型、属性。
- builtin:最核心的部分(module、基础类型)
- dialect:一组相关的 op/type/attr 的集合(如
arith、linalg、flow) - op:具体操作(如
arith.addf、linalg.matmul)
IREE 的每一层都是一个 dialect:Flow、Stream、HAL、VM 都是自定义的 MLIR dialect。
2. IR 结构核心概念
| 概念 | 是什么 | 类比 TVM |
|---|---|---|
| Operation (op) | 基本单位:名字 + operands + results + attributes + regions | IR 节点 / expr |
| Value | SSA 值,只能赋值一次;是 op 的 result 或 block 的 argument | 变量/临时值 |
| Block | 顺序执行的一串 op,可带 block argument,以 terminator 结尾 | 基本块 |
| Region | 一个或多个 block 的列表,支持嵌套(实现控制流) | 子图/作用域 |
| Type | 值的类型:i32 f32 tensor<2x4xf32> memref<...> vector<...> |
dtype |
| Attribute | 编译期常量/元数据:arith.constant 5 : i32 里的 5,或 op 的 {...} |
常量/annotations |
| Symbol | 命名实体(函数名、全局名),通过符号引用 | 函数名 |
Op 的通用语法(重点,读 IR 必会)
%result = linalg.matmul ins(%lhs, %rhs : tensor<2x4xf32>, tensor<4x2xf32>)
outs(%init : tensor<2x2xf32>) -> tensor<2x2xf32>
%result= 结果 SSA 值;-> tensor<2x2xf32>= 结果类型ins/outs= 输入 / 输出操作数(linalg 特有的命名,其他 op 直接用位置参数)(%lhs, %rhs : type1, type2)= 操作数及其类型
3. 什么是 Dialect?
Dialect = 一组相关的 op、类型、属性的集合,定义在同一个命名空间下。例如:
arith.addf是arithdialect 里的 op,名字用arith.前缀linalg.matmul是linalgdialect 里的 opflow.dispatch.region是 IREE 的flowdialect 里的 op
为什么要分 dialect?因为不同抽象层次需要不同的"语言":
- 高层用
linalg(表达"这是一个矩阵乘"的语义) - 低层用
llvm(表达"这是 fmuladd + 循环 + 指针") - 图级用
flow、异步用stream、设备用hal
编译过程就是从高层 dialect 逐层 lowering 到低层 dialect(这正是阶段 2 看到的 7 个阶段)。
4. 基础 dialect(IR 的积木)
func —— 函数
func.func @matmul(%arg0: tensor<2x4xf32>) -> tensor<2x2xf32> {
%0 = func.call @helper(%arg0) : (tensor<2x4xf32>) -> tensor<2x2xf32>
func.return %0 : tensor<2x2xf32>
}
- 定义函数
func.func、调用func.call、返回func.return
arith —— 算术/比较/转换
%0 = arith.constant 5 : i32
%1 = arith.addf %a, %b : f32 // 浮点加
%2 = arith.addi %x, %y : i32 // 整数加
%3 = arith.cmpf ogt, %a, %b : f32 // 浮点比较
%4 = arith.mulf %a, %b : f32 // 浮点乘
- 标量算术的"标准库",几乎所有计算最终都落到这里
scf —— 结构化控制流(for / if / while)
%0 = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> (tensor<...>) {
%1 = arith.addi %i, %acc : i32
scf.yield %1 : i32
}
%2 = scf.if %cond -> (f32) {
scf.yield %a : f32
} else {
scf.yield %b : f32
}
- "结构化"= 带明确的循环语义,便于分析和变换(IREE tiling 产生的就是
scf.for)
cf —— 非结构化控制流(底层分支)
cf.br ^bb1(%0 : i32) // 无条件跳转
cf.cond_br %cond, ^bb1(%0 : i32), ^bb2 // 条件跳转
- 最低层的控制流,
scf.for最终会 lower 成cf.br循环
affine —— 仿射循环与访存
%0 = affine.for %i = 0 to 10 {
%1 = affine.load %memref[%i, %i + 1] : memref<10x10xf32>
}
- 循环边界和下标必须是仿射表达式(
i、i+1、2*i+3),可做 polyhedral 分析 - IREE 历史上有用,现在更多用 scf + linalg
index 类型
index是"可移植的整数",循环变量、下标、shape 都用它- 编译到具体目标时再决定是 i32 还是 i64
5. 数据/计算 dialect(IREE codegen 的核心)
tensor —— 不可变张量(值语义)
%0 = tensor.empty() : tensor<2x2xf32> // 建空张量
%1 = tensor.extract_slice %t[0, 0][2, 2][1, 1] : tensor<4x4xf32> to tensor<2x2xf32>
%2 = tensor.insert_slice %s into %t[0, 0][2, 2][1, 1] : tensor<2x2xf32> into tensor<4x4xf32>
%3 = tensor.extract %t[0, 0] : tensor<2x2xf32> // 取标量
- 值语义:不可变,像纯函数式。编译期的高层抽象
memref —— 内存缓冲区(可变的 buffer)
%0 = memref.alloc() : memref<2x2xf32>
%1 = memref.load %0[%i, %j] : memref<2x2xf32> // 读
memref.store %v, %0[%i, %j] : memref<2x2xf32> // 写
- 可变,带 layout/stride,对应"真实内存"。bufferize 后 tensor 变 memref
vector —— 固定大小 SIMD 向量
%0 = vector.load %memref[%i, %j] : memref<2x2xf32>, vector<4xf32>
%1 = vector.transfer_read %t[0, 0], %pad : tensor<4x4xf32>, vector<4xf32>
- 固定大小的向量,用于表达 SIMD 指令(IREE 向量化后产生
vectorop)
linalg —— 结构化计算(IREE 的心跳)
// 命名 op:语义明确的高层计算
%0 = linalg.matmul ins(%lhs, %rhs : tensor<2x4xf32>, tensor<4x2xf32>)
outs(%init : tensor<2x2xf32>) -> tensor<2x2xf32>
%1 = linalg.fill ins(%cst : f32) outs(%empty : tensor<2x2xf32>) -> tensor<2x2xf32>
// generic op:自定义的"结构化计算",用 region 定义元素计算 + 迭代空间
%2 = linalg.generic {
indexing_maps = [...],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%a : tensor<...>) outs(%b : tensor<...>) {
^bb0(%in: f32, %out: f32):
%sum = arith.addf %in, %out : f32
linalg.yield %sum : f32
}
- 关键思想:linalg 把"计算"和"数据如何存/遍历"分离。 它只声明"这是矩阵乘"或"元素级计算 + 迭代空间",具体 tiling/向量化由后续 pass 做。 这就是 IREE codegen 的起点(对应 TVM 的 TE 层,但声明式更强)。
6. 易混淆概念对照(重点)
| 对比 | 区别 |
|---|---|
tensor vs memref |
值语义(不可变)vs 内存(可变)。编译期用 tensor,bufferize 后变 memref |
scf.for vs affine.for |
结构化(可标并行/归约)vs 仿射(可 polyhedral 分析) |
arith vs math |
基本算术(add/mul/cmp)vs 数学函数(math.sqrt、math.exp、math.absf) |
tensor.empty vs memref.alloc |
抽象空张量(值)vs 真实内存分配 |
7. Pass 基础设施(看懂 pipeline 的关键)
Pass = 遍历 IR 做变换的基本单位。IREE 那 731 次 pass 就是一个个变换。
两种主要机制(深入各层源码必懂):
(1) Pattern Rewriting(局部重写)
声明式:写"匹配 pattern → 替换成新 op"的规则,引擎(如 applyPatternsAndFoldGreedily)自动遍历应用。
适合局部的、等价的化简。
(2) Dialect Conversion(跨 dialect 转换)
三件套:ConversionTarget(哪些 op 合法/非法)+ TypeConverter(类型怎么换)+
patterns(非法 op 怎么重写成合法 op)。
适合把整个 dialect 换成另一个(如 Flow→Stream、tensor→memref)。
你在 pipeline 里看到的
ConvertToStreamPass、ConvertToLLVMPass、IREEComprehensiveBufferizePass都是 Dialect Conversion;CanonicalizerPass是 Pattern Rewriting。
常用观察工具
iree-compile --mlir-print-ir-after-all in.mlir -o out.vmfb 2> dump.txt # 每个 pass 后 dump
iree-compile --mlir-print-ir-after=ConvertToStreamPass --mlir-print-ir-module-scope in.mlir -o out.vmfb 2>&1 # 只看某个 pass
8. 与 TVM 0.6 对照
| MLIR | TVM 0.6 |
|---|---|
| Operation | IR node / expr |
| Pass | pass / IR visitor |
| Dialect Conversion | 逐层 lowering(Relay→TE→TIR) |
| Pattern Rewriting | pass 里的 rewrite 规则 |
| linalg.matmul/generic | TE 的 compute + 描述 |
| scf.for + vector | TIR 里的 loop + vectorize |
| memref | 缓冲区/物理内存 |
9. IREE 自己的 dialect(概览,详细见各层笔记)
| dialect | 职责 | 关键 op(在 dump 里见过) |
|---|---|---|
util |
通用工具(函数/全局包装) | util.func、util.global、util.call、util.return |
flow |
图级优化 + dispatch 划分 | flow.dispatch.region、flow.dispatch、flow.executable |
stream |
异步执行模型 | stream.executable、stream.binding、stream.tensor.import、stream.async.execute、stream.timepoint.await |
hal |
设备抽象 | hal.executable、hal.buffer_view、hal.command_buffer、hal.tensor.import/export |
vm |
VM 字节码 | vm.module、vm.func、vm.call |
iree_tensor_ext |
dispatch 张量 + 屏障 | iree_tensor_ext.dispatch.tensor.load/store、iree_tensor_ext.compute_barrier.start/end |
iree_linalg_ext |
扩展 linalg op(softmax/attention 等) | 复杂算子的结构化 op |
这些 dialect 的源码在
src/iree/compiler/src/iree/compiler/Dialect/<X>/, 每个 dialect 目录下有*.td(op 定义,TableGen 格式)和*.cpp(实现)。