MLIR OpBuilder

概述

在 MLIR 里，你不能随便 new 一个 op，必须通过 OpBuilder 这个"工厂"来创建。

为什么必须用 OpBuilder？

在 MLIR 中，一个 Operation 不是孤立存在的。它必须：

• 属于某个 Block
• 在 Block 中有明确的位置（插入点）
• 自动建立与操作数（operands）的 def-use 链
• 能被 Pass 的 Listener 追踪（以便回滚、通知）

如果你直接 new Operation(...)，这些关系都不会自动建立，IR 会处于非法状态。

所以 MLIR 强制要求：所有 op 必须通过 OpBuilder 创建。

OpBuilder 做三件事：

• 记录插入点：新 op 插到哪个块、哪个位置。
• 自动入图：创建出的 op 自动加入 IR，自动建立 def-use 链。
• 通知监听器：PatternRewriter 是一种特殊 builder，它记录所有创建/删除/修改，失败时可回滚。

用 OpBuilder 创建一个加法 op

假设我们有一个函数 @foo，里面已经有两条常量指令，我们想插入一条加法：

func.func @foo() -> i64 {
  %0 = arith.constant 10 : i64
  %1 = arith.constant 20 : i64
  // 这里想插入 %2 = arith.addi %0, %1 : i64
  return %0 : i64
}

步骤 1：拿到 OpBuilder 并设置插入点

void addTwoConstants(mlir::func::FuncOp funcOp) {
  // 创建 builder，绑定到函数第一个块的开头
  mlir::OpBuilder builder(&funcOp.getBody().front());
  
  // 或者：设置插入点到块末尾
  builder.setInsertionPointToEnd(&funcOp.getBody().front());
  
  // 或者：插到某条 op 前面
  // builder.setInsertionPoint(someOp);
  
  // 或者：插到某条 op 后面
  // builder.setInsertionPointAfter(someOp);
}

setInsertionPoint 就是移动 builder 的"光标"。创建出的 op 自动出现在光标位置。

步骤 2：用 create<T>() 创建 op

// 拿到已有的两条常量
auto cst0 = ...;  // %0 = arith.constant 10 : i64
auto cst1 = ...;  // %1 = arith.constant 20 : i64

// 创建加法 op
auto addOp = builder.create<mlir::arith::AddIOp>(
    loc,                        // Location：源码位置
    cst0.getResult(),           // 第一个操作数
    cst1.getResult()            // 第二个操作数
);
// addOp 已经插入到块末尾，并返回 arith::AddIOp 包装类
// addOp.getResult() 就是 %2

步骤 3：create<T> 的三段式模板

builder.create<mlir::arith::AddIOp>(loc, resultTypes..., operands...);
//           │                        │     │              │
//           │                        │     │              └─ 操作数：%0, %1
//           │                        │     └─ 结果类型（有时可省略，由 op 推断）
//           │                        └─ 源码位置
//           └─ 模板参数：要创建的 op 类型

create 内部会调用 AddIOp::build(builder, state, args...)。每个 op 都有一个静态 build 函数，决定"参数怎么变成操作数/结果/属性"。

两种构造 op 的路径

路径	build 从哪来	场景	例子
默认 build	tblgen 按 (ins)/(outs) 自动生成	大多数 op	arith.addi 只接受两个操作数，tblgen 自动生成 build
自定义 build	.td 里 let builders = OpBuilder\<...\> 或手写重载	需要语义化参数	只传 shape，自动造 type

自定义 build 的例子

假设我们定义了一个 simple.constant op，希望支持"只传一个整数值，自动推断类型为 i64"：

def Simple_ConstantOp : Simple_Op<"constant", [Pure]> {
  let arguments = (ins I64Attr:$value);
  let results = (outs I64:$result);
  
  // 自定义 builder：只传值，自动造 i64 类型
  let builders = [
    OpBuilder<(ins "int64_t":$value), [{
      build($_builder, $_state, $_builder.getI64Type(), 
            $_builder.getI64IntegerAttr(value));
    }]>
  ];
}

这样使用时：

// 默认 build（tblgen 生成）：需要传 Type + Attribute
auto cst1 = builder.create<simple::ConstantOp>(
    loc, builder.getI64Type(), builder.getI64IntegerAttr(42));

// 自定义 build：只传 42，自动推断类型
auto cst2 = builder.create<simple::ConstantOp>(loc, 42);

自定义 build 让调用更语义化、更简洁。

Location：每个 op 的"出生证明"

loc 记录 op 来自源码的哪个位置（文件:行:列），报错时能追溯到源头：

