[mlir][IR] Trigger notifyOperationModified for replacements

[springerm]

Summary

Each user of the original value is modified in-place. Therefore, the corresponding notification should be triggered.

Also fixes a bug where RewriterBase::mergeBlocks did not notify the GreedyPatternRewriteDriver when replacing uses of block arguments. This function does not trigger "operation replaced" notifications, so the GreedyPatternRewriteDriver was not made aware of such IR changes.

Depends On: D144193

[springerm]

update

[Mogball]

please add a test

[springerm]

address comments: add test case

mlir/include/mlir/IR/PatternMatch.h

for (OperandType &operand : llvm::make_early_inc_range(from->getUses())) {
       Operation *op = operand.getOwner();
       updateRootInPlace(op, [&]() { operand.set(to); });
     }
   }
 
   /// Find uses of `from` and replace them with `to` if the `functor` returns
   /// true. It also marks every modified uses and notifies the rewriter that an
   /// in-place operation modification is about to happen.
-  void
-  replaceUsesWithIf(Value from, Value to,
-                    llvm::unique_function<bool(OpOperand &) const> functor);
+  void replaceUsesWithIf(Value from, Value to,
+                         function_ref<bool(OpOperand &)> functor);
 
   /// Find uses of `from` and replace them with `to` except if the user is
   /// `exceptedUser`. It also marks every modified uses and notifies the
   /// rewriter that an in-place operation modification is about to happen.
   void replaceAllUsesExcept(Value from, Value to, Operation *exceptedUser) {
     return replaceUsesWithIf(from, to, [&](OpOperand &use) {
       Operation *user = use.getOwner();
       return user != exceptedUser;

mlir/lib/IR/PatternMatch.cpp

 /// The number of values in `newValues` is required to match the number of
 /// results of `op`.
 void RewriterBase::replaceOpWithIf(
     Operation *op, ValueRange newValues, bool *allUsesReplaced,
     llvm::unique_function<bool(OpOperand &) const> functor) {
   assert(op->getNumResults() == newValues.size() &&
          "incorrect number of values to replace operation");
 
-  // Notify the rewriter subclass that we're about to replace this root.
+  // Notify the listener that we're about to replace this op.
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
     rewriteListener->notifyOperationReplaced(op, newValues);
 
   // Replace each use of the results when the functor is true.
   bool replacedAllUses = true;
   for (auto it : llvm::zip(op->getResults(), newValues)) {
-    std::get<0>(it).replaceUsesWithIf(std::get<1>(it), functor);
+    replaceUsesWithIf(std::get<0>(it), std::get<1>(it), functor);
     replacedAllUses &= std::get<0>(it).use_empty();
   }
   if (allUsesReplaced)
     *allUsesReplaced = replacedAllUses;
 }
 
 /// This method replaces the uses of the results of `op` with the values in
 /// `newValues` when a use is nested within the given `block`. The number of
 /// values in `newValues` is required to match the number of results of `op`.
 /// If all uses of this operation are replaced, the operation is erased.
 void RewriterBase::replaceOpWithinBlock(Operation *op, ValueRange newValues,
                                         Block *block, bool *allUsesReplaced) {
   replaceOpWithIf(op, newValues, allUsesReplaced, [block](OpOperand &use) {
     return block->getParentOp()->isProperAncestor(use.getOwner());
   });
 }
 
 /// This method replaces the results of the operation with the specified list of
 /// values. The number of provided values must match the number of results of
 /// the operation.
 void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
-  // Notify the rewriter subclass that we're about to replace this root.
-  if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
-    rewriteListener->notifyOperationReplaced(op, newValues);
-
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
-  op->replaceAllUsesWith(newValues);
 
+  // Notify the listener that we're about to remove this op.
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
-    rewriteListener->notifyOperationRemoved(op);
-  op->erase();
+    rewriteListener->notifyOperationReplaced(op, newValues);
+
+  // Replace results one-by-one. Also notifies the listener of modifications.
+  for (auto it : llvm::zip(op->getResults(), newValues))
+    replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
+
+  // Erase the op. Also notifies the listener.
+  eraseOp(op);
 }
 
 /// 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");
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
     rewriteListener->notifyOperationRemoved(op);
void RewriterBase::mergeBlocks(Block *source, Block *dest,
   assert(llvm::all_of(source->getPredecessors(),
                       [dest](Block *succ) { return succ == dest; }) &&
          "expected 'source' to have no predecessors or only 'dest'");
   assert(argValues.size() == source->getNumArguments() &&
          "incorrect # of argument replacement values");
 
   // Replace all of the successor arguments with the provided values.
   for (auto it : llvm::zip(source->getArguments(), argValues))
-    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
+    replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
 
   // Splice the operations of the 'source' block into the 'dest' block and erase
   // it.
   dest->getOperations().splice(dest->end(), source->getOperations());
   source->dropAllUses();
   source->erase();
 }
 
 /// Find uses of `from` and replace them with `to` if the `functor` returns
 /// true. It also marks every modified uses and notifies the rewriter that an
 /// in-place operation modification is about to happen.
-void RewriterBase::replaceUsesWithIf(
-    Value from, Value to,
-    llvm::unique_function<bool(OpOperand &) const> functor) {
+void RewriterBase::replaceUsesWithIf(Value from, Value to,
+                                     function_ref<bool(OpOperand &)> functor) {
   for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
     if (functor(operand))
       updateRootInPlace(operand.getOwner(), [&]() { operand.set(to); });
   }
 }
 
 // Merge the operations of block 'source' before the operation 'op'. Source
 // block should not have existing predecessors or successors.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 void GreedyPatternRewriteDriver::notifyOperationReplaced(
     Operation *op, ValueRange replacement) {
   LLVM_DEBUG({
     logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   if (config.listener)
     config.listener->notifyOperationReplaced(op, replacement);
-  for (auto result : op->getResults())
-    for (auto *user : result.getUsers())
-      addToWorklist(user);
 }
 
 LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";