6.3 KiB
6.3 KiB
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：
linalgtensorvectormemrefarithscffunc
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：
linalgtensorvectormemrefarithscffunc - 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 侧的机制，但不要因为它熟就跳过其他层。