[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

Author: Vladislav Vinogradov

mlir/include/mlir/IR/PatternMatch.h

@@ -480,9 +480,10 @@ class PatternRewriter : public OpBuilder, public OpBuilder::Listener {
   /// 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.