[mlir] Improve GreedyPatternRewriteDriver and pass documentation

[matthias-springer]

Clarify what kind of IR modifications are allowed. Also improve the documentation of the greedy rewrite driver entry points.

Addressing comments in #76219.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

Clarify what kind of IR modifications are allowed. Also improve the documentation of the greedy rewrite driver entry points.

---

Full diff: https://github.com/llvm/llvm-project/pull/77614.diff

2 Files Affected:

• (modified) mlir/docs/PassManagement.md (+13-8)
• (modified) mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h (+57-25)

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 &region,
                              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,

[matthias-springer]

This is a bit sus. Example:

• Pass instance 1: currentOp->setOperand(0, currentOp->getBlock()->getArgument(0)
• Pass instance 2: print(currentOp->getBlock()->getArgument(0).getUses().size()

Instance 1 modifies the current op. That is allowed. (Or is it? The use list of a parent block argument changes. Does this count as "modifying per parent op"?).

Instance 2 reads the number of uses of a parent block argument. That is allowed. But there's a race condition if both instances run in parallel.

Maybe the only modifications that we should allow are "modifying attributes"?

[joker-eph]

You can't touch operands, this would race in the parent block (it counts as modifying the parent block scope).

[matthias-springer]

Ok, that's what I thought... So the only permissible change to the operation that the pass operates on is "modifying attributes". I updates the docs accordingly.

[jpienaar]

Thanks!