[Mlir-commits] [mlir] 5053fd9 - [mlir] Add a method to `RewriteBase` to replace a `Value` selectively.

[Mahesh Ravishankar]

Author: Mahesh Ravishankar
Date: 2023-01-16T05:03:40Z
New Revision: 5053fd95fdc3238567fe9cdb584905dc6cb45f85

URL: https://github.com/llvm/llvm-project/commit/5053fd95fdc3238567fe9cdb584905dc6cb45f85
DIFF: https://github.com/llvm/llvm-project/commit/5053fd95fdc3238567fe9cdb584905dc6cb45f85.diff

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

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.

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

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 5c7c96f5761bd..64eb66bf53b73 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -507,10 +507,21 @@ class RewriterBase : public OpBuilder, public OpBuilder::Listener {
   /// 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

diff  --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 2f1ad0c469834..b082b0d4cd6ef 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -317,15 +317,15 @@ 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) {
+/// 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); });
   }
 }