127 lines
6.3 KiB
Markdown
127 lines
6.3 KiB
Markdown
# 阶段 4b：Stream（异步调度）✅
|
||
|
||
> 2026-08-15。对应规划阶段 4 后半。源码：`src/iree/compiler/src/iree/compiler/Dialect/Stream/`。
|
||
> **这是 IREE 最有特色、也是 TVM 没有对应物的一层，值得重点理解。**
|
||
|
||
## 1. Stream 的职责：把同步图变成异步执行模型
|
||
|
||
Flow 层产出的还是「同步的图」（每个 op 依赖前一个的 result）。Stream 层把它变成**异步执行模型**：
|
||
|
||
- 每个计算（dispatch）变成可异步提交的工作
|
||
- 用 **timepoint**（类似 future/fence）显式表达「某批异步工作何时完成」
|
||
- 用 **resource** 管理内存的生命周期（分配/持有/释放）
|
||
- 用 **affinity** 标注每个 op 在哪个设备上执行
|
||
|
||
这一层是 IREE 支持多设备异步调度的核心。
|
||
|
||
## 2. 核心概念
|
||
|
||
### (1) timepoint —— 异步同步原语（灵魂）
|
||
|
||
`!stream.timepoint` = 一个「未来某个时刻会完成」的句柄，用来表达数据依赖和同步：
|
||
|
||
```mlir
|
||
// stream.async.execute 提交一批异步工作，返回 (results, timepoint)
|
||
%results, %result_timepoint = stream.async.execute ... {
|
||
%4 = stream.async.dispatch @matmul_dispatch_0::@...(%arg2[...], %arg3[...]) : ...
|
||
stream.yield %4 : !stream.resource<...>
|
||
} => !stream.timepoint
|
||
|
||
// stream.timepoint.await 等待 timepoint 完成，取出结果
|
||
%2 = stream.timepoint.await %result_timepoint => %results : !stream.resource<...>
|
||
```
|
||
|
||
### (2) resource —— 内存资源（带显式生命周期）
|
||
|
||
tensor（值语义）在 Stream 层变成 `!stream.resource<...>`（带大小），生命周期显式管理：
|
||
|
||
```mlir
|
||
%1 = stream.tensor.import on(#hal.device.affinity<@__device_0>) %arg0 : !hal.buffer_view -> tensor<2x4xf32> in !stream.resource<external>{%0}
|
||
```
|
||
|
||
resource 有：分配（alloc）、持有（retain）、释放（release）、子视图（subview）、
|
||
瞬态（transient）。生命周期由 **ARC（Automatic Reference Counting）** pass 自动管理。
|
||
|
||
### (3) affinity —— 设备亲和性
|
||
|
||
每个 op 标注 `on(#hal.device.affinity<@__device_0>)`，表示在哪个设备上执行。
|
||
跨设备的 op 需要 `stream.async.transfer` 来传输数据。
|
||
|
||
## 3. 四个抽象层次（关键！）
|
||
|
||
Stream dialect 内部有从高到低的层次，编译时逐步下降：
|
||
|
||
| 层次 | op 前缀 | 何时产生 | 说明 |
|
||
|---|---|---|---|
|
||
| tensor | `stream.tensor.*` | ConvertToStreamPass | 值语义张量（最接近 Flow） |
|
||
| async | `stream.async.*` + `stream.timepoint.*` | ScheduleExecutionPass | 异步执行 + timepoint 同步 |
|
||
| cmd | `stream.cmd.*` | 接近 HAL | 命令缓冲（dispatch/copy/fill/flush） |
|
||
| resource | `stream.resource.*` | 全流程 | 内存资源生命周期 |
|
||
|
||
## 4. 核心 op（`IR/StreamOps.td`）
|
||
|
||
- **tensor 层**：`stream.tensor.import/export`、`stream.tensor.dispatch`、`stream.tensor.sizeof`
|
||
- **async 层**：`stream.async.execute`、`stream.async.dispatch`、`stream.async.transfer`、
|
||
`stream.async.copy/fill/load/store`、`stream.async.barrier`
|
||
- **timepoint**：`stream.timepoint.await/barrier/join/import/export/immediate`
|
||
- **cmd 层**：`stream.cmd.dispatch/execute/copy/fill/flush/invalidate/serial/concurrent`
|
||
- **resource**：`stream.resource.alloc/dealloca/size/retain/release/subview/transients`
|
||
- **executable**：`stream.executable`、`stream.binding.subspan`
|
||
|
||
## 5. 关键 pass（`Transforms/` + `Analysis/`）
|
||
|
||
| pass | 文件 | 作用 |
|
||
|---|---|---|
|
||
| `ConvertToStreamPass` | ConvertToStream.cpp | Flow → Stream（核心入口） |
|
||
| `ScheduleExecutionPass` | ScheduleExecution.cpp | **调度**：图 → async.execute + timepoint |
|
||
| `ScheduleConcurrencyPass` | ScheduleConcurrency.cpp | 并发调度（并行 workgroup 重叠） |
|
||
| `PropagateTimepointsPass` | PropagateTimepoints.cpp | timepoint 传播（优化同步点） |
|
||
| `ScheduleAllocationPass` | ScheduleAllocation.cpp | 资源分配调度 |
|
||
| `AutomaticReferenceCountingPass` | AutomaticReferenceCounting.cpp | **ARC**：自动引用计数管理资源生命周期 |
|
||
| `ElideTimepointsPass` / `ElideAsyncCopiesPass` | 同名 | 消除冗余 timepoint / async copy |
|
||
| `ReuseAllocationsPass` / `RefineUsagePass` | 同名 | 内存复用 / 使用细化 |
|
||
|
||
**Analysis/**（调度前的分析）：
|
||
- `Affinity.cpp` —— 设备亲和性分析
|
||
- `Partitioning.cpp` —— 分区分析
|
||
- `ResourceHazards.cpp` / `ResourceUsage.cpp` —— 资源冲突 / 使用分析
|
||
|
||
## 6. IR 变换（阶段 2 已见，精简回顾）
|
||
|
||
```mlir
|
||
// ConvertToStreamPass：Flow → Stream，tensor 变 stream.resource，标上 affinity
|
||
%7 = stream.tensor.dispatch on(#hal.device.affinity<@__device_0>) @matmul_dispatch_0::@...(%2, %5) : ...
|
||
|
||
// ScheduleExecutionPass：包进 async.execute，产生 timepoint
|
||
%results, %result_timepoint = stream.async.execute ... { ... } => !stream.timepoint
|
||
%2 = stream.timepoint.await %result_timepoint => %results : !stream.resource<...>
|
||
```
|
||
|
||
## 7. 与 TVM 0.6 对照
|
||
|
||
| Stream 概念 | TVM 0.6 | 说明 |
|
||
|---|---|---|
|
||
| timepoint / async.execute | **无直接对应** | TVM 的异步主要在 runtime（线程池/stream），不在编译期 IR 里显式建模 |
|
||
| resource 生命周期 | **无**（TVM 的 buffer 生命周期由 runtime 管） | IREE 在编译期显式建模并 ARC |
|
||
| affinity | device 注解 | 类似但 IREE 更显式 |
|
||
| scheduling | 无（TVM 是同步图 + runtime 调度） | IREE 把调度变成 IR 变换 |
|
||
|
||
**核心差异**：TVM 把「图是同步的、异步是 runtime 的事」；IREE 把「异步调度」提升为**编译期的显式 IR 层**，
|
||
可以针对多设备做静态调度优化。这就是为什么 IREE 在多设备/异构场景更强。
|
||
|
||
## 8. 源码阅读路径
|
||
|
||
```
|
||
src/iree/compiler/src/iree/compiler/Dialect/Stream/
|
||
├── IR/StreamOps.td # op 定义（先读，认识 tensor/async/cmd/resource 四层）
|
||
├── Analysis/ # Affinity / Partitioning / ResourceHazards / ResourceUsage
|
||
├── Transforms/ConvertToStream.cpp # Flow → Stream
|
||
├── Transforms/ScheduleExecution.cpp # 异步调度（核心）
|
||
├── Transforms/PropagateTimepoints.cpp # timepoint 优化
|
||
└── Transforms/AutomaticReferenceCounting.cpp # 资源生命周期 ARC
|
||
```
|
||
|
||
> 学习建议：先读 `ScheduleExecution.cpp` 理解「怎么从同步图变异步」，
|
||
> 再读 `PropagateTimepoints.cpp` 理解「timepoint 怎么优化」，
|
||
> 最后读 `AutomaticReferenceCounting.cpp` 理解「资源生命周期怎么管」。
|