[mlir] GreedyPatternRewriteDriver: Support strict mode to GreedyPatternRewriteDriver

[springerm]

Summary

strictMode is moved to GreedyRewriteConfig to simplify the API and state of rewriter classes. The region-based GreedyPatternRewriteDriver now also supports strict mode.

MultiOpPatternRewriteDriver becomes simpler: fewer method must be overridden.

Depends On: D142614

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
   /// `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.
   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
+  ///   enqueued. All other ops are excluded.
+  /// * GreedyRewriteStrictness::ExistingOps: Only pre-existing ops (that were
+  ///   were on the worklist at the very beginning) enqueued. All other ops are
+  ///   excluded.
+  GreedyRewriteStrictness strictMode = GreedyRewriteStrictness::AnyOp;
 };
 
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
 /// Rewrite ops in the given region, which must be isolated from above, by
 /// repeatedly applying the highest benefit patterns in a greedy work-list
inline LogicalResult applyPatternsAndFoldGreedily(
   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
-/// of `strictMode`).
-///
-/// * GreedyRewriteStrictness::AnyOp: No ops are excluded.
-/// * GreedyRewriteStrictness::ExistingAndNewOps: Only pre-existing and newly
-///   created ops are simplified. All other ops are excluded.
-/// * GreedyRewriteStrictness::ExistingOps: Only pre-existing ops are
-///   simplified. All other ops are excluded.
+/// 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.
 ///
 /// 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.
 ///
 /// 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.
 LogicalResult
 applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                        const FrozenRewritePatternSet &patterns,
-                       GreedyRewriteStrictness strictMode,
                        GreedyRewriteConfig config = GreedyRewriteConfig(),
                        bool *changed = nullptr, bool *allErased = nullptr);
 
