[mlir] Add a method to `RewriteBase` to replace a `Value` selectively.

[mravishankar]

Summary

This method allows to selectively control from the caller when to
replace the uses of a Value. Still notifies the rewriter that the
user is updated in-place.

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

521–522

I'd just merge these two lines

mlir/include/mlir/IR/PatternMatch.h

void updateRootInPlace(Operation *root, CallableT &&callable) {
     finalizeRootUpdate(root);
   }
 
   /// Find uses of `from` and replace them with `to`. It also marks every
   /// modified uses and notifies the rewriter that an in-place operation
   /// modification is about to happen.
   void replaceAllUsesWith(Value from, Value 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 replaceUseIf(Value from, Value to,
+                    llvm::unique_function<bool(OpOperand &) const> 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);
+  void replaceAllUsesExcept(Value from, Value to, Operation *exceptedUser) {
+    return replaceUseIf(from, to, [&](OpOperand &use) {
+      Operation *user = use.getOwner();
+      return user != exceptedUser;
+    });
+  }
 
   /// Used to notify the rewriter that the IR failed to be rewritten because of
   /// a match failure, and provide a callback to populate a diagnostic with the
   /// reason why the failure occurred. This method allows for derived rewriters
   /// to optionally hook into the reason why a rewrite failed, and display it to
   /// users.
   template <typename CallbackT>
   std::enable_if_t<!std::is_convertible<CallbackT, Twine>::value, LogicalResult>

mlir/lib/IR/PatternMatch.cpp

 /// Find uses of `from` and replace it with `to`
 void RewriterBase::replaceAllUsesWith(Value from, Value to) {
   for (OpOperand &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` 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 RewriterBase::replaceAllUsesExcept(Value from, Value to,
-                                        Operation *exceptedUser) {
+/// 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::replaceUseIf(
+    Value from, Value to,
+    llvm::unique_function<bool(OpOperand &) const> functor) {
   for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
-    Operation *user = operand.getOwner();
-    if (user != exceptedUser)
-      updateRootInPlace(user, [&]() { operand.set(to); });
+    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.
 void RewriterBase::mergeBlockBefore(Block *source, Operation *op,
                                     ValueRange argValues) {
   assert(source->hasNoPredecessors() &&