[mlir][Transforms] Tighten replaceUsesOfBlockArgument (#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`.

Author: matthias-springer

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

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);
 }