-/// Applies the specified patterns on `op` alone while also trying to fold it,
-/// by selecting the highest benefits patterns in a greedy manner. Returns
-/// success if no more patterns can be matched. `erased` is set to true if `op`
-/// was folded away or erased as a result of becoming dead.
-///
-/// Returns success if the iterative process converged and no more patterns can
-/// be matched.
+/// Applies the specified patterns on `op` while also trying to fold it.
+/// This function is a shortcut for the ArrayRef<Operation *> overload and
+/// behaves the same way.
 inline LogicalResult
 applyOpPatternsAndFold(Operation *op, const FrozenRewritePatternSet &patterns,
                        GreedyRewriteConfig config = GreedyRewriteConfig(),
                        bool *erased = nullptr) {
-  return applyOpPatternsAndFold(ArrayRef(op), patterns,
-                                GreedyRewriteStrictness::ExistingOps, config,
+  return applyOpPatternsAndFold(ArrayRef(op), patterns, config,
                                 /*changed=*/nullptr, erased);
 }
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
   RewritePatternSet patterns(getContext());
   // Canonicalization patterns are needed so that affine.apply ops are composed
   // with the remaining affine.min/max ops.
   AffineMaxOp::getCanonicalizationPatterns(patterns, getContext());
   AffineMinOp::getCanonicalizationPatterns(patterns, getContext());
   patterns.insert<SimplifyAffineMinMaxOp<AffineMinOp>,
                   SimplifyAffineMinMaxOp<AffineMaxOp>>(getContext(), cstr);
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  GreedyRewriteConfig config;
+  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
   // Apply the simplification pattern to a fixpoint.
-  if (failed(
-          applyOpPatternsAndFold(targets, frozenPatterns,
-                                 GreedyRewriteStrictness::ExistingAndNewOps))) {
+  if (failed(applyOpPatternsAndFold(targets, frozenPatterns, config))) {
     auto diag = emitDefiniteFailure()
                 << "affine.min/max simplification did not converge";
     return diag;
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 void SimplifyBoundedAffineOpsOp::getEffects(

mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp

void AffineDataCopyGeneration::runOnOperation() {
 
   // Promoting single iteration loops could lead to simplification of
   // contained load's/store's, and the latter could anyway also be
   // canonicalized.
   RewritePatternSet patterns(&getContext());
   AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
   AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-  (void)applyOpPatternsAndFold(copyOps, frozenPatterns,
-                               GreedyRewriteStrictness::ExistingAndNewOps);
+  GreedyRewriteConfig config;
+  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
+  (void)applyOpPatternsAndFold(copyOps, frozenPatterns, config);
 }

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

for (auto attr : op->getAttrs()) {
         simplifyAndUpdateAttribute(op, attr.getName(), mapAttr);
       else if (auto setAttr = attr.getValue().dyn_cast<IntegerSetAttr>())
         simplifyAndUpdateAttribute(op, attr.getName(), setAttr);
     }
 
     if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
       opsToSimplify.push_back(op);
   });
-  (void)applyOpPatternsAndFold(opsToSimplify, frozenPatterns,
-                               GreedyRewriteStrictness::ExistingAndNewOps);
+  GreedyRewriteConfig config;
+  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
+  (void)applyOpPatternsAndFold(opsToSimplify, frozenPatterns, config);
 }

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

if (!opGroupQueue.empty()) {
                                   opGroupQueue, /*offset=*/0, forOp, b);
         lbShift = d * step;
       }
 
       if (res) {
         // Simplify/canonicalize the affine.for.
         RewritePatternSet patterns(res.getContext());
         AffineForOp::getCanonicalizationPatterns(patterns, res.getContext());
+        GreedyRewriteConfig config;
+        config.strictMode = GreedyRewriteStrictness::ExistingOps;
         bool erased;
-        (void)applyOpPatternsAndFold(res, std::move(patterns),
-                                     GreedyRewriteConfig(), &erased);
+        (void)applyOpPatternsAndFold(res, std::move(patterns), config, &erased);
         if (!erased && !prologue)
           prologue = res;
         if (!erased)
           epilogue = res;
       }
     } else {
       // Start of first interval.
       lbShift = d * step;

mlir/lib/Dialect/Affine/Utils/Utils.cpp

LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
 
   // Apply canonicalization patterns and folding - this is necessary for the
   // hoisting check to be correct (operands should be composed), and to be more
   // effective (no unused operands). Since the pattern rewriter's folding is
   // entangled with application of patterns, we may fold/end up erasing the op,
   // in which case we return with `folded` being set.
   RewritePatternSet patterns(ifOp.getContext());
   AffineIfOp::getCanonicalizationPatterns(patterns, ifOp.getContext());
-  bool erased;
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-  (void)applyOpPatternsAndFold(ifOp, frozenPatterns, GreedyRewriteConfig(),
-                               &erased);
+  GreedyRewriteConfig config;
+  config.strictMode = GreedyRewriteStrictness::ExistingOps;
+  bool erased;
+  (void)applyOpPatternsAndFold(ifOp, frozenPatterns, config, &erased);
   if (erased) {
     if (folded)
       *folded = true;
     return failure();
   }
   if (folded)
     *folded = false;
 

mlir/lib/Reducer/ReductionTreePass.cpp

for (const auto &op : enumerate(region.getOps())) {
     else
       opsInRange.push_back(&op.value());
   }
 
   // `applyOpPatternsAndFold` may erase the ops so we can't do the pattern
   // matching in above iteration. Besides, erase op not-in-range may end up in
   // invalid module, so `applyOpPatternsAndFold` should come before that
   // transform.
-  for (Operation *op : opsInRange)
+  for (Operation *op : opsInRange) {
     // `applyOpPatternsAndFold` returns whether the op is convered. Omit it
     // because we don't have expectation this reduction will be success or not.
-    (void)applyOpPatternsAndFold(op, patterns);
+    GreedyRewriteConfig config;
+    config.strictMode = GreedyRewriteStrictness::ExistingOps;
+    (void)applyOpPatternsAndFold(op, patterns, config);
+  }
 
   if (eraseOpNotInRange)
     for (Operation *op : opsNotInRange) {
       op->dropAllUses();
       op->erase();
     }
 }
 

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
   void removeFromWorklist(Operation *op);
 
   /// Notifies the driver that the specified operation may have been modified
   /// in-place.
   void finalizeRootUpdate(Operation *op) override;
 
 protected:
   /// Add the given operation to the worklist.
-  virtual void addSingleOpToWorklist(Operation *op);
+  void addSingleOpToWorklist(Operation *op);
 
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
   void notifyOperationInserted(Operation *op) override;
 
   // 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
protected:
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
 protected:
   /// Configuration information for how to simplify.
   const GreedyRewriteConfig config;
 
+  /// The list of ops we are restricting our rewrites to. These include the
+  /// supplied set of ops as well as new ops created while rewriting those ops
+  /// depending on `strictMode`. This set is not maintained when
+  /// `config.strictMode` is GreedyRewriteStrictness::AnyOp.
+  llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
+
 private:
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // namespace
 
