Is it safe to erase the current operation in the callback of walk traversal?

MLIR
1

For example, is the following code safe?

funcOp.walk([](Operation *op) {
  doSomething();
  if (condition)
    op->erase();
});

Do I have to write like this?

SmallVector<Operation *> toRemove;
funcOp.walk([&](Operation *op) {
  doSomething();
  if (condition)
    toRemove.push_back(op);
});

for (Operation *op : toRemove)
  op->erase();

I know the second implementation is safer than the first one. But is it necessary to start erasing after the walk process finishes?

Thanks in advance!

2

https://mlir.llvm.org/doxygen/classmlir_1_1Operation.html#a59740592240b950b8c8afcf4a2eb4113

A callback on a block or operation is allowed to erase that block or operation if either:

  • the walk is in post-order, or
  • the walk is in pre-order and the walk is skipped after the erasure.
3

Thanks!