[Mlir-commits] [mlir] [mlir][IR][WIP] Set insertion point when erasing an operation (PR #146955)

[Markus Böck]

================
@@ -152,12 +152,45 @@ void RewriterBase::replaceOp(Operation *op, Operation *newOp) {
   eraseOp(op);
 }
 
+/// Returns the given block iterator if it lies within the block `b`.
+/// Otherwise, otherwise finds the ancestor of the given block iterator that
+/// lies within `b`. Returns and "empty" iterator if the latter fails.
+///
+/// Note: This is a variant of Block::findAncestorOpInBlock that operates on
+/// block iterators instead of ops.
+static std::pair<Block *, Block::iterator>
+findAncestorIteratorInBlock(Block *b, Block *itBlock, Block::iterator it) {
+  // Case 1: The iterator lies within the block.
+  if (itBlock == b)
+    return std::make_pair(itBlock, it);
+
+  // Otherwise: Find ancestor iterator. Bail if we run out of parent ops.
+  Operation *parentOp = itBlock->getParentOp();
+  if (!parentOp)
+    return std::make_pair(static_cast<Block *>(nullptr), Block::iterator());
+  Operation *op = b->findAncestorOpInBlock(*parentOp);
+  if (!op)
+    return std::make_pair(static_cast<Block *>(nullptr), Block::iterator());
+  return std::make_pair(op->getBlock(), op->getIterator());
+}
+
 /// This method erases an operation that is known to have no uses. The uses of
 /// the given operation *must* be known to be dead.
 void RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
   auto *rewriteListener = dyn_cast_if_present<Listener>(listener);
 
+  // If the current insertion point is before/within the erased operation, we
+  // need to adjust the insertion point to be after the operation.
+  if (getInsertionBlock()) {
+    Block *insertionBlock;
+    Block::iterator insertionPoint;
+    std::tie(insertionBlock, insertionPoint) = findAncestorIteratorInBlock(
+        op->getBlock(), getInsertionBlock(), getInsertionPoint());
+    if (insertionBlock && insertionPoint == op->getIterator())
+      setInsertionPointAfter(op);
+  }
----------------
zero9178 wrote:

What do you think about simplifying this to just:
```
  if (getInsertionBlock())
    if (getInsertionPoint() == op->getIterator())
      setInsertionPointAfter(op);
```

I agree with you that we don't want to make this too expensive. I think the above should cover the cases discussed in https://github.com/llvm/llvm-project/pull/146908 and is a rather very small cost. 

I personally do not have expectations of the insertion point still working if the block is being deleted that one is inserting, but I think few users are thinking "My insertion point is before this op and the op is used as an anchor inside the builder to insert new ops, therefore I now need to set the insertion point again after I deleted the op". Even less so in patterns where the framework sets the insertion point for you!

https://github.com/llvm/llvm-project/pull/146955