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

[vinograd47]

Summary

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

[ftynse]

This has been bothering be for some time, thanks!

[vinograd47]

Rebased to latest main branch.

[vinograd47]

@ftynse I don't have commit access to the repo. Could you please land the change by yourself?

mlir/include/mlir/IR/PatternMatch.h

public:
   /// This method performs the final replacement for a pattern, where the
   /// results of the operation are updated to use the specified list of SSA
   /// values.
   virtual void replaceOp(Operation *op, ValueRange newValues);
 
   /// 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.
   virtual void eraseOp(Operation *op);
 
   /// This method erases all operations in a block.
   virtual void eraseBlock(Block *block);