216 lines
9.5 KiB
Markdown
216 lines
9.5 KiB
Markdown
# 阶段 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 必会）
|
||
|
||
```mlir
|
||
%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` 是 `arith` dialect 里的 op，名字用 `arith.` 前缀
|
||
- `linalg.matmul` 是 `linalg` dialect 里的 op
|
||
- `flow.dispatch.region` 是 IREE 的 `flow` dialect 里的 op
|
||
|
||
为什么要分 dialect？因为**不同抽象层次需要不同的"语言"**：
|
||
- 高层用 `linalg`（表达"这是一个矩阵乘"的语义）
|
||
- 低层用 `llvm`（表达"这是 fmuladd + 循环 + 指针"）
|
||
- 图级用 `flow`、异步用 `stream`、设备用 `hal`
|
||
|
||
编译过程就是**从高层 dialect 逐层 lowering 到低层 dialect**（这正是阶段 2 看到的 7 个阶段）。
|
||
|
||
## 4. 基础 dialect（IR 的积木）
|
||
|
||
### `func` —— 函数
|
||
```mlir
|
||
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` —— 算术/比较/转换
|
||
```mlir
|
||
%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）
|
||
```mlir
|
||
%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` —— 非结构化控制流（底层分支）
|
||
```mlir
|
||
cf.br ^bb1(%0 : i32) // 无条件跳转
|
||
cf.cond_br %cond, ^bb1(%0 : i32), ^bb2 // 条件跳转
|
||
```
|
||
- 最低层的控制流，`scf.for` 最终会 lower 成 `cf.br` 循环
|
||
|
||
### `affine` —— 仿射循环与访存
|
||
```mlir
|
||
%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` —— 不可变张量（值语义）
|
||
```mlir
|
||
%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）
|
||
```mlir
|
||
%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 向量
|
||
```mlir
|
||
%0 = vector.load %memref[%i, %j] : memref<2x2xf32>, vector<4xf32>
|
||
%1 = vector.transfer_read %t[0, 0], %pad : tensor<4x4xf32>, vector<4xf32>
|
||
```
|
||
- 固定大小的向量，用于表达 SIMD 指令（IREE 向量化后产生 `vector` op）
|
||
|
||
### `linalg` —— 结构化计算（**IREE 的心跳**）
|
||
```mlir
|
||
// 命名 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。
|
||
|
||
### 常用观察工具
|
||
|
||
```bash
|
||
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`（实现）。
|