[Mlir-commits] [mlir] [mlir][Transforms] `GreedyPatternRewriteDriver`: Check for out-of-scope IR modifications (PR #76219)
Matthias Springer
llvmlistbot at llvm.org
Wed Jan 10 09:06:55 PST 2024
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/76219
>From 2b3ec65b4c5d2eca49fd79bb277e482fa2fb80f6 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 10 Jan 2024 14:49:45 +0000
Subject: [PATCH 1/2] [mlir] Improve `GreedyPatternRewriteDriver` and pass
documentation
Clarify what kind of IR modifications are allowed. Also improve the documentation of the greedy rewrite driver entry points.
---
mlir/docs/PassManagement.md | 21 +++--
.../Transforms/GreedyPatternRewriteDriver.h | 82 +++++++++++++------
2 files changed, 70 insertions(+), 33 deletions(-)
diff --git a/mlir/docs/PassManagement.md b/mlir/docs/PassManagement.md
index 95a38207b7f854..ebe1be1e55cbb5 100644
--- a/mlir/docs/PassManagement.md
+++ b/mlir/docs/PassManagement.md
@@ -23,16 +23,21 @@ further below. All passes in MLIR derive from `OperationPass` and adhere to the
following restrictions; any noncompliance will lead to problematic behavior in
multithreaded and other advanced scenarios:
-* Must not modify any state referenced or relied upon outside the current
- operation being operated on. This includes adding or removing operations
- from the parent block, changing the attributes(depending on the contract
- of the current operation)/operands/results/successors of the current operation.
-* Must not modify the state of another operation not nested within the current
- operation being operated on.
- * Other threads may be operating on these operations simultaneously.
-* Must not inspect the state of sibling operations.
+* Must not inspect the state of operations that are siblings of the operation
+ that the pass operates on. Must neither access operations nested under those
+ siblings.
* Other threads may be modifying these operations in parallel.
* Inspecting the state of ancestor/parent operations is permitted.
+* Must not modify the state of operations other than the operation that the
+ pass operates on ("current operation") and its nested operations. This
+ includes adding, modifying or removing other operations from an ancestor
+ block.
+ * Other threads may be operating on these operations simultaneously.
+ * The attributes of the current operation may be modified freely.
+ * The operands of the current operation may be modified, as long as no
+ new operations are added to an ancestor block and no sibling operations
+ are accessed. (I.e., operands can only be changed to values defined by
+ ancestors.)
* Must not maintain mutable pass state across invocations of `runOnOperation`.
A pass may be run on many different operations with no guarantee of
execution order.
diff --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index b93ffd96bee5fa..763146aac15b9c 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -62,7 +62,8 @@ class GreedyRewriteConfig {
/// Only ops within the scope are added to the worklist. If no scope is
/// specified, the closest enclosing region around the initial list of ops
- /// is used as a scope.
+ /// (or the specified region, depending on which greedy rewrite entry point
+ /// is used) is used as a scope.
Region *scope = nullptr;
/// Strict mode can restrict the ops that are added to the worklist during
@@ -86,30 +87,54 @@ class GreedyRewriteConfig {
//===----------------------------------------------------------------------===//
/// Rewrite ops in the given region, which must be isolated from above, by
-/// repeatedly applying the highest benefit patterns in a greedy work-list
-/// driven manner.
+/// repeatedly applying the highest benefit patterns in a greedy worklist
+/// driven manner until a fixpoint is reached.
///
-/// This variant may stop after a predefined number of iterations, see the
-/// alternative below to provide a specific number of iterations before stopping
-/// in absence of convergence.
+/// The greedy rewrite may prematurely stop after a maximum number of
+/// iterations, which can be configured in the configuration parameter.
///
-/// Return success if the iterative process converged and no more patterns can
-/// be matched in the result operation regions. `changed` is set to true if the
-/// IR was modified at all.
+/// Also performs folding and simple dead-code elimination before attempting to
+/// match any of the provided patterns.
///
-/// Note: This does not apply patterns to the top-level operation itself.
-/// These methods also perform folding and simple dead-code elimination
-/// before attempting to match any of the provided patterns.
+/// A region scope can be set in the configuration parameter. By default, the
+/// scope is set to the specified region. Only in-scope ops are added to the
+/// worklist and only in-scope ops are allowed to be modified by the patterns.
///
-/// You may configure several aspects of this with GreedyRewriteConfig.
+/// Returns "success" if the iterative process converged (i.e., fixpoint was
+/// reached) and no more patterns can be matched within the region. `changed`
+/// is set to "true" if the IR was modified at all.
+///
+/// Note: This method does not apply patterns to the region's parent operation.
LogicalResult
applyPatternsAndFoldGreedily(Region ®ion,
const FrozenRewritePatternSet &patterns,
GreedyRewriteConfig config = GreedyRewriteConfig(),
bool *changed = nullptr);
-/// Rewrite ops in all regions of the given op, which must be isolated from
-/// above.
+/// Rewrite ops nested under the given operation, which must be isolated from
+/// above, by repeatedly applying the highest benefit patterns in a greedy
+/// worklist driven manner until a fixpoint is reached.
+///
+/// The greedy rewrite may prematurely stop after a maximum number of
+/// iterations, which can be configured in the configuration parameter.
+///
+/// Also performs folding and simple dead-code elimination before attempting to
+/// match any of the provided patterns.
+///
+/// This overload runs a separate greedy rewrite for each region of the
+/// specified op. A region scope can be set in the configuration parameter. By
+/// default, the scope is set to the region of the current greedy rewrite. Only
+/// in-scope ops are added to the worklist and only in-scope ops and the
+/// specified op itself are allowed to be modified by the patterns.
+///
+/// Note: The specified op may be modified, but it may not be removed by the
+/// patterns.
+///
+/// Returns "success" if the iterative process converged (i.e., fixpoint was
+/// reached) and no more patterns can be matched within the region. `changed`
+/// is set to "true" if the IR was modified at all.
+///
+/// Note: This method does not apply patterns to the given operation itself.
inline LogicalResult
applyPatternsAndFoldGreedily(Operation *op,
const FrozenRewritePatternSet &patterns,
@@ -129,27 +154,34 @@ applyPatternsAndFoldGreedily(Operation *op,
return failure(failed);
}
-/// Applies the specified rewrite patterns on `ops` while also trying to fold
-/// these ops.
+/// Rewrite the specified ops by repeatedly applying the highest benefit
+/// patterns in a greedy worklist driven manner until a fixpoint is reached.
+///
+/// The greedy rewrite may prematurely stop after a maximum number of
+/// iterations, which can be configured in the configuration parameter.
+///
+/// Also performs folding and simple dead-code elimination before attempting to
+/// match any of the provided patterns.
///
/// Newly created ops and other pre-existing ops that use results of rewritten
-/// ops or supply operands to such ops are simplified, unless such ops are
+/// ops or supply operands to such ops are also processed, unless such ops are
/// excluded via `config.strictMode`. Any other ops remain unmodified (i.e.,
/// regardless of `strictMode`).
///
/// In addition to strictness, a region scope can be specified. Only ops within
/// the scope are simplified. This is similar to `applyPatternsAndFoldGreedily`,
-/// where only ops within the given regions are simplified. If no scope is
-/// specified, it is assumed to be the first common enclosing region of the
-/// given ops.
+/// where only ops within the given region/op are simplified by default. If no
+/// scope is specified, it is assumed to be the first common enclosing region of
+/// the given ops.
///
/// Note that ops in `ops` could be erased as result of folding, becoming dead,
/// or via pattern rewrites. If more far reaching simplification is desired,
-/// applyPatternsAndFoldGreedily should be used.
+/// `applyPatternsAndFoldGreedily` should be used.
///
-/// Returns success if the iterative process converged and no more patterns can
-/// be matched. `changed` is set to true if the IR was modified at all.
-/// `allOpsErased` is set to true if all ops in `ops` were erased.
+/// Returns "success" if the iterative process converged (i.e., fixpoint was
+/// reached) and no more patterns can be matched. `changed` is set to "true" if
+/// the IR was modified at all. `allOpsErased` is set to "true" if all ops in
+/// `ops` were erased.
LogicalResult
applyOpPatternsAndFold(ArrayRef<Operation *> ops,
const FrozenRewritePatternSet &patterns,
>From 0f649fd3c11600787f66720bce667162716c1136 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 10 Jan 2024 17:05:22 +0000
Subject: [PATCH 2/2] [mlir][Transforms] `GreedyPatternRewriteDriver`: Check
for out-of-scope IR modifications
This commit adds an additional "expensive check" (only enabled with `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`) that looks for out-of-scope IR modifications.
`GreedyRewriteConfig::scope` specifies the `Region *` within which the greedy pattern rewrite operates. Operations that are out-of-scope are not added to the worklist. The new expensive check triggers an fatal error if:
* Op is inserted into out-of-scope region.
* Op is removed from out-of-scope region.
* Op is modified in out-of-scope region.
This change also tightens the greedy pattern rewriter entry points and makes sure that the specified `scope` is an `IsolatedFromAbove` region.
Note: `TileAllocation` (`ArmSME` dialect) must now be a module pass because it modifies `func.func` ops (adds attributes). This is forbidden for function passes (in which the scope of the greedy rewrite is set to the region of the function by default) because only function bodies are allowed to be modified. (TODO: Should we allow this? Is there something special about functions?)
---
.../Transforms/GreedyPatternRewriteDriver.h | 35 +++++-
.../Utils/GreedyPatternRewriteDriver.cpp | 107 +++++++++++++++---
2 files changed, 120 insertions(+), 22 deletions(-)
diff --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index 763146aac15b9c..f22c2c462bdee4 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -60,12 +60,30 @@ class GreedyRewriteConfig {
static constexpr int64_t kNoLimit = -1;
- /// Only ops within the scope are added to the worklist. If no scope is
- /// specified, the closest enclosing region around the initial list of ops
- /// (or the specified region, depending on which greedy rewrite entry point
- /// is used) is used as a scope.
+ /// Only ops within the scope are allowed to be modified and are added to the
+ /// worklist. If out-of-scope IR is modified, an assertion will fail inside
+ /// the greedy pattern rewrite driver if expensive checks are enabled (as long
+ /// as rewrite patterns use the rewriter API correctly).
+ ///
+ /// The scope region must be isolated from above. This ensures that
+ /// out-of-scope ops are not affected by rewrites.
+ ///
+ /// The op that owns the scope region must be isolated from above. If no scope
+ /// is specified, it is set as follows:
+ /// * Single op greedy rewrite: a greedy rewrite is performed for every region
+ /// of the op. See below. (The op must be isolated from above.)
+ /// * Multi op greedy rewrite: the closest enclosing IsolatedFromAbove region
+ /// around the initial list of ops. If there is no such region, the scope
+ /// is `nullptr`. This is because multi-op greedy rewrites are allowed to
+ /// modify top-level ops. (They are not allowed to erase top-level ops.)
+ /// * Single region greedy rewrite: the specified region. (The op that owns
+ /// the region must be isolated from above.)
Region *scope = nullptr;
+ /// By default, only in-scope ops may be modified. If this flag is set to
+ /// "true", the `scope.getParentOp()` may also be modified (but not removed).
+ bool allowScopeParentOpModifications = false;
+
/// Strict mode can restrict the ops that are added to the worklist during
/// the rewrite.
///
@@ -140,6 +158,15 @@ applyPatternsAndFoldGreedily(Operation *op,
const FrozenRewritePatternSet &patterns,
GreedyRewriteConfig config = GreedyRewriteConfig(),
bool *changed = nullptr) {
+ // If no explicit scope is specified, it is set to the respective region
+ // (per each loop iteration below). In that case, this greedy pattern rewriter
+ // entry point also allows modifications to `op`.
+ //
+ // Note: The region-based entry point does not allow modifications to
+ // `region.getParentOp()`.
+ if (!config.scope)
+ config.allowScopeParentOpModifications = true;
+
bool anyRegionChanged = false;
bool failed = false;
for (Region ®ion : op->getRegions()) {
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 67c2d9d59f4c92..cca99ab9d5acd4 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -324,6 +324,11 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
private:
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ /// Return "true" if the given op is guaranteed to be out of scope.
+ bool isOutOfScope(Operation *op) const;
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
/// Look over the provided operands for any defining operations that should
/// be re-added to the worklist. This function should be called when an
/// operation is modified or removed, as it may trigger further
@@ -375,6 +380,28 @@ GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
}
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+bool GreedyPatternRewriteDriver::isOutOfScope(Operation *op) const {
+ // No op is out of scope if no scope was set.
+ if (!config.scope)
+ return false;
+ // Check if the given op and the scope region are part of the same IR tree.
+ // The parent op into which the given op was inserted may be unlinked, in
+ // which case we do not consider the given op to be out of scope. (That parent
+ // op will likely be inserted later, together with all its nested ops.)
+ Region *r = config.scope;
+ while (r) {
+ if (r->findAncestorOpInRegion(*op) || r->getParentOp() == op)
+ break;
+ r = r->getParentRegion();
+ }
+ if (!r)
+ return false;
+ // Op is out of scope if it is not within the scope region.
+ return !config.scope->findAncestorOpInRegion(*op);
+}
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
bool GreedyPatternRewriteDriver::processWorklist() {
#ifndef NDEBUG
const char *logLineComment =
@@ -579,6 +606,8 @@ void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
addSingleOpToWorklist(op);
return;
}
+ // TODO: Unlinked ops are currently not added to the worklist if a `scope`
+ // is specified.
if (region == nullptr)
return;
} while ((op = region->getParentOp()));
@@ -600,6 +629,13 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
logger.startLine() << "** Insert : '" << op->getName() << "'(" << op
<< ")\n";
});
+
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ if (config.scope && isOutOfScope(op))
+ llvm::report_fatal_error(
+ "greedy pattern rewrite inserted op into region that is out of scope");
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
if (config.listener)
config.listener->notifyOperationInserted(op);
if (config.strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
@@ -608,10 +644,28 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
}
void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
+ // TODO: This notification should also be triggered when moving an op into
+ // this op.
LLVM_DEBUG({
logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
<< ")\n";
});
+
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ if (config.scope) {
+ if (op == config.scope.getParentOp()) {
+ if (!config.allowScopeParentOpModifications) {
+ llvm::report_fatal_error(
+ "greedy pattern rewrite modified scope parent op, but scope parent "
+ "modifications are forbidden");
+ }
+ } else if (isOutOfScope(op)) {
+ llvm::report_fatal_error("greedy pattern rewrite modified op within "
+ "region that is out of scope");
+ }
+ }
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
if (config.listener)
config.listener->notifyOperationModified(op);
addToWorklist(op);
@@ -637,16 +691,11 @@ void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
<< ")\n";
});
-#ifndef NDEBUG
- // Only ops that are within the configured scope are added to the worklist of
- // the greedy pattern rewriter. Moreover, the parent op of the scope region is
- // the part of the IR that is taken into account for the "expensive checks".
- // A greedy pattern rewrite is not allowed to erase the parent op of the scope
- // region, as that would break the worklist handling and the expensive checks.
- if (config.scope && config.scope->getParentOp() == op)
- llvm_unreachable(
- "scope region must not be erased during greedy pattern rewrite");
-#endif // NDEBUG
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ if (config.scope && isOutOfScope(op))
+ llvm::report_fatal_error(
+ "greedy pattern rewrite removed op from region that is out of scope");
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
if (config.listener)
config.listener->notifyOperationRemoved(op);
@@ -800,16 +849,22 @@ LogicalResult
mlir::applyPatternsAndFoldGreedily(Region ®ion,
const FrozenRewritePatternSet &patterns,
GreedyRewriteConfig config, bool *changed) {
- // The top-level operation must be known to be isolated from above to
- // prevent performing canonicalizations on operations defined at or above
- // the region containing 'op'.
- assert(region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- "patterns can only be applied to operations IsolatedFromAbove");
-
// Set scope if not specified.
if (!config.scope)
config.scope = ®ion;
+ // Make sure that the specified region on which the greedy rewrite should
+ // operate is in scope.
+ assert(config.scope->isAncestor(®ion) && "input region must be in scope");
+
+ // The scope of a greedy pattern rewrite must be IsolatedFromAbove. Ops that
+ // are out of scope are never added to the worklist and any out-of-scope IR
+ // modifications trigger an assertion when expensive expensive checks are
+ // enabled (as long as the rewriter API is used correctly).
+ assert(
+ config.scope->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+ "greedy pattern rewrite scope must be IsolatedFromAbove");
+
#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
if (failed(verify(config.scope->getParentOp())))
llvm::report_fatal_error(
@@ -886,7 +941,8 @@ LogicalResult MultiOpPatternRewriteDriver::simplify(ArrayRef<Operation *> ops,
return success(worklist.empty());
}
-/// Find the region that is the closest common ancestor of all given ops.
+/// Find the IsolateFromAbove region that is the closest common ancestor of all
+/// given ops.
///
/// Note: This function returns `nullptr` if there is a top-level op among the
/// given list of ops.
@@ -896,6 +952,7 @@ static Region *findCommonAncestor(ArrayRef<Operation *> ops) {
if (ops.size() == 1)
return ops.front()->getParentRegion();
+ // Find the closest region that contains all ops.
Region *region = ops.front()->getParentRegion();
ops = ops.drop_front();
int sz = ops.size();
@@ -912,6 +969,12 @@ static Region *findCommonAncestor(ArrayRef<Operation *> ops) {
break;
region = region->getParentRegion();
}
+
+ // Find the closest IsolatedFromAbove region.
+ while (region &&
+ !region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
+ region = region->getParentRegion();
+
return region;
}
@@ -932,8 +995,16 @@ LogicalResult mlir::applyOpPatternsAndFold(
// there is a top-level op among `ops`.
config.scope = findCommonAncestor(ops);
} else {
- // If a scope was provided, make sure that all ops are in scope.
+ // If a scope was provided, make sure that it is IsolatedFromAbove and that
+ // all ops are in scope.
#ifndef NDEBUG
+ // The scope of a greedy pattern rewrite must be IsolatedFromAbove. Ops that
+ // are out of scope are never added to the worklist and any out-of-scope IR
+ // modifications trigger an assertion when expensive expensive checks are
+ // enabled (as long as the rewriter API is used correctly).
+ assert(
+ config.scope->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+ "greedy pattern rewrite scope must be IsolatedFromAbove");
bool allOpsInScope = llvm::all_of(ops, [&](Operation *op) {
return static_cast<bool>(config.scope->findAncestorOpInRegion(*op));
});
More information about the Mlir-commits
mailing list