[mlir] specify the values when notifying about op replacement

[ftynse]

Summary

It is useful for PatternRewriter listeners to know the values that are
replacing the op in addition to only the fact of the op being replaced
for being able to keep track of changes or for debugging.

mlir/include/mlir/IR/PatternMatch.h

explicit RewriterBase(const OpBuilder &otherBuilder)
     setListener(this);
   }
   ~RewriterBase() override;
 
   /// These are the callback methods that subclasses can choose to implement if
   /// they would like to be notified about certain types of mutations.
 
   /// Notify the rewriter that the specified operation is about to be replaced
-  /// with another set of operations. This is called before the uses of the
-  /// operation have been changed.
-  virtual void notifyRootReplaced(Operation *op) {}
+  /// with the set of values potentially produced by new operations. This is
+  /// called before the uses of the operation have been changed.
+  virtual void notifyRootReplaced(Operation *op, ValueRange replacement) {}
 
   /// This is called on an operation that a rewrite is removing, right before
   /// the operation is deleted. At this point, the operation has zero uses.
   virtual void notifyOperationRemoved(Operation *op) {}
 
   /// Notify the rewriter that the pattern failed to match the given operation,
   /// and provide a callback to populate a diagnostic with the reason why the
   /// failure occurred. This method allows for derived rewriters to optionally

mlir/lib/IR/PatternMatch.cpp

 /// 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.
-  notifyRootReplaced(op);
+  notifyRootReplaced(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);
     replacedAllUses &= std::get<0>(it).use_empty();
   }
   if (allUsesReplaced)
void RewriterBase::replaceOpWithinBlock(Operation *op, ValueRange newValues,
   });
 }
 
 /// 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.
-  notifyRootReplaced(op);
+  notifyRootReplaced(op, newValues);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
   notifyOperationRemoved(op);
   op->erase();
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

protected:
 
   // If an operation is about to be removed, make sure it is not in our
   // worklist anymore because we'd get dangling references to it.
   void notifyOperationRemoved(Operation *op) override;
 
   // When the root of a pattern is about to be replaced, it can trigger
   // simplifications to its users - make sure to add them to the worklist
   // before the root is changed.
-  void notifyRootReplaced(Operation *op) override;
+  void notifyRootReplaced(Operation *op, ValueRange replacement) override;
 
   /// PatternRewriter hook for erasing a dead operation.
   void eraseOp(Operation *op) override;
 
   /// PatternRewriter hook for notifying match failure reasons.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) override;

 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
 }
 
-void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op) {
+void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op,
+                                                    ValueRange replacement) {
   LLVM_DEBUG({
     logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   for (auto result : op->getResults())
     for (auto *user : result.getUsers())
       addToWorklist(user);
 }
protected:
   /// If an operation is about to be removed, mark it so that we can let clients
   /// know.
   void notifyOperationRemoved(Operation *op) override {
     opErasedViaPatternRewrites = true;
   }
 
   // When a root is going to be replaced, its removal will be notified as well.
   // So there is nothing to do here.
-  void notifyRootReplaced(Operation *op) override {}
+  void notifyRootReplaced(Operation *op, ValueRange replacement) override {}
 
 private:
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;