diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h --- a/mlir/include/mlir/IR/PatternMatch.h +++ b/mlir/include/mlir/IR/PatternMatch.h @@ -502,11 +502,16 @@ finalizeRootUpdate(root); } - /// Find uses of `from` and replace it with `to`. It also marks every modified + /// 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 --- a/mlir/lib/IR/PatternMatch.cpp +++ b/mlir/lib/IR/PatternMatch.cpp @@ -317,6 +317,19 @@ } } +/// 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,