# 在 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/` |