note_AI_compiler/iree/11-writing-a-pass.md

在 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 末尾加：

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

#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 列表（按字母序）加：

"HelloWorld.cpp"

2.4 重新构建

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. 关键结构解析（理解这套模板）

#define GEN_PASS_DEF_HELLOWORLDPASS
#include ".../Passes.h.inc"

• TableGen 从 Passes.td 生成 Passes.h.inc，里面定义了 impl::HelloWorldPassBase<HelloWorldPass> 基类
• GEN_PASS_DEF_* 宏告诉 .inc 文件「我要生成这个 pass 的基类定义」

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）：

# 在构建目录里
build/tools/iree-opt --iree-flow-hello-world input.mlir

先写一个测试输入 test.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

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 第 7 节）：

// 一个 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/