[mlir] GreedyPatternRewriteDriver: Ignore scope when rewriting top-level ops

[springerm]

Summary

Top-level ModuleOps cannot be transformed with the GreedyPatternRewriteDriver since D141945 because they do not have an enclosing region that could be used as a scope. Make the scope optional inside GreedyPatternRewriteDriver, so that top-level ops can be processed when they are on the initial list of ops.

Note: This does not allow users to bypass the scoping mechanism by setting config.scope = nullptr.

[mikeurbach]

Thanks! I think this makes sense to me.

[rkayaith]

I'm curious about the use case here, wouldn't replacing the top-level op cause a crash? The pass manager and tools hold a pointer to the top level op, so it doesn't seem safe to apply rewrites to it.

[mehdi_amini]

LGTM, should the description has Fixes #60462 as well?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

276

Technically you're doing more pointer chasing here aren't you? The op->getParentOp() will do some redundant traversal from op->getParentRegion().
This may very much not matter...

[mehdi_amini]

In D143151#4100391, @rkayaith wrote:

I'm curious about the use case here, wouldn't replacing the top-level op cause a crash? The pass manager and tools hold a pointer to the top level op, so it doesn't seem safe to apply rewrites to it.

Absolutely, but you can rewrite in place or rewrite downwards. For example a pattern may match a my_hw_specific_module op and normalize attribute, and lookup the entry point and the global variables and do something about it, without replacing the top-level my_hw_specific_module op.

[mehdi_amini]

Also, this is in line with a module pass: independently of the pattern-rewriting you can modify the ModuleOp in a Module pass but you can't replace it!

[rkayaith]

In D143151#4100560, @mehdi_amini wrote:

In D143151#4100391, @rkayaith wrote:

I'm curious about the use case here, wouldn't replacing the top-level op cause a crash? The pass manager and tools hold a pointer to the top level op, so it doesn't seem safe to apply rewrites to it.

Absolutely, but you can rewrite in place or rewrite downwards. For example a pattern may match a my_hw_specific_module op and normalize attribute, and lookup the entry point and the global variables and do something about it, without replacing the top-level my_hw_specific_module op.

In D143151#4100561, @mehdi_amini wrote:

Also, this is in line with a module pass: independently of the pattern-rewriting you can modify the ModuleOp in a Module pass but you can't replace it!

I see, I ran into an issue because dialect conversion includes the parent op in what can be matched: https://github.com/llvm/llvm-project/issues/60491
But I see now that this change only affects applyOpPatternsAndFold, and not applyPatternsAndFoldGreedily. Since applyOpPatternsAndFold is a more targeted API I guess it would be less likely to run into a similar issue.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
 
   /// This specifies the maximum number of rewrites within an iteration. Use
   /// `kNoLimit` to disable this limit.
   int64_t maxNumRewrites = kNoLimit;
 
   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 is used as a scope.
+  /// specified, the closest enclosing region around the initial list of ops
+  /// is used as a scope.
   Region *scope = nullptr;
 
   /// Strict mode can restrict the ops that are added to the worklist during
   /// the rewrite.
   ///
   /// * GreedyRewriteStrictness::AnyOp: No ops are excluded.
   /// * GreedyRewriteStrictness::ExistingAndNewOps: Only pre-existing ops (that
   ///   were on the worklist at the very beginning) and newly created ops are

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

#endif
   PatternApplicator matcher;
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), folder(ctx), config(config), matcher(patterns) {
-  assert(config.scope && "scope is not specified");
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 }
 
 bool GreedyPatternRewriteDriver::processWorklist() {
 #ifndef NDEBUG
#endif
   }
 
   return changed;
 }
 
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
   // Gather potential ancestors while looking for a "scope" parent region.
   SmallVector<Operation *, 8> ancestors;
+  do {
     ancestors.push_back(op);
-  while (Region *region = op->getParentRegion()) {
-      if (config.scope == region) {
-        // All gathered ops are in fact ancestors.
-        for (Operation *op : ancestors)
-          addSingleOpToWorklist(op);
-        break;
-      }
-    op = region->getParentOp();
-    if (!op)
-      break;
-    ancestors.push_back(op);
+    if (config.scope == op->getParentRegion()) {
+      // Scope (can be `nullptr`) was reached. Stop traveral and enqueue ops.
+      return;
     }
+  } while ((op = op->getParentOp()));
 }
 
 void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
   if (config.strictMode == GreedyRewriteStrictness::AnyOp ||
       strictModeFilteredOps.contains(op)) {
     // Check to see if the worklist already contains this op.
     if (worklistMap.count(op))
       return;
LogicalResult MultiOpPatternRewriteDriver::simplify(ArrayRef<Operation *> ops,
   bool result = processWorklist();
   if (changed)
     *changed = result;
 
   return success(worklist.empty());
 }
 
 /// Find the 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.
 static Region *findCommonAncestor(ArrayRef<Operation *> ops) {
   assert(!ops.empty() && "expected at least one op");
   // Fast path in case there is only one op.
   if (ops.size() == 1)
     return ops.front()->getParentRegion();
 
   Region *region = ops.front()->getParentRegion();
   ops = ops.drop_front();
   int sz = ops.size();
   llvm::BitVector remainingOps(sz, true);
-  do {
+  while (region) {
     int pos = -1;
     // Iterate over all remaining ops.
     while ((pos = remainingOps.find_first_in(pos + 1, sz)) != -1) {
       // Is this op contained in `region`?
       if (region->findAncestorOpInRegion(*ops[pos]))
         remainingOps.reset(pos);
     }
     if (remainingOps.none())
       break;
-  } while ((region = region->getParentRegion()));
-  assert(region && "could not find common parent region");
+    region = region->getParentRegion();
+  }
   return region;
 }
 
 LogicalResult mlir::applyOpPatternsAndFold(
     ArrayRef<Operation *> ops, const FrozenRewritePatternSet &patterns,
     GreedyRewriteConfig config, bool *changed, bool *allErased) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
     if (allErased)
       *allErased = true;
     return success();
   }
 
   // Determine scope of rewrite.
   if (!config.scope) {
-    // Compute scope if none was provided.
+    // Compute scope if none was provided. The scope will remain `nullptr` if
+    // 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.
 #ifndef NDEBUG
     bool allOpsInScope = llvm::all_of(ops, [&](Operation *op) {
       return static_cast<bool>(config.scope->findAncestorOpInRegion(*op));
     });
     assert(allOpsInScope && "ops must be within the specified scope");