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 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(from.getOwner(), from, converter); mapping.map(from, to); }