first commit

2026-08-20 17:20:09 +08:00
commit 8049b6f58d
14 changed files with 1649 additions and 0 deletions
@@ -0,0 +1,181 @@
+# 在 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/` |