[mlir] Add `replaceAllUsesExcept` to rewriter

[dcaballe]

Summary

This patch adds replaceAllUsesExcept to the rewriter class.
The implementation is copy-pasted from Value + calling
updateRootInPlace to notify the listeners about the
corresponding IR changes.

[Mogball]

mlir/lib/IR/PatternMatch.cpp

327

please drop trivial braces

[rriddle]

LGTM

mlir/include/mlir/IR/PatternMatch.h

public:
   /// callable.
   template <typename CallableT>
   void updateRootInPlace(Operation *root, CallableT &&callable) {
     startRootUpdate(root);
     callable();
     finalizeRootUpdate(root);
   }
 
-  /// Find uses of `from` and replace it with `to`. It also marks every modified
-  /// uses and notifies the rewriter that an in-place operation modification is
-  /// about to happen.
+  /// 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` 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);
+
   /// 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>
   notifyMatchFailure(Location loc, CallbackT &&reasonCallback) {

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) {
+  for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
+    Operation *user = operand.getOwner();
+    if (user != exceptedUser)
+      updateRootInPlace(user, [&]() { 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() &&
          "expected 'source' to have no predecessors");
   assert(source->hasNoSuccessors() &&
          "expected 'source' to have no successors");