[mlir] Return new Operation from `Rewriter::replaceOpWithNewOp`

It will allow to perform additional manipulation with the newly created Operation.
For example, custom attributes propagation/changes.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D95525

diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 1a306e6..04904da 100644 (file)
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -480,9 +480,10 @@ public:
   /// Replaces the result op with a new op that is created without verification.
   /// The result values of the two ops must be the same types.
   template <typename OpTy, typename... Args>
-  void replaceOpWithNewOp(Operation *op, Args &&... args) {
+  OpTy replaceOpWithNewOp(Operation *op, Args &&... args) {
     auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
     replaceOpWithResultsOfAnotherOp(op, newOp.getOperation());
+    return newOp;
   }
 
   /// This method erases an operation that is known to have no uses.