[Mlir-commits] [mlir] 9ccf613 - [mlir][Transforms][NFC] Store per-pattern IR modifications in separate state (#145319)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 25 23:54:42 PDT 2025
Author: Matthias Springer
Date: 2025-06-26T08:54:39+02:00
New Revision: 9ccf613b34263ad0fd3fb3155aeb2ac9df3df952
URL: https://github.com/llvm/llvm-project/commit/9ccf613b34263ad0fd3fb3155aeb2ac9df3df952
DIFF: https://github.com/llvm/llvm-project/commit/9ccf613b34263ad0fd3fb3155aeb2ac9df3df952.diff
LOG: [mlir][Transforms][NFC] Store per-pattern IR modifications in separate state (#145319)
This commit adds extra state to `ConversionPatternRewriterImpl` to store
all modified / newly-created operations and moved / newly-created blocks
in separate lists on a per-pattern basis.
This is in preparation of the One-Shot Dialect Conversion refactoring:
the new driver will no longer maintain a list of all IR rewrites, so
information about newly-created operations (which is needed to trigger
recursive legalization) must be retained in a different data structure.
This commit is also expected to improve the performance of the existing
driver. The previous implementation iterated over all new IR
modifications and then filtered them by type. It also required an
additional pointer indirection (through `std::unique_ptr<IRRewrite>`) to
retrieve the operation/block pointer.
Added:
Modified:
mlir/lib/Transforms/Utils/DialectConversion.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 7cfe7250d02c3..955a106c21941 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1080,6 +1080,16 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
/// to modify/access them is invalid rewriter API usage.
SetVector<Operation *> replacedOps;
+ /// A set of operations that were created by the current pattern.
+ SetVector<Operation *> patternNewOps;
+
+ /// A set of operations that were modified by the current pattern.
+ SetVector<Operation *> patternModifiedOps;
+
+ /// A set of blocks that were inserted (newly-created blocks or moved blocks)
+ /// by the current pattern.
+ SetVector<Block *> patternInsertedBlocks;
+
/// A mapping of all unresolved materializations (UnrealizedConversionCastOp)
/// to the corresponding rewrite objects.
DenseMap<UnrealizedConversionCastOp, UnresolvedMaterializationRewrite *>
@@ -1571,6 +1581,7 @@ void ConversionPatternRewriterImpl::notifyOperationInserted(
if (!previous.isSet()) {
// This is a newly created op.
appendRewrite<CreateOperationRewrite>(op);
+ patternNewOps.insert(op);
return;
}
Operation *prevOp = previous.getPoint() == previous.getBlock()->end()
@@ -1655,6 +1666,8 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
}
});
+ patternInsertedBlocks.insert(block);
+
if (!previous) {
// This is a newly created block.
appendRewrite<CreateBlockRewrite>(block);
@@ -1852,6 +1865,8 @@ void ConversionPatternRewriter::finalizeOpModification(Operation *op) {
assert(!impl->wasOpReplaced(op) &&
"attempting to modify a replaced/erased op");
PatternRewriter::finalizeOpModification(op);
+ impl->patternModifiedOps.insert(op);
+
// There is nothing to do here, we only need to track the operation at the
// start of the update.
#ifndef NDEBUG
@@ -1964,21 +1979,25 @@ class OperationLegalizer {
/// Legalize the resultant IR after successfully applying the given pattern.
LogicalResult legalizePatternResult(Operation *op, const Pattern &pattern,
ConversionPatternRewriter &rewriter,
- RewriterState &curState);
+ const SetVector<Operation *> &newOps,
+ const SetVector<Operation *> &modifiedOps,
+ const SetVector<Block *> &insertedBlocks);
/// Legalizes the actions registered during the execution of a pattern.
LogicalResult
legalizePatternBlockRewrites(Operation *op,
ConversionPatternRewriter &rewriter,
ConversionPatternRewriterImpl &impl,
- RewriterState &state, RewriterState &newState);
- LogicalResult legalizePatternCreatedOperations(
- ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
- RewriterState &state, RewriterState &newState);
- LogicalResult legalizePatternRootUpdates(ConversionPatternRewriter &rewriter,
- ConversionPatternRewriterImpl &impl,
- RewriterState &state,
- RewriterState &newState);
+ const SetVector<Block *> &insertedBlocks,
+ const SetVector<Operation *> &newOps);
+ LogicalResult
+ legalizePatternCreatedOperations(ConversionPatternRewriter &rewriter,
+ ConversionPatternRewriterImpl &impl,
+ const SetVector<Operation *> &newOps);
+ LogicalResult
+ legalizePatternRootUpdates(ConversionPatternRewriter &rewriter,
+ ConversionPatternRewriterImpl &impl,
+ const SetVector<Operation *> &modifiedOps);
//===--------------------------------------------------------------------===//
// Cost Model
@@ -2131,6 +2150,15 @@ OperationLegalizer::legalize(Operation *op,
return failure();
}
+/// Helper function that moves and returns the given object. Also resets the
+/// original object, so that it is in a valid, empty state again.
+template <typename T>
+static T moveAndReset(T &obj) {
+ T result = std::move(obj);
+ obj = T();
+ return result;
+}
+
LogicalResult
OperationLegalizer::legalizeWithFold(Operation *op,
ConversionPatternRewriter &rewriter) {
@@ -2192,6 +2220,9 @@ OperationLegalizer::legalizeWithPattern(Operation *op,
RewriterState curState = rewriterImpl.getCurrentState();
auto onFailure = [&](const Pattern &pattern) {
assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates");
+ rewriterImpl.patternNewOps.clear();
+ rewriterImpl.patternModifiedOps.clear();
+ rewriterImpl.patternInsertedBlocks.clear();
LLVM_DEBUG({
logFailure(rewriterImpl.logger, "pattern failed to match");
if (rewriterImpl.config.notifyCallback) {
@@ -2212,7 +2243,13 @@ OperationLegalizer::legalizeWithPattern(Operation *op,
// successfully applied.
auto onSuccess = [&](const Pattern &pattern) {
assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates");
- auto result = legalizePatternResult(op, pattern, rewriter, curState);
+ SetVector<Operation *> newOps = moveAndReset(rewriterImpl.patternNewOps);
+ SetVector<Operation *> modifiedOps =
+ moveAndReset(rewriterImpl.patternModifiedOps);
+ SetVector<Block *> insertedBlocks =
+ moveAndReset(rewriterImpl.patternInsertedBlocks);
+ auto result = legalizePatternResult(op, pattern, rewriter, newOps,
+ modifiedOps, insertedBlocks);
appliedPatterns.erase(&pattern);
if (failed(result)) {
if (!rewriterImpl.config.allowPatternRollback)
@@ -2253,10 +2290,11 @@ bool OperationLegalizer::canApplyPattern(Operation *op, const Pattern &pattern,
return true;
}
-LogicalResult
-OperationLegalizer::legalizePatternResult(Operation *op, const Pattern &pattern,
- ConversionPatternRewriter &rewriter,
- RewriterState &curState) {
+LogicalResult OperationLegalizer::legalizePatternResult(
+ Operation *op, const Pattern &pattern, ConversionPatternRewriter &rewriter,
+ const SetVector<Operation *> &newOps,
+ const SetVector<Operation *> &modifiedOps,
+ const SetVector<Block *> &insertedBlocks) {
auto &impl = rewriter.getImpl();
assert(impl.pendingRootUpdates.empty() && "dangling root updates");
@@ -2274,12 +2312,10 @@ OperationLegalizer::legalizePatternResult(Operation *op, const Pattern &pattern,
#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
// Legalize each of the actions registered during application.
- RewriterState newState = impl.getCurrentState();
- if (failed(legalizePatternBlockRewrites(op, rewriter, impl, curState,
- newState)) ||
- failed(legalizePatternRootUpdates(rewriter, impl, curState, newState)) ||
- failed(legalizePatternCreatedOperations(rewriter, impl, curState,
- newState))) {
+ if (failed(legalizePatternBlockRewrites(op, rewriter, impl, insertedBlocks,
+ newOps)) ||
+ failed(legalizePatternRootUpdates(rewriter, impl, modifiedOps)) ||
+ failed(legalizePatternCreatedOperations(rewriter, impl, newOps))) {
return failure();
}
@@ -2289,20 +2325,14 @@ OperationLegalizer::legalizePatternResult(Operation *op, const Pattern &pattern,
LogicalResult OperationLegalizer::legalizePatternBlockRewrites(
Operation *op, ConversionPatternRewriter &rewriter,
- ConversionPatternRewriterImpl &impl, RewriterState &state,
- RewriterState &newState) {
- SmallPtrSet<Operation *, 16> operationsToIgnore;
+ ConversionPatternRewriterImpl &impl,
+ const SetVector<Block *> &insertedBlocks,
+ const SetVector<Operation *> &newOps) {
+ SmallPtrSet<Operation *, 16> alreadyLegalized;
// If the pattern moved or created any blocks, make sure the types of block
// arguments get legalized.
- for (int i = state.numRewrites, e = newState.numRewrites; i != e; ++i) {
- BlockRewrite *rewrite = dyn_cast<BlockRewrite>(impl.rewrites[i].get());
- if (!rewrite)
- continue;
- Block *block = rewrite->getBlock();
- if (isa<BlockTypeConversionRewrite, EraseBlockRewrite,
- ReplaceBlockArgRewrite, InlineBlockRewrite>(rewrite))
- continue;
+ for (Block *block : insertedBlocks) {
// Only check blocks outside of the current operation.
Operation *parentOp = block->getParentOp();
if (!parentOp || parentOp == op || block->getNumArguments() == 0)
@@ -2322,41 +2352,26 @@ LogicalResult OperationLegalizer::legalizePatternBlockRewrites(
continue;
}
- // Otherwise, check that this operation isn't one generated by this pattern.
- // This is because we will attempt to legalize the parent operation, and
- // blocks in regions created by this pattern will already be legalized later
- // on. If we haven't built the set yet, build it now.
- if (operationsToIgnore.empty()) {
- for (unsigned i = state.numRewrites, e = impl.rewrites.size(); i != e;
- ++i) {
- auto *createOp =
- dyn_cast<CreateOperationRewrite>(impl.rewrites[i].get());
- if (!createOp)
- continue;
- operationsToIgnore.insert(createOp->getOperation());
+ // Otherwise, try to legalize the parent operation if it was not generated
+ // by this pattern. This is because we will attempt to legalize the parent
+ // operation, and blocks in regions created by this pattern will already be
+ // legalized later on.
+ if (!newOps.count(parentOp) && alreadyLegalized.insert(parentOp).second) {
+ if (failed(legalize(parentOp, rewriter))) {
+ LLVM_DEBUG(logFailure(
+ impl.logger, "operation '{0}'({1}) became illegal after rewrite",
+ parentOp->getName(), parentOp));
+ return failure();
}
}
-
- // If this operation should be considered for re-legalization, try it.
- if (operationsToIgnore.insert(parentOp).second &&
- failed(legalize(parentOp, rewriter))) {
- LLVM_DEBUG(logFailure(impl.logger,
- "operation '{0}'({1}) became illegal after rewrite",
- parentOp->getName(), parentOp));
- return failure();
- }
}
return success();
}
LogicalResult OperationLegalizer::legalizePatternCreatedOperations(
ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
- RewriterState &state, RewriterState &newState) {
- for (int i = state.numRewrites, e = newState.numRewrites; i != e; ++i) {
- auto *createOp = dyn_cast<CreateOperationRewrite>(impl.rewrites[i].get());
- if (!createOp)
- continue;
- Operation *op = createOp->getOperation();
+ const SetVector<Operation *> &newOps) {
+ for (Operation *op : newOps) {
if (failed(legalize(op, rewriter))) {
LLVM_DEBUG(logFailure(impl.logger,
"failed to legalize generated operation '{0}'({1})",
@@ -2369,12 +2384,8 @@ LogicalResult OperationLegalizer::legalizePatternCreatedOperations(
LogicalResult OperationLegalizer::legalizePatternRootUpdates(
ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl,
- RewriterState &state, RewriterState &newState) {
- for (int i = state.numRewrites, e = newState.numRewrites; i != e; ++i) {
- auto *rewrite = dyn_cast<ModifyOperationRewrite>(impl.rewrites[i].get());
- if (!rewrite)
- continue;
- Operation *op = rewrite->getOperation();
+ const SetVector<Operation *> &modifiedOps) {
+ for (Operation *op : modifiedOps) {
if (failed(legalize(op, rewriter))) {
LLVM_DEBUG(logFailure(
impl.logger, "failed to legalize operation updated in-place '{0}'",
More information about the Mlir-commits
mailing list