121 lines
6.3 KiB
Markdown
121 lines
6.3 KiB
Markdown
# IREE 学习规划
|
||
|
||
> 目标：系统掌握 IREE **整体**，建立完整的 pipeline 心智模型。
|
||
> codegen 只是其中一个环节（恰好有经验，可快速过），不作为重点。
|
||
|
||
## 0. 先回答：是否应该先学 MLIR？
|
||
|
||
**结论：要学，但用「MLIR 基础 + 边拆 IREE 边补 MLIR」的方式，不要从头啃完 MLIR 教程。**
|
||
|
||
理由：
|
||
- IREE 的**每一层都是 MLIR dialect**（Flow、Stream、HAL、以及 codegen 用到的 Linalg/Vector），
|
||
不只是 codegen 依赖 MLIR。没有 MLIR 基础，整个 pipeline 都看不懂。
|
||
- 但已有 TVM 0.6 经验，IR/pass/调度/访存等**编译器核心概念已具备**，MLIR 只是换表达方式，
|
||
不需要从零学编译原理，只需补「MLIR 的表达方式 + 基础设施」。
|
||
|
||
**MLIR 需要掌握的核心（聚焦这些即可，不必面面俱到）：**
|
||
- Operation / Value / Block / Region / SSA —— 对应熟悉的 IR node / expression
|
||
- Dialect 与 OpBuilder
|
||
- **Pass + Pattern Rewriting / Dialect Conversion** —— MLIR 的两种 pass 机制
|
||
- 核心 dialect：`linalg` `tensor` `vector` `memref` `arith` `scf` `func`
|
||
|
||
## 1. IREE 完整 pipeline（心智模型，先建立这个）
|
||
|
||
```
|
||
PyTorch / TF / JAX / ONNX
|
||
↓ (torch-mlir / iree-import 系列工具)
|
||
StableHLO / TOSA ← MLIR 前端方言（跨框架的统一入口）
|
||
↓
|
||
IREE Flow dialect ← 图级优化：type 推导、算子融合、划分 dispatch region（executable）
|
||
↓
|
||
IREE Stream dialect ← 异步执行模型：timepoint / fence 调度，把图变成可异步执行的流
|
||
↓
|
||
IREE HAL dialect ← 设备抽象：device / command buffer / allocator / executable 加载
|
||
↓
|
||
Codegen（executable） ← Linalg/Vector/Tensor → LLVM IR / SPIR-V（设备相关代码生成）
|
||
↓
|
||
IREE VM ← 序列化成紧凑字节码，供运行时加载
|
||
↓
|
||
iree-runtime（C） ← 加载 VM 模块 + HAL driver（CPU/CUDA/Vulkan/...）执行
|
||
```
|
||
|
||
## 2. 关键概念映射（TVM 0.6 → IREE/MLIR）
|
||
|
||
| TVM 0.6 | IREE / MLIR | 说明 |
|
||
|---|---|---|
|
||
| Relay（图级 IR、fusion、算子融合） | Flow dialect + StableHLO/TOSA | 图级优化、形成 dispatch region |
|
||
| TE / TIR + schedule 原语（split/fuse/reorder/unroll） | Linalg + Vector + tiling/vectorization | tensor-level → loop-level 代码生成 |
|
||
| codegen（LLVM/CUDA/OpenCL） | 翻译到 LLVM / SPIR-V | 后端代码生成 |
|
||
| runtime（graph runtime / 早期 VM） | HAL + VM | 设备抽象 + 字节码执行 |
|
||
| （TVM 无直接对应） | Stream dialect | 异步调度，IREE 特色之一 |
|
||
|
||
## 3. 学习阶段（系统覆盖整体）
|
||
|
||
### 阶段 0：环境准备
|
||
- 安装 IREE（推荐 nightly Python 包 `iree-compiler` + `iree-runtime`，或从源码编译）
|
||
- 跑通一个最小例子：`iree-compile` 编译 → `iree-run-module` 执行
|
||
- 熟悉工具链：`iree-compile` / `iree-run-module` / `iree-dump-module`
|
||
|
||
### 阶段 1：MLIR 基础（IREE 全流程的地基）
|
||
- MLIR 核心概念：Operation / Value / Block / Region / Dialect / Type / Attribute
|
||
- 核心 dialect：`linalg` `tensor` `vector` `memref` `arith` `scf` `func`
|
||
- Pass 与 Pattern Rewriting / Dialect Conversion
|
||
- 工具：`mlir-opt` / `mlir-translate`，用 `--mlir-print-ir-before-all` / `--mlir-print-ir-after-all` 观察 pass 效果
|
||
- 参考：MLIR 官方 Toy tutorial + 语言参考 + linalg 文档
|
||
|
||
### 阶段 2：IREE 架构总览与完整管线
|
||
- 用 `--mlir-print-ir-after-all` 追踪一个简单模型从输入到 VM 的完整变换
|
||
- 认识每一层 dialect 的职责与边界（对上「心智模型」图）
|
||
- 理解各层之间如何衔接（输入 → Flow → Stream → HAL → executable → VM）
|
||
|
||
### 阶段 3：前端输入（StableHLO / TOSA）
|
||
- 各框架如何进入 IREE：torch-mlir、iree-import-tf / iree-import-onnx
|
||
- StableHLO 与 TOSA 的角色、差异、如何选择
|
||
- 输入的图结构 → 进入 Flow 前的形态
|
||
|
||
### 阶段 4：编译中段（Flow + Stream）
|
||
- Flow：type 推导、fusion、dispatch region（executable）的形成
|
||
- Stream：异步执行模型，timepoint / fence 如何表示数据依赖与调度
|
||
- 图级优化与 TVM Relay 的对照
|
||
|
||
### 阶段 5：HAL 设备抽象
|
||
- device / command buffer / allocator / executable 的抽象设计
|
||
- HAL 如何屏蔽 CPU / CUDA / Vulkan 等后端差异
|
||
- 与 TVM codegen 的 device 抽象对照
|
||
|
||
### 阶段 6：Codegen（可执行体 lowering，有经验可快进）
|
||
- executable 的 lowering：Linalg → Vector → LLVM IR / SPIR-V
|
||
- 后端选择：LLVMCPU、LLVMGPU(CUDA)、SPIR-V(Vulkan) 等
|
||
- tile-and-fuse / tiling / vectorization / 访存优化
|
||
- 与 TVM TE/TIR + schedule 的对应关系（你已熟，重点是补 MLIR 侧的机制）
|
||
|
||
### 阶段 7：运行时（VM + iree-runtime）
|
||
- IREE VM 字节码格式与执行模型
|
||
- iree-runtime（C）：模块加载、HAL driver、调度执行
|
||
- 与 TVM runtime 对比
|
||
|
||
### 阶段 8：实践项目（巩固，可选）
|
||
- 写一个自定义 pass（例如自定义 fusion / tiling pass）
|
||
- 加一个自定义 backend，或写一个自定义 dialect 练手
|
||
- 追踪并优化一个具体算子的端到端 lowering
|
||
|
||
## 4. 学习资源
|
||
|
||
- **IREE 官方文档**：<https://iree.dev/>
|
||
- **IREE 代码库**：<https://github.com/iree-org/iree>
|
||
- 中段：`compiler/src/iree/compiler/Dialect/Flow/`、`.../Dialect/Stream/`、`.../Dialect/HAL/`
|
||
- codegen：`compiler/src/iree/compiler/Codegen/`
|
||
- **MLIR 官方文档**：<https://mlir.llvm.org/>
|
||
- **MLIR Toy tutorial**：<https://mlir.llvm.org/docs/Tutorials/Toy/>
|
||
- **Linalg 文档**：<https://mlir.llvm.org/docs/Dialects/Linalg/>
|
||
- **论文**：*Composable and Modular Code Generation in MLIR*（IREE codegen 的核心设计）
|
||
- **torch-mlir**：<https://github.com/llvm/torch-mlir>（前端 → StableHLO/TOSA 参考）
|
||
|
||
## 5. 学习方法建议
|
||
|
||
- **以 IR dump 为主线**：每个阶段都用 `--mlir-print-ir-after-all` / `--mlir-print-ir-before-all`
|
||
看 IR 变化，这是理解 pass pipeline 最快的方式。
|
||
- **对照 TVM**：每遇到一个 IREE 概念，都问自己「这在 TVM 里对应什么」，把已有经验迁移过来。
|
||
- **按层推进，不跳层**：前端 → 中段 → HAL → codegen → 运行时，每层弄懂职责和边界后再进下一层。
|
||
codegen 有经验，可以快速过、只补 MLIR 侧的机制，但不要因为它熟就跳过其他层。
|