# 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 官方文档**: - **IREE 代码库**: - 中段:`compiler/src/iree/compiler/Dialect/Flow/`、`.../Dialect/Stream/`、`.../Dialect/HAL/` - codegen:`compiler/src/iree/compiler/Codegen/` - **MLIR 官方文档**: - **MLIR Toy tutorial**: - **Linalg 文档**: - **论文**:*Composable and Modular Code Generation in MLIR*(IREE codegen 的核心设计) - **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 侧的机制,但不要因为它熟就跳过其他层。