[mlir] GreedyPatternRewriteDriver: Entry point takes single region

[springerm]

Summary

The rewrite driver is typically applied to a single region or all regions of the same op. There is no longer an overload to apply the rewrite driver to a list of regions.

This simplifies the rewrite driver implementation because the scope is now a single region as opposed to a list of regions.

Note: This change is not NFC because config.maxIterations and config.maxNumRewrites is now counted for each region separately.

Depends On: D141949

[springerm]

update

[mehdi_amini]

Note: This change is not NFC because config.maxIterations and config.maxNumRewrites is now counted for each region separately.

I think it isn't NFC also because regions are processed separately rather than together right?
Now we'll be reaching fix-point on regions in isolation and the worklist filtering won't add operations outside of the current region.

[springerm]

In D142611#4084357, @mehdi_amini wrote:

Note: This change is not NFC because config.maxIterations and config.maxNumRewrites is now counted for each region separately.

I think it isn't NFC also because regions are processed separately rather than together right?
Now we'll be reaching fix-point on regions in isolation and the worklist filtering won't add operations outside of the current region.

Expanded commit message.

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 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.
 ///
 /// 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,
+    Region &region, const FrozenRewritePatternSet &patterns,
     GreedyRewriteConfig config = GreedyRewriteConfig());
 
-/// Rewrite the given regions, which must be isolated from above.
+/// Rewrite ops in all regions of the given op, which must be isolated from
+/// above.
 inline LogicalResult applyPatternsAndFoldGreedily(
     Operation *op, const FrozenRewritePatternSet &patterns,
     GreedyRewriteConfig config = GreedyRewriteConfig()) {
-  return applyPatternsAndFoldGreedily(op->getRegions(), patterns, config);
+  bool failed = false;
+  for (Region &region : op->getRegions())
+    failed |= applyPatternsAndFoldGreedily(region, patterns, config).failed();
+  return failure(failed);
 }
 
 /// Applies the specified rewrite patterns on `ops` while also trying to fold
 /// these ops.
 ///
 /// 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
 /// excluded via `strictMode`. Any other ops remain unmodified (i.e., regardless

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
         patterns);
 
     // Add constant folding patterns.
     populateConstantFoldLinalgOperations(patterns, defaultControlFn);
 
     // Use TopDownTraversal for compile time reasons
     GreedyRewriteConfig grc;
     grc.useTopDownTraversal = true;
-    (void)applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns),
-                                       grc);
+    (void)applyPatternsAndFoldGreedily(op, std::move(patterns), grc);
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgElementwiseOpFusionPass() {
   return std::make_unique<LinalgElementwiseOpFusionPass>();
 }

mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp

struct ExpandStridedMetadataPass final
   void runOnOperation() override;
 };
 
 } // namespace
 
 void ExpandStridedMetadataPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   memref::populateExpandStridedMetadataPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(),
-                                     std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 std::unique_ptr<Pass> memref::createExpandStridedMetadataPass() {
   return std::make_unique<ExpandStridedMetadataPass>();
 }

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

struct FoldMemRefAliasOpsPass final
   void runOnOperation() override;
 };
 
 } // namespace
 
 void FoldMemRefAliasOpsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   memref::populateFoldMemRefAliasOpPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(),
-                                     std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 std::unique_ptr<Pass> memref::createFoldMemRefAliasOpsPass() {
   return std::make_unique<FoldMemRefAliasOpsPass>();
 }

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

void memref::populateResolveShapedTypeResultDimsPatterns(
   patterns.add<DimOfShapedTypeOpInterface<memref::DimOp>,
                DimOfShapedTypeOpInterface<tensor::DimOp>>(
       patterns.getContext());
 }
 
 void ResolveRankedShapeTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
-  if (failed(applyPatternsAndFoldGreedily(getOperation()->getRegions(),
-                                          std::move(patterns))))
+  if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
 }
 
 void ResolveShapedTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
-  if (failed(applyPatternsAndFoldGreedily(getOperation()->getRegions(),
-                                          std::move(patterns))))
+  if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<Pass> memref::createResolveShapedTypeResultDimsPass() {
   return std::make_unique<ResolveShapedTypeResultDimsPass>();
 }
 
 std::unique_ptr<Pass> memref::createResolveRankedShapeTypeResultDimsPass() {
   return std::make_unique<ResolveRankedShapeTypeResultDimsPass>();
 }

mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp

struct LowerVectorMaskPass
 
   void runOnOperation() override {
     Operation *op = getOperation();
     MLIRContext *context = op->getContext();
 
     RewritePatternSet loweringPatterns(context);
     populateVectorMaskLoweringPatternsForSideEffectingOps(loweringPatterns);
 
-    if (failed(applyPatternsAndFoldGreedily(op->getRegions(),
-                                            std::move(loweringPatterns))))
+    if (failed(applyPatternsAndFoldGreedily(op, std::move(loweringPatterns))))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
   }
 };
 

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 namespace {
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns in a roughly "bottom up" way.
 class GreedyPatternRewriteDriver : public PatternRewriter {
 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
                                       const GreedyRewriteConfig &config,
-                                      const DenseSet<Region *> &scope);
+                                      const Region &scope);
 
-  /// Simplify the operations within the given regions.
-  bool simplify(MutableArrayRef<Region> regions) &&;
+  /// Simplify the ops within the given region.
+  bool simplify(Region &region) &&;
 
   /// Add the given operation and its ancestors to the worklist.
   void addToWorklist(Operation *op);
 
   /// Pop the next operation from the worklist.
   Operation *popFromWorklist();
 
   /// If the specified operation is in the worklist, remove it.