error: 'simple.add' op 操作数类型不匹配
  note: see current operation: %2 = simple.add %0, %1 : i64
  note: at file.mlir:12:3   

Location 是调试的生命线。没有它，Pass 报错时你只知道"某条 op 有问题"，但不知道它从哪来。

创建 op 时必须传loc：

auto loc = builder.getUnknownLoc();  // 不知道位置时用这个
// 或者从已有 op 继承
auto loc = existingOp.getLoc();
// 或者在 Pass 里从当前 op 拿
auto loc = op->getLoc();

Pattern Rewriter：带"撤销追踪"的 builder

在 Pattern Rewrite（比如把 add(x, 0) 优化成 x）中，不能直接用 OpBuilder，必须用 PatternRewriter：

class AddZeroFolder : public mlir::OpRewritePattern<mlir::arith::AddIOp> {
  mlir::LogicalResult matchAndRewrite(
      mlir::arith::AddIOp op, mlir::PatternRewriter &rewriter) const override {
    
    // 检查是否有操作数是 0
    auto cst = op.getRhs().getDefiningOp<mlir::arith::ConstantOp>();
    if (!cst || !cst.getValue().cast<mlir::IntegerAttr>().getValue().isZero())
      return mlir::failure();
    
    // 用 rewriter 替换：把 add 的结果直接用 %0 替代
    rewriter.replaceOp(op, op.getLhs());
    return mlir::success();
  }
};

为什么必须用 rewriter？ 因为：

• rewriter.replaceOp 会记录"这条 op 被替换了"，如果后续 pattern 失败，框架可以回滚。
• rewriter.create 创建的新 op 也会被记录，支持撤销。
• 直接用 builder 或 new 会绕过追踪，导致 IR 状态不一致。

完整例子串联：从创建到替换

假设我们要在 @foo 里插入一条加法，然后用 Pattern Rewrite 消除它：

原始 IR：

func.func @foo() -> i64 {
  %0 = arith.constant 10 : i64
  %1 = arith.constant 0 : i64
  %2 = arith.addi %0, %1 : i64    // 这条可以被优化掉
  return %2 : i64
}

步骤 1：用 OpBuilder 插入一条新加法

void insertAdd(mlir::func::FuncOp funcOp) {
  mlir::OpBuilder builder(&funcOp.getBody().front());
  builder.setInsertionPointToEnd(&funcOp.getBody().front());
  
  auto cst0 = funcOp.getBody().front().front();  // %0
  auto cst1 = ...;  // %1
  
  auto addOp = builder.create<mlir::arith::AddIOp>(
      funcOp.getLoc(),
      cst0.getResult(),
      cst1.getResult());
  // 现在块末尾多了一条 %3 = arith.addi %0, %1
}

步骤 2：用 Pattern Rewriter 消除 add 0

class AddZeroFolder : public mlir::OpRewritePattern<mlir::arith::AddIOp> {
  mlir::LogicalResult matchAndRewrite(
      mlir::arith::AddIOp op, mlir::PatternRewriter &rewriter) const override {
    auto rhs = op.getRhs().getDefiningOp<mlir::arith::ConstantOp>();
    if (!rhs || !rhs.getValue().cast<mlir::IntegerAttr>().getValue().isZero())
      return mlir::failure();
    rewriter.replaceOp(op, op.getLhs());  // %2 直接用 %0 替代
    return mlir::success();
  }
};

优化后 IR：

func.func @foo() -> i64 {
  %0 = arith.constant 10 : i64
  %1 = arith.constant 0 : i64
  return %0 : i64    // %2 被 %0 替代，%1 变成死代码
}

总结

概念	含义	关键点
OpBuilder	创建 op 的唯一工厂	记录插入点、自动入图、通知监听器
create(args...)	创建 op 的模板方法	内部调用 T::build()，三段式：loc + 结果类型 + 操作数
默认 build	tblgen 自动生成	按 (ins)/(outs) 生成
自定义 build	手写或 .td 声明	支持语义化参数（如只传 shape）
Insertion Point	builder 的"光标"	setInsertionPointToStart/End/Before/After
Location	op 的"出生证明"	报错时追溯源码位置
PatternRewriter	带撤销追踪的 builder	Pattern Rewrite 必须用它，不能直接 new

一句话：**OpBuilder 是 MLIR 中创建 op 的唯一入口，它替你处理了"插到哪、怎么连边、如何通知"这些机械细节，让你只专注于表达意图。 **

而 PatternRewriter 是它的"安全版"，在重写 IR 时提供回滚保护。