[Mlir-commits] [mlir] 77603e2 - [mlir] Add `replaceAllUsesExcept` to rewriter

[Diego Caballero]

Author: Diego Caballero
Date: 2022-12-06T07:42:15Z
New Revision: 77603e28ce05d3a70bcf7fb2880e7b0c6c547810

URL: https://github.com/llvm/llvm-project/commit/77603e28ce05d3a70bcf7fb2880e7b0c6c547810
DIFF: https://github.com/llvm/llvm-project/commit/77603e28ce05d3a70bcf7fb2880e7b0c6c547810.diff

LOG: [mlir] Add `replaceAllUsesExcept` to rewriter

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.

Reviewed By: Mogball

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/PatternMatch.h
    mlir/lib/IR/PatternMatch.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index e548dc9e708bc..7790f96f6bd92 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -502,11 +502,16 @@ class RewriterBase : public OpBuilder, public OpBuilder::Listener {
     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

diff  --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index d3072b506c027..1a2b234c3a19c 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -317,6 +317,18 @@ void RewriterBase::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 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,