Skip to content

[mlir][Transforms] Tighten replaceUsesOfBlockArgument #155227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 25, 2025
+24 −0
Merged

[mlir][Transforms] Tighten replaceUsesOfBlockArgument #155227

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions mlir/include/mlir/Transforms/DialectConversion.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions mlir/lib/Transforms/Utils/DialectConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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);
}
Expand Down