auto insertKnownConstant = [&](Operation *op) {
     // accidentally reversing the constant order during processing.
     Attribute constValue;
     if (matchPattern(op, m_Constant(&constValue)))
       if (!folder.insertKnownConstant(op, constValue))
         return true;
     return false;
   };
 
+  // Populate strict mode ops.
+  if (config.strictMode != GreedyRewriteStrictness::AnyOp) {
+    strictModeFilteredOps.clear();
+    region.walk([&](Operation *op) { strictModeFilteredOps.insert(op); });
+  }
+
   bool changed = false;
   int64_t iteration = 0;
   do {
     // Check if the iteration limit was reached.
     if (iteration++ >= config.maxIterations &&
         config.maxIterations != GreedyRewriteConfig::kNoLimit)
       break;
 
while (Region *region = op->getParentRegion()) {
     op = region->getParentOp();
     if (!op)
       break;
     ancestors.push_back(op);
   }
 }
 
 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;
 
     worklistMap[op] = worklist.size();
     worklist.push_back(op);
   }
+}
 
 Operation *GreedyPatternRewriteDriver::popFromWorklist() {
   auto *op = worklist.back();
   worklist.pop_back();
 
   // This operation is no longer in the worklist, keep worklistMap up to date.
   if (op)
     worklistMap.erase(op);
void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
   }
 }
 
 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+  if (config.strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
+    strictModeFilteredOps.insert(op);
   addToWorklist(op);
 }
 
 void GreedyPatternRewriteDriver::finalizeRootUpdate(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
                        << ")\n";
   });
logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
 
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
+
+  if (config.strictMode != GreedyRewriteStrictness::AnyOp)
+    strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op,
                                                     ValueRange replacement) {
   LLVM_DEBUG({
     logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
                        << ")\n";
   });

 /// 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,
-      GreedyRewriteStrictness strictMode, const GreedyRewriteConfig &config,
+      const GreedyRewriteConfig &config,
       llvm::SmallDenseSet<Operation *, 4> *survivingOps = nullptr)
       : GreedyPatternRewriteDriver(ctx, patterns, config),
