【infra之路】AI 编译器学习笔记（六）：重拾 C++——lambda 与 MLIR 中的回调

读 MLIR 代码时，walk、llvm::any_of、构造 scf.for 的循环体，到处是 [&](...) { ... }。这一篇把 lambda 讲清楚，并把它和前几篇反复出现的"纸条"问题连起来。API 以 LLVM 20 为准。

---

一、为什么需要 lambda

很多 API 的设计是"你给我一段代码，我在合适的时候替你执行"。walk 负责遍历所有 op，对每个 op 做什么由你决定。lambda 让你把这段代码直接写在调用处：

// 不用 lambda：先定义，再传名字
void printName(Operation *op) { llvm::outs() << op->getName() << "\n"; }
func.walk(printName);

// 用 lambda：就地写
func.walk([](Operation *op) { llvm::outs() << op->getName() << "\n"; });

二、语法

[捕获列表](参数列表) { 函数体 }

相比普通函数，只是把"返回类型 + 函数名"换成了 []。返回类型通常由编译器从 return 推出，需要时可写成 -> 类型 放在参数列表后。lambda 也可以存进变量：

auto square = [](int x) { return x * x; };
int y = square(5);   // 25

三、捕获

函数体默认看不到外部局部变量，需要在方括号里声明：

int count = 0;
func.walk([&](Operation *op) {
  if (isa<ReluOp>(op)) ++count;   // 改的是外面的 count
});

写法	含义
[]	不捕获
[&]	用到的外部变量都按引用捕获
[=]	用到的外部变量都按值拷贝
[&count, ctx]	逐个指定：count 按引用，ctx 按值
[this]	捕获当前对象，可调成员函数

计数的例子必须用 [&]。用 [=] 会编译报错，因为按值捕获的副本默认只读；即使加 mutable，改的也只是副本，外面的 count 仍是 0。

四、按引用捕获也是"纸条"

[&] 捕获时，lambda 里存的只是指向外部变量的引用。lambda 被调用时，那些变量必须还活着。

std::function<int()> makeCounter() {
  int n = 0;
  return [&]() { return ++n; };   // n 在函数返回时销毁
}

改法：按值捕获 [n]() mutable { return ++n; }，让 lambda 自己带一份。

再看一个更隐蔽的例子：

std::function<void()> cb;
{
  SmallVector<int> data = {1, 2, 3};
  cb = [&]() { llvm::outs() << data.size(); };
}
cb();   // data 已销毁，未定义行为

std::function 确实拥有那个 lambda，但它只拥有 lambda 对象本身，lambda 引用的东西它不管。保管好了纸条，房子照样会被拆。

经验规则 ：lambda 当场用完的（如传给 walk），放心用 [&]；要存起来以后调用的，想清楚它引用的东西届时是否还在。

五、function_ref 与 std::function

函数接收 lambda 时，参数类型常见两种：

• std::function<void(Operation *)>：拷贝并拥有一份 lambda，可以存起来。
• llvm::function_ref<void(Operation *)>：不拥有，只指向调用者的 lambda，开销更小。

void(Operation *) 描述的是"接收 Operation *、不返回值的函数"。

function_ref 和 StringRef、ArrayRef 是一家人。调用时传入的 lambda 通常是临时对象：

forEachUser(v, [&](Operation *u) { ... });   // lambda 在分号处销毁

在 forEachUser 内部调用 fn 是安全的；但如果把 fn 存进全局变量留到以后调用，它就悬空了。LLVM 的约定：只在调用期间使用的参数写 function_ref，需要保存的写 std::function。

六、MLIR 中的三种典型用法

1. 遍历：walk

参数写成具体 op 类型，walk 会只对这种 op 调用 lambda，相当于内部替你做了 dyn_cast：

int bnum = 0;
func.walk([&](BiasAddOp) { bnum++; });   // 参数没用到，可以不起名字

需要中途停止时返回 WalkResult：

bool hasRelu = false;
func.walk([&](ReluOp) {
  hasRelu = true;
  return WalkResult::interrupt();
});

2. 判断：llvm::any_of / all_of

bool allF32 = llvm::all_of(op->getOperands(), [&](Value v) {
  return v.getType().isF32();
});

3. 构造带 region 的 op：body builder

scf.for、linalg.generic 内部有一段代码。创建时传入一个 lambda，框架建好外壳后调用它填充内容：

rewriter.create<scf::ForOp>(loc, lb, ub, step, ValueRange{},
    [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {
      Value x = b.create<memref::LoadOp>(loc, input, ValueRange{iv});
      b.create<memref::StoreOp>(loc, x, output, ValueRange{iv});
      b.create<scf::YieldOp>(loc);   // 提供 body builder 时需自己补上 yield
    });

lambda 的参数由框架传入：b 的插入点已设在循环体内，iv 是循环变量。

---

小结

• lambda 是就地写的匿名函数，[] 里声明如何使用外部变量。
• 当场用完用 [&]；要保存的 lambda 注意引用的生命周期。
• 引用、StringRef、ArrayRef、function_ref、按引用捕获都是不拥有数据的"纸条"，判断标准始终是：指向谁，用的时候它还在不在。
• MLIR 中 lambda 主要出现在 walk、any_of/all_of 和 body builder 三处。

至此，读 MLIR 代码所需的 C++ 基础告一段落，接下来回到编译器主线。