175 lines
8.7 KiB
Markdown
175 lines
8.7 KiB
Markdown
# 阶段 2：IREE 架构总览与完整管线 ✅
|
||
|
||
> 2026-08-14。用 `--mlir-print-ir-after-all` 追踪一个 linalg matmul 从输入到 VM 的完整变换，
|
||
> 把 231 个 pass 划分成 7 个阶段，建立 IREE 整体的心智模型。
|
||
|
||
## 1. 一句话总结
|
||
|
||
IREE 是一个**「图 → 异步调度 → 设备代码 → 字节码」的完整编译栈**，每一层都用 MLIR dialect 表达：
|
||
输入方言（linalg/tosa/stablehlo）→ Flow → Stream → HAL → VM，其中 codegen（可执行体 lowering）
|
||
在 Stream 和 HAL 之间发生。
|
||
|
||
## 2. 完整 pipeline 全景（231 个 pass → 7 个阶段）
|
||
|
||
| 阶段 | 关键 dialect | 代表性 pass | 职责 | TVM 对应 |
|
||
|---|---|---|---|---|
|
||
| A 输入 | (输入方言) + util + hal.tensor | `WrapEntryPointsPass`、`InlinerPass` | 统一入口、包 ABI（buffer_view） | 前端导入 |
|
||
| B 设备 | hal.device.target | `AssignLegacyTargetDevicesPass`、`MaterializeTargetDevicesPass` | 目标设备分配 | device plan |
|
||
| C Flow | **flow** | `FormDispatchRegionsPass`、`OutlineDispatchRegionsPass` | 图级优化、fusion、dispatch 划分 | Relay 图优化 + fusion |
|
||
| D Stream | **stream** | `ConvertToStreamPass`、`ScheduleExecutionPass` | 异步调度（timepoint/fence） | （无，IREE 特色） |
|
||
| E Codegen | linalg/vector/llvm | `TileAndFuse`、`GenericVectorization`、`Bufferize`、`ConvertToLLVM` | 可执行体 lowering 到设备代码 | TE/TIR + schedule + codegen |
|
||
| F HAL | **hal** | `ConvertToHALPass`、`SerializeAllExecutablesPass` | 设备抽象、可执行体序列化 | runtime device 抽象 |
|
||
| G VM | **vm** | （尾部 pass + 序列化） | 字节码序列化 | VM |
|
||
|
||
> 完整 231 个 pass 的去重序列见 [examples/matmul_pipeline_dump.txt](../../examples/matmul_pipeline_dump.txt)
|
||
> （32702 行完整 dump，可用 `grep 'IR Dump After'` 提取 pass 序列）。
|
||
|
||
## 3. 关键 IR 变换（每个 dialect 边界）
|
||
|
||
### 边界 1：输入 → ABI（WrapEntryPointsPass）
|
||
用户写的 `func(tensor)` 被包成对外 `util.func public`，参数变成 `!hal.buffer_view`，
|
||
内部用 `hal.tensor.import` / `hal.tensor.export` 在 buffer_view 和 tensor 之间转换：
|
||
|
||
```mlir
|
||
util.func public @matmul(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view {
|
||
%0 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<2x4xf32>
|
||
%1 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<4x2xf32>
|
||
%2 = util.call @_matmul(%0, %1) : (...) -> tensor<2x2xf32>
|
||
%3 = hal.tensor.export %2 "output0" : tensor<2x2xf32> -> !hal.buffer_view
|
||
util.return %3 : !hal.buffer_view
|
||
}
|
||
```
|
||
|
||
### 边界 2：Flow dispatch 划分（FormDispatchRegionsPass → OutlineDispatchRegionsPass）
|
||
`linalg.matmul` 被包进 `flow.dispatch.region`（→ 相当于 TVM 里划分 kernel），
|
||
再 outline 成独立的 `flow.executable`，主函数里只剩一个 `flow.dispatch` 调用：
|
||
|
||
```mlir
|
||
// FormDispatchRegionsPass 之后：
|
||
%6 = flow.dispatch.region -> (tensor<2x2xf32>) {
|
||
%9 = linalg.matmul ins(%1, %3 : ...) outs(%8 : ...) -> tensor<2x2xf32>
|
||
flow.return %9 : tensor<2x2xf32>
|
||
}
|
||
|
||
// OutlineDispatchRegionsPass 之后：
|
||
flow.executable private @matmul_dispatch_0 {
|
||
flow.executable.export public @matmul_dispatch_0_matmul_2x2x4_f32 workgroups() -> (index, index, index) { ... }
|
||
builtin.module {
|
||
func.func @matmul_dispatch_0_matmul_2x2x4_f32(%arg0: !iree_tensor_ext.dispatch.tensor<...>, ...) {
|
||
// 可执行体内部：linalg 计算 + dispatch.tensor.load/store
|
||
}
|
||
}
|
||
}
|
||
util.func public @matmul(...) {
|
||
%2 = flow.dispatch @matmul_dispatch_0::@matmul_dispatch_0_matmul_2x2x4_f32(%0, %1) : (...) -> tensor<2x2xf32>
|
||
}
|
||
```
|
||
|
||
### 边界 3：Flow → Stream（ConvertToStreamPass）
|
||
可执行体变 `stream.executable`，绑定变 `!stream.binding`；主函数变成异步资源流：
|
||
|
||
```mlir
|
||
// 可执行体函数签名：!stream.binding（抽象设备缓冲区）
|
||
func.func @matmul_dispatch_0_matmul_2x2x4_f32(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: !stream.binding)
|
||
|
||
// 主函数：tensor → stream.resource + 设备亲和性标注
|
||
%1 = stream.tensor.import on(#hal.device.affinity<@__device_0>) %arg0 : !hal.buffer_view -> tensor<2x4xf32> in !stream.resource<external>{%0}
|
||
%7 = stream.tensor.dispatch on(#hal.device.affinity<@__device_0>) @matmul_dispatch_0::@matmul_dispatch_0_matmul_2x2x4_f32(%2, %5) : (...) -> tensor<2x2xf32> in !stream.resource<*>{%6}
|
||
```
|
||
|
||
### 边界 4：Stream 异步调度（ScheduleExecutionPass）
|
||
dispatch 被包进 `stream.async.execute`，返回 `!stream.timepoint`，再 `await`：
|
||
|
||
```mlir
|
||
%results, %result_timepoint = stream.async.execute on(...) with(%0 as %arg2, %1 as %arg3) -> !stream.resource<...> {
|
||
%4 = stream.async.dispatch @matmul_dispatch_0::@matmul_dispatch_0_matmul_2x2x4_f32(%arg2[...], %arg3[...]) : ...
|
||
stream.yield %4 : !stream.resource<...>
|
||
} => !stream.timepoint
|
||
%2 = stream.timepoint.await %result_timepoint => %results : !stream.resource<...>
|
||
```
|
||
|
||
> `timepoint` 是 IREE 异步模型的核心：类似 future/fence，表示「某批异步工作何时完成」，
|
||
> 用来表达数据依赖和跨设备同步。这是 TVM 没有的层。
|
||
|
||
### 边界 5：Codegen（可执行体 lowering，ConvertToLLVMPass）
|
||
可执行体里的 linalg matmul 被 tiling → vectorization → bufferize → LLVM：
|
||
|
||
```mlir
|
||
// 最终：显式循环 + fmuladd + 向量 load
|
||
llvm.func @matmul_dispatch_0_matmul_2x2x4_f32(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) -> i32 {
|
||
// 从 iree_hal_executable_dispatch_state_v0_t 取 bindings
|
||
// ^bb1/^bb2/^bb3 显式循环，llvm.intr.fmuladd 做 MAC，vector<1xf32> 向量化
|
||
}
|
||
```
|
||
|
||
> 注意：codegen 在 Stream 和 HAL 之间发生（pass 130~199），因为可执行体要先编译成设备代码，
|
||
> 才能在 HAL 层被序列化。关键 pass：`LLVMCPUSelectLoweringStrategyPass`（选 tiling/向量化策略）、
|
||
> `TileAndDistributeToWorkgroups`、`LLVMCPUTileAndFuse`、`GenericVectorization`、
|
||
> `IREEComprehensiveBufferize`、`ConvertToLLVM`。
|
||
|
||
### 边界 6：Stream → HAL（ConvertToHALPass）
|
||
可执行体包成 `hal.executable.variant`，带 pipeline layout（缓冲区绑定）和 ordinal：
|
||
|
||
```mlir
|
||
hal.executable private @matmul_dispatch_0 {
|
||
hal.executable.variant public @embedded_elf_x86_64 target(#executable_target_embedded_elf_x86_64) {
|
||
hal.executable.export public @matmul_dispatch_0_matmul_2x2x4_f32 ordinal(0) layout(#pipeline_layout)
|
||
attributes {workgroup_size = [1 : index, 1 : index, 1 : index]}
|
||
builtin.module { llvm.func @matmul_dispatch_0_matmul_2x2x4_f32(...) { ... } }
|
||
}
|
||
}
|
||
// #pipeline_layout 定义了 3 个 storage_buffer 绑定（2 ReadOnly + 1 输出）
|
||
```
|
||
|
||
### 边界 7：最终 VM（字节码模块）
|
||
序列化后的 VM 模块：声明类型、依赖 `hal` 模块、import 一堆 `hal.*` 函数：
|
||
|
||
```
|
||
Required Types: i32, i64, !hal.allocator, !hal.buffer, !hal.buffer_view,
|
||
!hal.command_buffer, !hal.device, !hal.executable, !hal.fence, !vm.buffer
|
||
Imported Functions:
|
||
[ 7] hal.command_buffer.dispatch(...) -> ()
|
||
[11] hal.device.queue.execute.indirect(...) -> ()
|
||
[16] hal.fence.await(i32, i64, tuple<!vm.ref<?>>...) -> (i32)
|
||
...
|
||
```
|
||
|
||
运行时 `iree-runtime` 加载这个 VM 模块、解析 hal 模块 import、用 HAL driver 执行。
|
||
|
||
## 4. TVM 0.6 对照总结
|
||
|
||
| IREE 层 | 作用 | TVM 0.6 对应 | 差异点 |
|
||
|---|---|---|---|
|
||
| Flow | 图级优化 + dispatch 划分 | Relay（fusion、算子划分） | 划分出的 dispatch = 一个 kernel |
|
||
| Stream | 异步调度 | 无 | timepoint/fence 显式建模异步 |
|
||
| Codegen | 可执行体 lowering | TE/TIR + schedule + codegen | MLIR dialect lowering，非 schedule 原语 |
|
||
| HAL | 设备抽象 | runtime device + codegen 的 device 层 | 更抽象，统一 CPU/GPU/Vulkan |
|
||
| VM | 字节码 | graph runtime / 早期 VM | 紧凑字节码 + 模块依赖 |
|
||
|
||
## 5. 复现命令
|
||
|
||
```bash
|
||
# 完整 dump（含每个 pass 的 IR）
|
||
iree-compile --iree-hal-target-backends=llvm-cpu --mlir-print-ir-after-all \
|
||
examples/matmul.mlir -o /tmp/matmul.vmfb 2> dump.txt
|
||
|
||
# 只看 pass 序列
|
||
grep -oE 'IR Dump After .*//' dump.txt | sed 's/IR Dump After //; s/ (.*//; s/ \/\/$//' \
|
||
| awk '!seen[$0]++' | nl
|
||
|
||
# 只看某个 pass 之后的 IR
|
||
iree-compile --iree-hal-target-backends=llvm-cpu \
|
||
--mlir-print-ir-after=ConvertToStreamPass --mlir-print-ir-module-scope \
|
||
examples/matmul.mlir -o /tmp/matmul.vmfb 2>&1 | head -80
|
||
|
||
# 看最终 VM 模块
|
||
iree-dump-module /tmp/matmul.vmfb
|
||
```
|
||
|
||
## 6. 下一步
|
||
|
||
- 阶段 1（MLIR 基础）可选补：现在读懂了 dialect 转换，回头看 Operation/Value/Pass/PatternRewriting
|
||
会很快，可先快速过一遍再深入各层。
|
||
- 阶段 3（前端）、4（Flow+Stream）、5（HAL）、6（codegen）、7（运行时）逐层深入，
|
||
每一层都对着 `src/iree/compiler/src/iree/compiler/Dialect/<X>/` 的源码读。
|