protected:
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
 protected:
   /// Configuration information for how to simplify.
   const GreedyRewriteConfig config;
 
   /// Only ops within this scope are simplified.
-  const DenseSet<Region *> scope;
+  const Region &scope;
 
 private:
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
-    const GreedyRewriteConfig &config, const DenseSet<Region *> &scope)
+    const GreedyRewriteConfig &config, const Region &scope)
     : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config),
       scope(scope) {
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 }
 
-bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) && {
+bool GreedyPatternRewriteDriver::simplify(Region &region) && {
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
 
   /// A utility function to log a process result for the given reason.
   auto logResult = [&](StringRef result, const llvm::Twine &msg = {}) {
     logger.unindent();
     logger.startLine() << "} -> " << result;
if (iteration++ >= config.maxIterations &&
         config.maxIterations != GreedyRewriteConfig::kNoLimit)
       break;
 
     worklist.clear();
     worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
-      for (auto &region : regions) {
         region.walk([&](Operation *op) {
           if (!insertKnownConstant(op))
             addToWorklist(op);
         });
-      }
     } else {
       // Add all nested operations to the worklist in preorder.
-      for (auto &region : regions) {
         region.walk<WalkOrder::PreOrder>([&](Operation *op) {
           if (!insertKnownConstant(op)) {
             worklist.push_back(op);
             return WalkResult::advance();
           }
           return WalkResult::skip();
         });
-      }
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());
       // Remember the reverse index.
       for (size_t i = 0, e = worklist.size(); i != e; ++i)
         worklistMap[worklist[i]] = i;
     }
 
#endif
         changed = true;
         ++numRewrites;
       }
     }
 
     // After applying patterns, make sure that the CFG of each of the regions
     // is kept up to date.
     if (config.enableRegionSimplification)
-      changed |= succeeded(simplifyRegions(*this, regions));
+      changed |= succeeded(simplifyRegions(*this, region));
   } while (changed);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
   // Gather potential ancestors while looking for a "scope" parent region.
   SmallVector<Operation *, 8> ancestors;
   ancestors.push_back(op);
   while (Region *region = op->getParentRegion()) {
-    if (scope.contains(region)) {
+    if (&scope == region) {
       // All gathered ops are in fact ancestors.
       for (Operation *op : ancestors)
         addSingleOpToWorklist(op);
       break;
     }
     op = region->getParentOp();
     if (!op)
       break;

 
 /// 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.
 ///
 LogicalResult
-mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
+mlir::applyPatternsAndFoldGreedily(Region &region,
                                    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) &&
+  assert(region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
          "patterns can only be applied to operations IsolatedFromAbove");
 
-  // Limit ops on the worklist to this scope.
-  DenseSet<Region *> scope;
-  for (Region &r : regions)
-    scope.insert(&r);
-
   // Start the pattern driver.
-  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config,
-                                    scope);
-  bool converged = std::move(driver).simplify(regions);
+  GreedyPatternRewriteDriver driver(region.getContext(), patterns, config,
+                                    region);
+  bool converged = std::move(driver).simplify(region);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite did not converge after scanning "
                  << config.maxIterations << " times\n";
   });
   return success(converged);
 }
 
 //===----------------------------------------------------------------------===//
 // MultiOpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// This is a specialized GreedyPatternRewriteDriver to apply patterns and
 /// perform folding for a supplied set of ops. It repeatedly simplifies while
 /// restricting the rewrites to only the provided set of ops or optionally
 /// to those directly affected by it (result users or operand providers). Parent
 /// ops are not considered.
 class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
 public:
   explicit MultiOpPatternRewriteDriver(
       MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
-      const DenseSet<Region *> &scope, GreedyRewriteStrictness strictMode,
+      const Region &scope, GreedyRewriteStrictness strictMode,
       llvm::SmallDenseSet<Operation *, 4> *survivingOps = nullptr)
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig(), scope),
         strictMode(strictMode), survivingOps(survivingOps) {}
 
   /// Performs the specified rewrites on `ops` while also trying to fold these
   /// ops. `strictMode` controls which other ops are simplified. Only ops
   /// within the given scope region are added to the worklist. If no scope is
   /// specified, it assumed to be closest common region of all `ops`.
bool allOpsInScope = llvm::all_of(ops, [&](Operation *op) {
       return static_cast<bool>(scope->findAncestorOpInRegion(*op));
     });
     assert(allOpsInScope && "ops must be within the specified scope");
 #endif // NDEBUG
   }
 
   // Start the pattern driver.
   llvm::SmallDenseSet<Operation *, 4> surviving;
-  DenseSet<Region *> scopeSet;
-  scopeSet.insert(scope);
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
-                                     scopeSet, strictMode,
+                                     *scope, strictMode,
                                      allErased ? &surviving : nullptr);
   LogicalResult converged = std::move(driver).simplifyLocally(ops, changed);
   if (allErased)
     *allErased = surviving.empty();
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite did not converge after "
                  << GreedyRewriteConfig().maxNumRewrites << " rewrites";
   });
   return converged;
 }

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestSelectiveReplacementPatternDriver
   }
   StringRef getDescription() const final {
     return "Test selective replacement in the PatternRewriter";
   }
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     mlir::RewritePatternSet patterns(context);
     patterns.add<TestSelectiveOpReplacementPattern>(context);
-    (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(),
-                                       std::move(patterns));
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//