[Mlir-commits] [mlir] [mlir][Transforms] Tighten `replaceUsesOfBlockArgument` (PR #155227)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/155227

Improve the documentation of `replaceUsesOfBlockArgument` to clarify its semantics is rollback mode. Add an assertion to make sure that the same block argument is not replaced multiple times. That's an API violation and messes with the internal state of the conversion driver.

This commit is in preparation of adding full support for `RewriterBase::replaceAllUsesWith`.


>From 1bb109ea1faffc7afb46e8a8c92eefb088431cbe Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Mon, 25 Aug 2025 09:32:06 +0000
Subject: [PATCH] [mlir][Transforms] Tighten `replaceUsesOfBlockArgument`

---
 .../mlir/Transforms/DialectConversion.h        |  6 ++++++
 .../lib/Transforms/Utils/DialectConversion.cpp | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 536b23f5c33c1..f23a70601fc0a 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -782,6 +782,12 @@ class ConversionPatternRewriter final : public PatternRewriter {
 
   /// Replace all the uses of the block argument `from` with `to`. This
   /// function supports both 1:1 and 1:N replacements.
+  ///
+  /// Note: If `allowPatternRollback` is set to "true", this function replaces
+  /// all current and future uses of the block argument. This same block
+  /// block argument must not be replaced multiple times. Uses are not replaced
+  /// immediately but in a delayed fashion. Patterns may still see the original
+  /// uses when inspecting IR.
   void replaceUsesOfBlockArgument(BlockArgument from, ValueRange to);
 
   /// Return the converted value of 'key' with a type defined by the type
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 7494ca9ec3784..e3248204d6694 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1132,6 +1132,11 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
   IRRewriter notifyingRewriter;
 
 #ifndef NDEBUG
+  /// A set of replaced block arguments. This set is for debugging purposes
+  /// only and it is maintained only if `allowPatternRollback` is set to
+  /// "true".
+  DenseSet<BlockArgument> replacedArgs;
+
   /// A set of operations that have pending updates. This tracking isn't
   /// strictly necessary, and is thus only active during debug builds for extra
   /// verification.
@@ -1892,6 +1897,19 @@ void ConversionPatternRewriterImpl::replaceUsesOfBlockArgument(
     return;
   }
 
+#ifndef NDEBUG
+  // Make sure that a block argument is not replaced multiple times. In
+  // rollback mode, `replaceUsesOfBlockArgument` replaces not only all current
+  // uses of the given block argument, but also all future uses that may be
+  // introduced by future pattern applications. Therefore, it does not make
+  // sense to call `replaceUsesOfBlockArgument` multiple times with the same
+  // block argument. Doing so would overwrite the mapping and mess with the
+  // internal state of the dialect conversion driver.
+  assert(!replacedArgs.contains(from) &&
+         "attempting to replace a block argument that was already replaced");
+  replacedArgs.insert(from);
+#endif // NDEBUG
+
   appendRewrite<ReplaceBlockArgRewrite>(from.getOwner(), from, converter);
   mapping.map(from, to);
 }