[mlir] GreedyPatternRewriter: Relax IsolatedFromAbove restriction

[springerm]

Summary

A scope was added to the GreedyPatternRewriteDriver in D140304. This scope prevents ops from being added to the worklist if they are not inside regions.

The reason why IsIsolatedFromAbove was necessary:

// Collects all the operands and result uses of the given `op` into work
// list. Also remove `op` and nested ops from worklist.
originalOperands.assign(op->operand_begin(), op->operand_end());
auto preReplaceAction = [&](Operation *op) {
  // Add the operands to the worklist for visitation.
  addOperandsToWorklist(originalOperands);
  // ...
}

Depends On: D140304

[nicolasvasilache]

Fantastic, this will greatly improve usability of the rewriter.

[mehdi_amini]

The reason why IsIsolatedFromAbove was necessary:

This isn't the only reason: any pattern operates on use/def lists and as such isn't restricted to the current region if not isolated.

[springerm]

In D141366#4050941, @mehdi_amini wrote:

The reason why IsIsolatedFromAbove was necessary:

This isn't the only reason: any pattern operates on use/def lists and as such isn't restricted to the current region if not isolated.

My thinking was: The number of uses of a value defined outside of the region may change in number, but the ops themselves are not modified in any way. Values defined inside the region cannot be used outside of the region.

Example: We call applyPatternsAndFoldGreedily on the region of the scf.execute_region.

%0 = "outside_op"() : () -> (i32)
%1 = scf.execute_region -> (i32) {
  %2 = "inside_op"(%0) : (i32) -> (i32)
  scf.yield %2 : i32
}

We may rewrite/replace/erase inside_op and %0 may loose one of its users. But outside_op is not modified in any way. If it is "pure" op, it's possible that it could now DCE away, but that's not something that applyPatternsAndFoldGreedily would do.

[springerm]

Actually, one thing that could happen: A rewrite pattern of inside_op could modify the outside_op (or erase it, etc.). inside_op could just be the anchor point of the pattern. This could happen even if the region is isolated from above. Nothing prevents the pattern from retrieving the parent of inside_op and doing something with it. E.g., a swapping pattern that moves an op out of a region could do that (even if the region is isolated from above). When the region is not isolated from above that's much easier though.

We could guard against this by listening to IR changes. If an op outside of the region gets modifed/erased/etc. (as indicated by finalizeRootUpdate/notifyOperationRemoved) we could assert. That would also safeguard against illegal IR changes even when the region is isolated from above.

[springerm]

How are "use lists" of Values modeled? What's happening when a Value (e.g., OpResult of outside_op) looses a user? Does this trigger an update of some internal data structure of outside_op? If so, are these data structures thread-safe?

[mehdi_amini]

In D141366#4054440, @springerm wrote:

Actually, one thing that could happen: A rewrite pattern of inside_op could modify the outside_op (or erase it, etc.). inside_op could just be the anchor point of the pattern.

Yes this is the kind of thing I had in mind: the kind of patterns that would rewrite a subgraph as a whole.

This could happen even if the region is isolated from above. Nothing prevents the pattern from retrieving the parent of inside_op and doing something with it. E.g., a swapping pattern that moves an op out of a region could do that (even if the region is isolated from above). When the region is not isolated from above that's much easier though.

Yes the difference is whether to be intentional about it or not: patterns walking use-def chains is standard, walking parent and going through an isolated region is already dangerous (we can do it, but such pattern will put restrictions on the kind of pass that can run it).

We could guard against this by listening to IR changes. If an op outside of the region gets modifed/erased/etc. (as indicated by finalizeRootUpdate/notifyOperationRemoved) we could assert. That would also safeguard against illegal IR changes even when the region is isolated from above.

Defending against anything unexpected is welcome, we could and should do better. In this case I'm not sure it is desirable to have this situation as it seems too much a footgun and also the only way for a pattern to be written safely would require to thread through the "scopes" to the pattern itself where everything would have to be checked.

[mehdi_amini]

In D141366#4054442, @springerm wrote:

How are "use lists" of Values modeled? What's happening when a Value (e.g., OpResult of outside_op) looses a user? Does this trigger an update of some internal data structure of outside_op? If so, are these data structures thread-safe?

See the schema at the bottom of https://mlir.llvm.org/docs/Tutorials/UnderstandingTheIRStructure/

This is not thread-safe: the whole threading model is organized around isolated regions in MLIR (hence why you can't run a pass on a non-isolated operation for example).

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
 
   static constexpr int64_t kNoLimit = -1;
 };
 
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
-/// Rewrite the regions of the specified operation, which must be isolated from
-/// above, by repeatedly applying the highest benefit patterns in a greedy
-/// work-list driven manner.
+/// Rewrite the regions of the specified operation by repeatedly applying the
+/// highest benefit patterns in a greedy work-list driven manner. Only ops that
+/// are inside one of given regions are added to the work-list.
 ///
 /// 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.
 ///
 /// Return success if the iterative process converged and no more patterns can
 /// be matched in the result operation regions.
 ///
 /// 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.
 ///
 /// You may configure several aspects of this with GreedyRewriteConfig.
 LogicalResult applyPatternsAndFoldGreedily(
     MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
     GreedyRewriteConfig config = GreedyRewriteConfig());
 
-/// Rewrite the given regions, which must be isolated from above.
+/// Rewrite the given regions.
 inline LogicalResult applyPatternsAndFoldGreedily(
     Operation *op, const FrozenRewritePatternSet &patterns,
     GreedyRewriteConfig config = GreedyRewriteConfig()) {
   return applyPatternsAndFoldGreedily(op->getRegions(), patterns, config);
 }
 
 /// Applies the specified patterns on `op` alone while also trying to fold it,
 /// by selecting the highest benefits patterns in a greedy manner. Returns

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
   });
   return failure();
 }
 
-/// Rewrite the regions of the specified operation, which must be isolated from
-/// above, by repeatedly applying the highest benefit patterns in a greedy
-/// work-list driven manner. Return success if no more patterns can be matched
-/// in the result operation regions. Note: This does not apply patterns to the
-/// top-level operation itself.
+/// Rewrite the regions of the specified operation by repeatedly applying the
+/// highest benefit patterns in a greedy work-list driven manner. Return success
+/// if no more patterns can be matched in the result operation regions.
+/// Note: This does not apply patterns to the top-level operation itself.
 ///
 LogicalResult
 mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
                                    const FrozenRewritePatternSet &patterns,
                                    GreedyRewriteConfig config) {
   if (regions.empty())
     return success();
 
-  // 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'.
-  auto regionIsIsolated = [](Region &region) {
-    return region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>();
-  };
-  (void)regionIsIsolated;
-  assert(llvm::all_of(regions, regionIsIsolated) &&
-         "patterns can only be applied to operations IsolatedFromAbove");
-
   // Start the pattern driver.
   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config);
   bool converged = driver.simplify(regions);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite did not converge after scanning "
                  << config.maxIterations << " times\n";
   });
   return success(converged);