92 lines
4.8 KiB
Markdown
92 lines
4.8 KiB
Markdown
# 阶段 4a：Flow（图级优化 + dispatch 划分）✅
|
||
|
||
> 2026-08-15。对应规划阶段 4 前半。源码：`src/iree/compiler/src/iree/compiler/Dialect/Flow/`。
|
||
> IR 变换示例见 [02-architecture.md](02-architecture.md) 边界 2。
|
||
|
||
## 1. Flow 的职责
|
||
|
||
Flow 是 IREE 的**图级 IR**，在输入方言（linalg）和 Stream（异步）之间。两大职责：
|
||
|
||
1. **图级优化**：算子融合（fusion）、reshape 折叠、类型推导、量化转换、全局常量处理等
|
||
2. **dispatch 划分**：把图切成一个个 **dispatch region**，再 outline 成独立的 `flow.executable`
|
||
—— 每个 dispatch = 一个可独立编译执行的 kernel（对应 TVM 里划分出的一个 kernel）
|
||
|
||
## 2. 核心 op（定义在 `IR/FlowOps.td`）
|
||
|
||
| op | 作用 |
|
||
|---|---|
|
||
| `flow.tensor.empty/load/store/slice/reshape/update` | Flow 层的张量操作（值语义） |
|
||
| `flow.tensor.barrier` | 张量屏障，标记 dispatch 边界（前身是 `iree_tensor_ext.compute_barrier`） |
|
||
| `flow.tensor.transfer` | 跨设备张量传输 |
|
||
| `flow.dispatch.region` | **dispatch region**：一个可执行单元（核心） |
|
||
| `flow.dispatch` / `flow.dispatch.workgroups` | 调用 dispatch（启动一个 kernel） |
|
||
| `flow.executable` / `flow.executable.export` | 可执行体 + 导出点 |
|
||
| `flow.dispatch.workgroup_count/id/size` | workgroup 相关内建 |
|
||
| `flow.channel.*` + `flow.collective.*` | 多设备通信（AllGather/AllReduce/AllToAll/SendRecv） |
|
||
| `flow.parameter.load/write` | 参数（权重）处理 |
|
||
|
||
## 3. 关键 pass（`Transforms/` 目录）
|
||
|
||
| pass | 文件 | 作用 |
|
||
|---|---|---|
|
||
| `ConvertToFlowPass` | ConvertToFlow.cpp | 输入方言 → Flow |
|
||
| `FusionPreprocessingPass` / `ElementwiseOpFusionPass` / `FuseMultiUseElementwiseProducerPass` | （dispatch creation） | **算子融合**（elementwise 融合进 reduce/matmul） |
|
||
| `FormDispatchRegionsPass` | FormDispatchRegions.cpp | **形成 dispatch region**（核心） |
|
||
| `CloneProducersIntoDispatchRegionsPass` | （dispatch creation） | 把 producer 克隆进 dispatch region |
|
||
| `ConvertDispatchRegionsToWorkgroupsPass` | ConvertRegionToWorkgroups.cpp | dispatch region → workgroup 语义 |
|
||
| `OutlineDispatchRegionsPass` | OutlineDispatchRegions.cpp | **outline 成 flow.executable** |
|
||
| `DeduplicateExecutablesPass` | DeduplicateExecutables.cpp | 去重相同 executable |
|
||
| `CleanupTensorShapesPass` / `CaptureDynamicDimsPass` | 同名 .cpp | 形状清理 / 动态维度捕获 |
|
||
|
||
## 4. dispatch region 形成（核心概念）
|
||
|
||
**dispatch region = 一个要交给 codegen 独立编译的计算单元**（对应一个 kernel 启动）。
|
||
|
||
形成逻辑（`FormDispatchRegionsPass`）：
|
||
1. 以「root op」（如 `linalg.matmul`、reduction）为锚点
|
||
2. 把可融合的 elementwise producer 克隆进同一个 region（fusion）
|
||
3. 用 `flow.tensor.barrier` 标记 region 的输入/输出边界
|
||
4. 后续 outline 成 `flow.executable`
|
||
|
||
IR 变化（阶段 2 已见，精简）：
|
||
|
||
```mlir
|
||
// FormDispatchRegionsPass 之后：matmul + fill 被包进 dispatch.region
|
||
%6 = flow.dispatch.region -> (tensor<2x2xf32>) {
|
||
%8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<2x2xf32>) -> tensor<2x2xf32>
|
||
%9 = linalg.matmul ins(%1, %3 : ...) outs(%8 : ...) -> tensor<2x2xf32>
|
||
flow.return %9 : tensor<2x2xf32>
|
||
}
|
||
|
||
// OutlineDispatchRegionsPass 之后：outline 成 executable，主函数只剩 flow.dispatch 调用
|
||
flow.executable private @matmul_dispatch_0 { ... } // 独立的可执行体
|
||
%2 = flow.dispatch @matmul_dispatch_0::@matmul_dispatch_0_matmul_2x2x4_f32(%0, %1) : ...
|
||
```
|
||
|
||
## 5. 与 TVM 0.6 对照
|
||
|
||
| IREE Flow | TVM 0.6 Relay |
|
||
|---|---|
|
||
| 图级优化（fusion/reshape/type propagation） | Relay 的 fusion / 图优化 pass |
|
||
| dispatch region 形成 | 确定哪些算子合成一个 kernel（relay fusion 后的分组） |
|
||
| flow.executable | 一个编译好的 kernel |
|
||
| flow.tensor.*（值语义张量） | Relay 的 TensorType 表达式 |
|
||
| flow.collective.* | Relay 里的通信算子（较晚版本才有） |
|
||
|
||
## 6. 源码阅读路径
|
||
|
||
```
|
||
src/iree/compiler/src/iree/compiler/Dialect/Flow/
|
||
├── IR/FlowOps.td # op 定义（TableGen）—— 先读这个，认识所有 op
|
||
├── IR/FlowOps.cpp # op 实现（verifier、folder、builder）
|
||
├── IR/FlowInterfaces.td # 接口定义（如 Dispatchable/Partitionable）
|
||
├── Transforms/FormDispatchRegions.cpp # dispatch 划分核心算法
|
||
├── Transforms/OutlineDispatchRegions.cpp
|
||
├── Transforms/ConvertRegionToWorkgroups.cpp
|
||
└── Conversion/TensorToFlow/ # tensor → flow 转换
|
||
```
|
||
|
||
> 学习建议：先读 `FlowOps.td` 认 op，再看 `FormDispatchRegions.cpp` 理解划分算法。
|
||
> 划分算法在 IREE 3.x 已重构为基于「partitionable interface + workload」的方式
|
||
> （不是老的 DispatchableOpInterface），读的时候注意版本。
|