From 5053fd95fdc3238567fe9cdb584905dc6cb45f85 Mon Sep 17 00:00:00 2001 From: Mahesh Ravishankar Date: Wed, 4 Jan 2023 20:32:18 +0000 Subject: [PATCH] [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 --- mlir/include/mlir/IR/PatternMatch.h | 13 ++++++++++++- mlir/lib/IR/PatternMatch.cpp | 16 ++++++++-------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h index 5c7c96f..64eb66b 100644 --- a/mlir/include/mlir/IR/PatternMatch.h +++ b/mlir/include/mlir/IR/PatternMatch.h @@ -507,10 +507,21 @@ public: /// 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 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 2f1ad0c..b082b0d 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 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); }); } } -- 2.7.4