4.8 KiB
4.8 KiB
阶段 4a:Flow(图级优化 + dispatch 划分)✅
2026-08-15。对应规划阶段 4 前半。源码:
src/iree/compiler/src/iree/compiler/Dialect/Flow/。 IR 变换示例见 02-architecture.md 边界 2。
1. Flow 的职责
Flow 是 IREE 的图级 IR,在输入方言(linalg)和 Stream(异步)之间。两大职责:
- 图级优化:算子融合(fusion)、reshape 折叠、类型推导、量化转换、全局常量处理等
- 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):
- 以「root op」(如
linalg.matmul、reduction)为锚点 - 把可融合的 elementwise producer 克隆进同一个 region(fusion)
- 用
flow.tensor.barrier标记 region 的输入/输出边界 - 后续 outline 成
flow.executable
IR 变化(阶段 2 已见,精简):
// 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),读的时候注意版本。