[mlir][IR] Trigger notifyOperationRemoved callback for nested ops

[springerm]

Summary

When cloning an op, the notifyOperationInserted callback is triggered for all nested ops. Similarly, the notifyOperationRemoved callback should be triggered for all nested ops when removing an op.

Listeners may inspect the IR during a notifyOperationRemoved callback. Therefore, when multiple ops are removed in a single RewriterBase::eraseOp call, the notifications must be triggered in an order in which the ops could have been removed one-by-one:

• Op removals must be interleaved with notifyOperationRemoved callbacks. A callback is triggered right before the respective op is removed.
• Ops are removed post-order and in reverse order. Other traversal orders could delete an op that still has uses. (This is not avoidable in graph regions and with cyclic block graphs.)

[Mogball]

Is this meant to fix something?

[springerm]

I have a use case for it. I could work around it. But it's more of a consistency issue. (We notify of nested ops when they are created, so we should also notify when they are removed.) Not many people seem to be using the listeners, so I think nobody noticed this until now.

[springerm]

The GreedyPatternRewriteDriver can be simplified in the same go.

[springerm]

Also simplify GreedyPatternRewriteDriver as part of this change.

[rriddle]

How does the interact with dialect conversion?

[springerm]

In D144193#4134153, @rriddle wrote:

How does the interact with dialect conversion?

Dialect conversion bypasses this notification entirely. It overrides eraseOp and replaceOp. (There may be some opportunities to reuse a larger part of the RewriterBase infrastructure in various places of the dialect conversion.)

[springerm]

change traversal order: bottom to top

[springerm]

@rriddle There's an issue that I overlooked. The order in which the notifications should be triggered. It should be bottom-to-top. Because ops can only be removed bottom-to-top. And we should trigger the callback in the same order as if the ops were removed one-by-one. Does this sounds reasonable to you? I just noticed this because I had failures in the GreedyPatternRewriteDriver.

I think we don't have a way to walk ops bottom-to-top at the moment. Thus the custom traversal. I'm going to prepare a change that adds another parameter to Operation::walk and rebase this revision on top of it.

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

Add test case

[springerm]

interleave notifications and op removal

[springerm]

@Mogball @rriddle @mehdi_amini I made some significant changes to this revision since it was accepted. Could you take another look?

[springerm]

update

[mehdi_amini]

mlir/lib/IR/PatternMatch.cpp

283

Technically this return true is valid only for registered dialect. But seems like this is a debug function so not that important I guess

290

Blocks are not necessarily sorted topologically, how do you handle that?

[springerm]

mlir/lib/IR/PatternMatch.cpp

290

I didn't know that such IR was valid.

I think ideally we would walk the IR by following the "reverse successor" relationship of blocks. Not sure if it is really necessary here, there may be a simpler solution. Definitely something needs to be changed because I found an example that crahes.

Related to this: I've seen a use case in the DialectConversion where IR must be walked according to the "forward successor" relationship.

We could utilize the new Iterator template of Operation::walk: By providing two new iterators: ForwardDominanceIterator and ReverseDominanceIterator. They work like ForwardIterator and ReverseIterator but walk blocks according to successors/reverse successors.

These iterators would be more heavyweight than simple forward/reverse iterators (i.e., they need to keep track of visited blocks etc. in case of cycles). Also, the current detail::walk implementations would have to be moved from Visitors.cpp to Visitors.h, which is already tricky without including Operation.h in Visitors.h (cyclic include).

[springerm]

fix unstructured control flow

[springerm]

This should now handle arbitrary successor/predecessor structures and/or graph regions correctly.

[springerm]

This does not work correctly if an op in a dead region has nested ops. I was trying to do without, but we really need a ReverseDominanceIterator that also enumerates dead blocks.

[springerm]

Perform a custom iteration over regions/blocks/ops. This is much simpler than trying to make ReverseDominanceIterator enumerate dead (unreachable) blocks.

[springerm]

mlir/lib/IR/PatternMatch.cpp

290

This is now fixed by iterating blocks post-order.

[mehdi_amini]

LGTM, but @rriddle should have another look

[springerm]

rebase and trigger new CI run

[springerm]

rebase

[springerm]

also drop uses of bbArg

[springerm]

Committed as: 695a5a6a66396b83263bbb3f1946fbaf41e422c3

mlir/lib/IR/PatternMatch.cpp

 void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootReplaced(op, newValues);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
-  notifyOperationRemoved(op);
+  op->walk<WalkOrder::PostOrder, BottomUpIterator>(
+      [&](Operation *op) { notifyOperationRemoved(op); });
   op->erase();
 }
 
 /// 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");
-  notifyOperationRemoved(op);
+  op->walk<WalkOrder::PostOrder, BottomUpIterator>(
+      [&](Operation *op) { notifyOperationRemoved(op); });
   op->erase();
 }
 
 void RewriterBase::eraseBlock(Block *block) {
   for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
     assert(op.use_empty() && "expected 'op' to have no uses");
     eraseOp(&op);
   }
   block->erase();
 }
 
 /// Merge the operations of block 'source' into the end of block 'dest'.
 /// 'source's predecessors must be empty or only contain 'dest`.
 /// 'argValues' is used to replace the block arguments of 'source' after
 /// merging.
 void RewriterBase::mergeBlocks(Block *source, Block *dest,
                                ValueRange argValues) {
   assert(llvm::all_of(source->getPredecessors(),
                       [dest](Block *succ) { return succ == dest; }) &&

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

LLVM_DEBUG({
             logger.startLine(),
             OpPrintingFlags().printGenericOpForm().elideLargeElementsAttrs());
         logger.getOStream() << "\n\n";
       }
     });
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
-      notifyOperationRemoved(op);
-      op->erase();
+      eraseOp(op);
       changed = true;
 
       LLVM_DEBUG(logResultWithLine("success", "operation is trivially dead"));
       continue;
     }
 
     // Collects all the operands and result uses of the given `op` into work
     // list. Also remove `op` and nested ops from worklist.
     originalOperands.assign(op->operand_begin(), op->operand_end());
     auto preReplaceAction = [&](Operation *op) {
       // Add the operands to the worklist for visitation.
       addOperandsToWorklist(originalOperands);
 
       // Add all the users of the result to the worklist so we make sure
       // to revisit them.
       for (auto result : op->getResults())
         for (auto *userOp : result.getUsers())
           addToWorklist(userOp);
 
-      notifyOperationRemoved(op);
+      op->walk<WalkOrder::PostOrder, BottomUpIterator>(
+          [&](Operation *op) { notifyOperationRemoved(op); });
     };
 
     // Add the given operation to the worklist.
     auto collectOps = [this](Operation *op) { addToWorklist(op); };
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,

 
 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
 
   addOperandsToWorklist(op->getOperands());
-  op->walk([this](Operation *operation) {
-    removeFromWorklist(operation);
-    folder.notifyRemoval(operation);
-  });
+  removeFromWorklist(op);
+  folder.notifyRemoval(op);
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
     strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op,
                                                     ValueRange replacement) {
   LLVM_DEBUG({