182 lines
6.1 KiB
Markdown
182 lines
6.1 KiB
Markdown
# 在 IREE 里写一个自定义 Pass
|
||
|
||
> 2026-08-16。基于源码构建的 iree-compile，写 pass 的完整指南。
|
||
> 以在 **Flow dialect** 里加一个最小 pass 为例，代码可直接复制。
|
||
|
||
## 1. Pass 的三件套
|
||
|
||
IREE 的 pass 需要改 3 个地方（都在某个 dialect 的 `Transforms/` 下）：
|
||
|
||
| 文件 | 作用 |
|
||
|---|---|
|
||
| `Passes.td` | TableGen 声明 pass（名字、作用的操作类型、选项） |
|
||
| `HelloWorld.cpp`（新） | pass 实现（`runOnOperation`） |
|
||
| `CMakeLists.txt` | 把新 .cpp 加入编译 |
|
||
|
||
以 Flow 为例，路径是：
|
||
`src/iree/compiler/src/iree/compiler/Dialect/Flow/Transforms/`
|
||
|
||
## 2. 最小示例：HelloWorld pass（打印所有 op）
|
||
|
||
### 2.1 在 `Passes.td` 里声明（加一行）
|
||
|
||
在 `Passes.td` 末尾加：
|
||
```td
|
||
def HelloWorldPass :
|
||
Pass<"iree-flow-hello-world", "mlir::ModuleOp"> {
|
||
let summary = "Prints all operations in the module (learning example).";
|
||
}
|
||
```
|
||
|
||
- `iree-flow-hello-world` = pass 的 CLI 名字（`iree-opt --iree-flow-hello-world` 用）
|
||
- `mlir::ModuleOp` = 这个 pass 作用于整个 module（`getOperation()` 返回 `ModuleOp`）
|
||
- 传 `""` 则作用于任意 op（泛型 pass）
|
||
|
||
### 2.2 新建 `HelloWorld.cpp`
|
||
|
||
```cpp
|
||
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
|
||
|
||
#include "mlir/IR/BuiltinOps.h"
|
||
#include "mlir/Pass/Pass.h"
|
||
#include "llvm/Support/raw_ostream.h"
|
||
|
||
namespace mlir::iree_compiler::IREE::Flow {
|
||
|
||
#define GEN_PASS_DEF_HELLOWORLDPASS
|
||
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h.inc"
|
||
|
||
namespace {
|
||
|
||
struct HelloWorldPass : impl::HelloWorldPassBase<HelloWorldPass> {
|
||
using impl::HelloWorldPassBase<HelloWorldPass>::HelloWorldPassBase;
|
||
|
||
void runOnOperation() override {
|
||
// 遍历 module 里所有 op 并打印名字
|
||
getOperation().walk([&](Operation *op) {
|
||
llvm::errs() << " [op] " << op->getName() << "\n";
|
||
});
|
||
}
|
||
};
|
||
|
||
} // namespace
|
||
} // namespace mlir::iree_compiler::IREE::Flow
|
||
```
|
||
|
||
### 2.3 在 `CMakeLists.txt` 里加入
|
||
|
||
在 `Transforms/CMakeLists.txt` 的 `SRCS` 列表（按字母序）加：
|
||
```cmake
|
||
"HelloWorld.cpp"
|
||
```
|
||
|
||
### 2.4 重新构建
|
||
|
||
```bash
|
||
export PATH="/home/ylhan/ai/ai_compiler/.venv/bin:$PATH"
|
||
cmake --build /home/ylhan/ai/ai_compiler/src/iree-build/ --target iree-opt -j 16
|
||
```
|
||
（只重建 `iree-opt`，增量编译很快，几秒到几十秒）
|
||
|
||
## 3. 关键结构解析（理解这套模板）
|
||
|
||
```cpp
|
||
#define GEN_PASS_DEF_HELLOWORLDPASS
|
||
#include ".../Passes.h.inc"
|
||
```
|
||
- TableGen 从 `Passes.td` 生成 `Passes.h.inc`，里面定义了 `impl::HelloWorldPassBase<HelloWorldPass>` 基类
|
||
- `GEN_PASS_DEF_*` 宏告诉 `.inc` 文件「我要生成这个 pass 的基类定义」
|
||
|
||
```cpp
|
||
struct HelloWorldPass : impl::HelloWorldPassBase<HelloWorldPass> { ... }
|
||
```
|
||
- 你的 pass 继承生成的基类（CRTP 模式），`impl::` 是 `IREE::Flow::impl` 命名空间
|
||
|
||
**不需要手动注册** —— `Passes.cpp` 里的 `registerPasses()` 是从 TableGen 生成的，
|
||
会自动 `PassRegistration` 所有在 `Passes.td` 里声明的 pass。
|
||
|
||
## 4. 用 `iree-opt` 单测 pass
|
||
|
||
`iree-opt` 是测试 pass 的核心工具（不跑完整编译管线，只跑你指定的 pass）：
|
||
|
||
```bash
|
||
# 在构建目录里
|
||
build/tools/iree-opt --iree-flow-hello-world input.mlir
|
||
```
|
||
|
||
先写一个测试输入 `test.mlir`：
|
||
```mlir
|
||
func.func @main(%arg0: tensor<2x4xf32>, %arg1: tensor<4x2xf32>) -> tensor<2x2xf32> {
|
||
%c0 = arith.constant 0.0 : f32
|
||
%empty = tensor.empty() : tensor<2x2xf32>
|
||
%init = linalg.fill ins(%c0 : f32) outs(%empty : tensor<2x2xf32>) -> tensor<2x2xf32>
|
||
%r = linalg.matmul ins(%arg0, %arg1 : tensor<2x4xf32>, tensor<4x2xf32>)
|
||
outs(%init : tensor<2x2xf32>) -> tensor<2x2xf32>
|
||
return %r : tensor<2x2xf32>
|
||
}
|
||
```
|
||
|
||
预期输出：打印 module 里所有 op 的名字（`func.func`, `arith.constant`, `tensor.empty`, `linalg.fill`, `linalg.matmul`, ...）。
|
||
|
||
### 常用调试 flag
|
||
|
||
```bash
|
||
iree-opt --iree-flow-hello-world --mlir-print-ir-after-all test.mlir # 看 pass 后的 IR
|
||
iree-opt --pass-pipeline='builtin.module(iree-flow-hello-world)' test.mlir # 显式指定 pipeline
|
||
```
|
||
|
||
## 5. 进阶：真正做变换的 pass（Pattern Rewriting）
|
||
|
||
HelloWorld 只是「看」，真正的 pass 是「改」。核心是 **Pattern Rewriting**（见 [03-mlir-basics.md](03-mlir-basics.md) 第 7 节）：
|
||
|
||
```cpp
|
||
// 一个 pattern：把 arith.addi %x, 0 简化成 %x
|
||
struct SimplifyAddZero : OpRewritePattern<arith::AddIOp> {
|
||
using OpRewritePattern<arith::AddIOp>::OpRewritePattern;
|
||
|
||
LogicalResult matchAndRewrite(arith::AddIOp op,
|
||
PatternRewriter &rewriter) const override {
|
||
if (auto cst = op.getRhs().getDefiningOp<arith::ConstantOp>()) {
|
||
if (cst.value() == 0) {
|
||
rewriter.replaceOp(op, op.getLhs());
|
||
return success();
|
||
}
|
||
}
|
||
return failure();
|
||
}
|
||
};
|
||
|
||
struct MySimplifyPass : impl::MySimplifyPassBase<MySimplifyPass> {
|
||
void runOnOperation() override {
|
||
RewritePatternSet patterns(&getContext());
|
||
patterns.add<SimplifyAddZero>(&getContext());
|
||
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
|
||
return signalPassFailure();
|
||
}
|
||
}
|
||
};
|
||
```
|
||
|
||
- `OpRewritePattern<OpTy>`：声明「匹配某个 op，重写」
|
||
- `matchAndRewrite`：匹配逻辑 + 重写逻辑，`rewriter.replaceOp` 替换 op
|
||
- `applyPatternsGreedily`：贪婪遍历 IR 应用 patterns（对应 IREE 里的 `CanonicalizePass`）
|
||
|
||
## 6. 下一步方向
|
||
|
||
1. 先把 HelloWorld 跑通，理解「声明 → 实现 → 构建 → 单测」的闭环
|
||
2. 改写一个真正的 pattern（如上）
|
||
3. 深入读 IREE 现有的 pass 源码，照着改：
|
||
- 简单 pattern：`Flow/Transforms/Canonicalize.cpp`
|
||
- 复杂 pass：`Codegen/LLVMCPU/LLVMCPUTile.cpp`
|
||
4. 写 pass 时配合 `--mlir-print-ir-before-all` / `--mlir-print-ir-after-all` 看 IR 变化
|
||
|
||
## 7. 参考：pass 位置速查
|
||
|
||
| 想改的层 | 路径 |
|
||
|---|---|
|
||
| Flow（图级） | `compiler/src/iree/compiler/Dialect/Flow/Transforms/` |
|
||
| Stream（异步） | `compiler/src/iree/compiler/Dialect/Stream/Transforms/` |
|
||
| HAL（设备） | `compiler/src/iree/compiler/Dialect/HAL/Transforms/` |
|
||
| Codegen（CPU） | `compiler/src/iree/compiler/Codegen/LLVMCPU/` |
|
||
| Codegen（GPU） | `compiler/src/iree/compiler/Codegen/LLVMGPU/` |
|