Skip to content

[mlir][Transforms] Track erased ops separately #83051

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
Feb 26, 2024
Merged

[mlir][Transforms] Track erased ops separately #83051

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
23 changes: 17 additions & 6 deletions mlir/lib/Transforms/Utils/DialectConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,9 +153,9 @@ namespace {
/// This is useful when saving and undoing a set of rewrites.
struct RewriterState {
RewriterState(unsigned numRewrites, unsigned numIgnoredOperations,
unsigned numErased)
unsigned numErased, unsigned numReplacedOps)
: numRewrites(numRewrites), numIgnoredOperations(numIgnoredOperations),
numErased(numErased) {}
numErased(numErased), numReplacedOps(numReplacedOps) {}

/// The current number of rewrites performed.
unsigned numRewrites;
Expand All @@ -165,6 +165,9 @@ struct RewriterState {

/// The current number of erased operations/blocks.
unsigned numErased;

/// The current number of replaced ops that are scheduled for erasure.
unsigned numReplacedOps;
};

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -954,6 +957,9 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
/// operation was ignored.
SetVector<Operation *> ignoredOps;

// A set of operations that were erased.
SetVector<Operation *> replacedOps;

/// The current type converter, or nullptr if no type converter is currently
/// active.
const TypeConverter *currentTypeConverter = nullptr;
Expand Down Expand Up @@ -1152,7 +1158,7 @@ void ConversionPatternRewriterImpl::applyRewrites() {

RewriterState ConversionPatternRewriterImpl::getCurrentState() {
return RewriterState(rewrites.size(), ignoredOps.size(),
eraseRewriter.erased.size());
eraseRewriter.erased.size(), replacedOps.size());
}

void ConversionPatternRewriterImpl::resetState(RewriterState state) {
Expand All @@ -1165,6 +1171,9 @@ void ConversionPatternRewriterImpl::resetState(RewriterState state) {

while (eraseRewriter.erased.size() != state.numErased)
eraseRewriter.erased.pop_back();

while (replacedOps.size() != state.numReplacedOps)
replacedOps.pop_back();
}

void ConversionPatternRewriterImpl::undoRewrites(unsigned numRewritesToKeep) {
Expand Down Expand Up @@ -1228,9 +1237,11 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
return success();
}

// TODO: This function is a misnomer. It does not actually check if `op` is in
// `ignoredOps`.
bool ConversionPatternRewriterImpl::isOpIgnored(Operation *op) const {
// Check to see if this operation or the parent operation is ignored.
return ignoredOps.count(op->getParentOp()) || ignoredOps.count(op);
return ignoredOps.count(op->getParentOp()) || replacedOps.count(op);
}

void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) {
Expand Down Expand Up @@ -1479,7 +1490,7 @@ void ConversionPatternRewriterImpl::notifyOperationInserted(
void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
ValueRange newValues) {
assert(newValues.size() == op->getNumResults());
assert(!ignoredOps.contains(op) && "operation was already replaced");
assert(!replacedOps.contains(op) && "operation was already replaced");

// Track if any of the results changed, e.g. erased and replaced with null.
bool resultChanged = false;
Expand All @@ -1500,7 +1511,7 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,

// Mark this operation as recursively ignored so that we don't need to
// convert any nested operations.
ignoredOps.insert(op);
replacedOps.insert(op);
markNestedOpsIgnored(op);
}

Expand Down