-        strictMode(strictMode), survivingOps(survivingOps) {}
+        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`.
   ///
   /// Note that ops in `ops` could be erased as a result of folding, becoming
   /// dead, or via pattern rewrites. The return value indicates convergence.
   ///
   /// All erased ops are stored in `erased`.
   LogicalResult simplifyLocally(ArrayRef<Operation *> op,
                                 bool *changed = nullptr) &&;
 
-protected:
-  void addSingleOpToWorklist(Operation *op) override {
-    if (strictMode == GreedyRewriteStrictness::AnyOp ||
-        strictModeFilteredOps.contains(op))
-      GreedyPatternRewriteDriver::addSingleOpToWorklist(op);
-  }
-
 private:
-  void notifyOperationInserted(Operation *op) override {
-    if (strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
-      strictModeFilteredOps.insert(op);
-    GreedyPatternRewriteDriver::notifyOperationInserted(op);
-  }
-
   void notifyOperationRemoved(Operation *op) override {
     GreedyPatternRewriteDriver::notifyOperationRemoved(op);
     if (survivingOps)
       survivingOps->erase(op);
-    if (strictMode != GreedyRewriteStrictness::AnyOp)
-      strictModeFilteredOps.erase(op);
   }
 
-  /// `strictMode` control which ops are added to the worklist during
-  /// simplification.
-  const GreedyRewriteStrictness strictMode = GreedyRewriteStrictness::AnyOp;
-
-  /// The list of ops we are restricting our rewrites to. These include the
-  /// supplied set of ops as well as new ops created while rewriting those ops
-  /// depending on `strictMode`. This set is not maintained when `strictMode`
-  /// is GreedyRewriteStrictness::AnyOp.
-  llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
-
   /// An optional set of ops that survived the rewrite. This set is populated
   /// at the beginning of `simplifyLocally` with the inititally provided list
   /// of ops.
   llvm::SmallDenseSet<Operation *, 4> *const survivingOps = nullptr;
 };
 
 } // namespace
 
 LogicalResult
 MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
                                              bool *changed) && {
   if (survivingOps) {
     survivingOps->clear();
     survivingOps->insert(ops.begin(), ops.end());
   }
 
-  if (strictMode != GreedyRewriteStrictness::AnyOp) {
+  if (config.strictMode != GreedyRewriteStrictness::AnyOp) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
   if (changed)
     *changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
     addSingleOpToWorklist(op);
 
   // These are scratch vectors used in the folding loop below.
   SmallVector<Value, 8> originalOperands, resultValues;
   int64_t numRewrites = 0;
   while (!worklist.empty() &&
          (numRewrites < config.maxNumRewrites ||
           config.maxNumRewrites == GreedyRewriteConfig::kNoLimit)) {
     Operation *op = popFromWorklist();
 
     // Nulls get added to the worklist when operations are removed, ignore
     // them.
     if (op == nullptr)
       continue;
 
-    assert((strictMode == GreedyRewriteStrictness::AnyOp ||
+    assert((config.strictMode == GreedyRewriteStrictness::AnyOp ||
             strictModeFilteredOps.contains(op)) &&
            "unexpected op was inserted under strict mode");
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       notifyOperationRemoved(op);
       op->erase();
       if (changed)
if (remainingOps.none())
       break;
   } while ((region = region->getParentRegion()));
   assert(region && "could not find common parent region");
   return region;
 }
 
 LogicalResult mlir::applyOpPatternsAndFold(
     ArrayRef<Operation *> ops, const FrozenRewritePatternSet &patterns,
-    GreedyRewriteStrictness strictMode, GreedyRewriteConfig config,
-    bool *changed, bool *allErased) {
+    GreedyRewriteConfig config, bool *changed, bool *allErased) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
     if (allErased)
       *allErased = true;
     return success();
   }
 
#ifndef NDEBUG
     });
     assert(allOpsInScope && "ops must be within the specified scope");
 #endif // NDEBUG
   }
 
   // Start the pattern driver.
   llvm::SmallDenseSet<Operation *, 4> surviving;
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
-                                     strictMode, config,
-                                     allErased ? &surviving : nullptr);
+                                     config, 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 "
                  << config.maxNumRewrites << " rewrites";
   });
   return converged;
 }

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

for (Operation *op : copyOps) {
     patterns.clear();
     if (isa<AffineLoadOp>(op)) {
       AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
     } else {
       assert(isa<AffineStoreOp>(op) && "expected affine store op");
       AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
     }
   }
-  (void)applyOpPatternsAndFold(copyOps, std::move(patterns),
-                               GreedyRewriteStrictness::ExistingAndNewOps);
+  GreedyRewriteConfig config;
+  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
+  (void)applyOpPatternsAndFold(copyOps, std::move(patterns), config);
 }
 
 namespace mlir {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>();
 }
 } // namespace mlir

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

void runOnOperation() override {
     getOperation()->walk([&](Operation *op) {
       StringRef opName = op->getName().getStringRef();
       if (opName == "test.insert_same_op" ||
           opName == "test.replace_with_new_op" || opName == "test.erase_op") {
         ops.push_back(op);
       }
     });
 
-    GreedyRewriteStrictness mode;
+    GreedyRewriteConfig config;
     if (strictMode == "AnyOp") {
-      mode = GreedyRewriteStrictness::AnyOp;
+      config.strictMode = GreedyRewriteStrictness::AnyOp;
     } else if (strictMode == "ExistingAndNewOps") {
-      mode = GreedyRewriteStrictness::ExistingAndNewOps;
+      config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
     } else if (strictMode == "ExistingOps") {
-      mode = GreedyRewriteStrictness::ExistingOps;
+      config.strictMode = GreedyRewriteStrictness::ExistingOps;
     } else {
       llvm_unreachable("invalid strictness option");
     }
 
     // Check if these transformations introduce visiting of operations that
     // are not in the `ops` set (The new created ops are valid). An invalid
     // operation will trigger the assertion while processing.
     bool changed = false;
     bool allErased = false;
-    (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns), mode,
-                                 GreedyRewriteConfig(), &changed, &allErased);
+    (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns), config,
+                                 &changed, &allErased);
     Builder b(ctx);
     getOperation()->setAttr("pattern_driver_changed", b.getBoolAttr(changed));
     getOperation()->setAttr("pattern_driver_all_erased",
                             b.getBoolAttr(allErased));
   }
 
   Option<std::string> strictMode{
       *this, "strictness",