[mlir] GreedyPatternRewriteDriver: Add new strict mode option

[springerm]

Summary

There are now three options:

• AnyOp (previously false)
• ExistingAndNewOps (previously true)
• ExistingOps: this one is new.

The last option corresponds to what the applyOpPatternsAndFold(Operation*, ...) overload is doing. It is now also supported on the applyOpPatternsAndFold(ArrayRef<Operation *>, ...) overload.

[springerm]

rebase

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

605

mlir/test/Transforms/test-strict-pattern-driver.mlir

1

Since we're here, can you remove the allow-unregistered-dialect?

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 
 #ifndef MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_
 #define MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_
 
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 
 namespace mlir {
 
+/// This enum controls which ops are put on the worklist during a greedy
+/// pattern rewrite.
+enum class GreedyRewriteStrictness {
+  /// No restrictions wrt. which ops are processed.
+  AnyOp,
+  /// Only pre-existing and newly created ops are processed.
+  ExistingAndNewOps,
+  /// Only pre-existing ops are processed.
+  ExistingOps
+};
+
 /// This class allows control over how the GreedyPatternRewriteDriver works.
 class GreedyRewriteConfig {
 public:
   /// This specifies the order of initial traversal that populates the rewriters
   /// worklist.  When set to true, it walks the operations top-down, which is
   /// generally more efficient in compile time.  When set to false, its initial
   /// traversal of the region tree is bottom up on each block, which may match
   /// larger patterns when given an ambiguous pattern set.

 ///
 /// Returns success if the iterative process converged and no more patterns can
 /// be matched.
 LogicalResult applyOpPatternsAndFold(Operation *op,
                                      const FrozenRewritePatternSet &patterns,
                                      bool *erased = nullptr);
 
 /// Applies the specified rewrite patterns on `ops` while also trying to fold
-/// these ops as well as any other ops that were in turn created due to such
-/// rewrites. Furthermore, any pre-existing ops in the IR outside of `ops`
-/// remain completely unmodified if `strict` is set to true. If `strict` is
-/// false, other operations that use results of rewritten ops or supply operands
-/// to such ops are in turn simplified; any other ops still remain unmodified
-/// (i.e., regardless of `strict`). Note that ops in `ops` could be erased as a
-/// result of folding, becoming dead, or via pattern rewrites. If more far
-/// reaching simplification is desired, applyPatternsAndFoldGreedily should be
-/// used.
+/// 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.
+///
+/// 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.
 LogicalResult applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                                      const FrozenRewritePatternSet &patterns,
-                                     bool strict, bool *changed = nullptr);
+                                     GreedyRewriteStrictness strictMode,
+                                     bool *changed = nullptr);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

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

SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
   // 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));
   // Apply the simplification pattern to a fixpoint.
   if (failed(
-          applyOpPatternsAndFold(targets, frozenPatterns, /*strict=*/true))) {
+          applyOpPatternsAndFold(targets, frozenPatterns,
+                                 GreedyRewriteStrictness::ExistingAndNewOps))) {
     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, /*strict=*/true);
+  (void)applyOpPatternsAndFold(copyOps, frozenPatterns,
+                               GreedyRewriteStrictness::ExistingAndNewOps);
 }

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, /*strict=*/true);
+  (void)applyOpPatternsAndFold(opsToSimplify, frozenPatterns,
+                               GreedyRewriteStrictness::ExistingAndNewOps);
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 /// 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,
-                                       bool strict)
+                                       GreedyRewriteStrictness strictMode)
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
-        strictMode(strict) {}
+        strictMode(strictMode) {}
 
+  /// Performs the specified rewrites on `ops` while also trying to fold these
+  /// ops. `strictMode` controls which other ops are simplified.
+  ///
+  /// Note that ops in `ops` could be erased as a result of folding, becoming
+  /// dead, or via pattern rewrites. The return value indicates convergence.
   LogicalResult simplifyLocally(ArrayRef<Operation *> op,
                                 bool *changed = nullptr);
 
   void addToWorklist(Operation *op) override {
-    if (!strictMode || strictModeFilteredOps.contains(op))
+    if (strictMode == GreedyRewriteStrictness::AnyOp ||
+        strictModeFilteredOps.contains(op))
       GreedyPatternRewriteDriver::addSingleOpToWorklist(op);
   }
 
 private:
   void notifyOperationInserted(Operation *op) override {
-    if (strictMode)
+    if (strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
       strictModeFilteredOps.insert(op);
     GreedyPatternRewriteDriver::notifyOperationInserted(op);
   }
 
   void notifyOperationRemoved(Operation *op) override {
     GreedyPatternRewriteDriver::notifyOperationRemoved(op);
-    if (strictMode)
+    if (strictMode != GreedyRewriteStrictness::AnyOp)
       strictModeFilteredOps.erase(op);
   }
 
-  /// If `strictMode` is true, any pre-existing ops outside of
-  /// `strictModeFilteredOps` remain completely untouched by the rewrite driver.
-  /// If `strictMode` is false, operations that use results of (or supply
-  /// operands to) any rewritten ops stemming from the simplification of the
-  /// provided ops are in turn simplified; any other ops still remain untouched
-  /// (i.e., regardless of `strictMode`).
-  bool strictMode = false;
-
-  /// The list of ops we are restricting our rewrites to if `strictMode` is on.
-  /// These include the supplied set of ops as well as new ops created while
-  /// rewriting those ops. This set is not maintained when strictMode is off.
+  /// `strictMode` control which ops are added to the worklist during
+  /// simplificiation.
+  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;
 };
 
 } // namespace
 
-/// Performs the specified rewrites on `ops` while also trying to fold these ops
-/// as well as any other ops that were in turn created due to these rewrite
-/// patterns. Any pre-existing ops outside of `ops` remain completely
-/// unmodified if `strictMode` is true. If `strictMode` is false, other
-/// operations that use results of rewritten ops or supply operands to such ops
-/// are in turn simplified; any other ops still remain unmodified (i.e.,
-/// regardless of `strictMode`). Note that ops in `ops` could be erased as a
-/// result of folding, becoming dead, or via pattern rewrites. Returns true if
-/// at all any changes happened.
-// Unlike `OpPatternRewriteDriver::simplifyLocally` which works on a single op
-// or GreedyPatternRewriteDriver::simplify, this method just iterates until
-// the worklist is empty. As our objective is to keep simplification "local",
-// there is no strong rationale to re-add all operations into the worklist and
-// rerun until an iteration changes nothing. If more widereaching simplification
-// is desired, GreedyPatternRewriteDriver should be used.
 LogicalResult
 MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
                                              bool *changed) {
-  if (strictMode) {
+  if (strictMode != GreedyRewriteStrictness::AnyOp) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
   if (changed)
     *changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
     addToWorklist(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 || strictModeFilteredOps.contains(op)) &&
+    assert((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)
         *changed = true;
if (succeeded(matcher.matchAndRewrite(op, *this))) {
         *changed = true;
       ++numRewrites;
     }
   }
 
   return success(worklist.empty());
 }
 
-/// Rewrites only `op` using the supplied canonicalization patterns and
-/// folding. `erased` is set to true if the op is erased as a result of being
-/// folded, replaced, or dead.
 LogicalResult mlir::applyOpPatternsAndFold(
     Operation *op, const FrozenRewritePatternSet &patterns, bool *erased) {
   // Start the pattern driver.
   GreedyRewriteConfig config;
   OpPatternRewriteDriver driver(op->getContext(), patterns);
   bool opErased;
   LogicalResult converged =
       driver.simplifyLocally(op, config.maxNumRewrites, opErased);
   if (erased)
     *erased = opErased;
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite did not converge after "
                  << config.maxNumRewrites << " rewrites";
   });
   return converged;
 }
 
-LogicalResult
-mlir::applyOpPatternsAndFold(ArrayRef<Operation *> ops,
-                             const FrozenRewritePatternSet &patterns,
-                             bool strict, bool *changed) {
+LogicalResult mlir::applyOpPatternsAndFold(
+    ArrayRef<Operation *> ops, const FrozenRewritePatternSet &patterns,
+    GreedyRewriteStrictness strictMode, bool *changed) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
     return success();
   }
 
   // Start the pattern driver.
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
-                                     strict);
+                                     strictMode);
   return driver.simplifyLocally(ops, changed);
 }

mlir/test/Transforms/test-strict-pattern-driver.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -test-strict-pattern-driver %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect \
+// RUN:     -test-strict-pattern-driver="strictness=ExistingAndNewOps" \
+// RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-EN
 
-// CHECK-LABEL: func @test_erase
-//       CHECK:   test.arg0
-//       CHECK:   test.arg1
-//   CHECK-NOT:   test.erase_op
+// RUN: mlir-opt -allow-unregistered-dialect \
+// RUN:     -test-strict-pattern-driver="strictness=ExistingOps" \
+// RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-EX
+
+// CHECK-EN-LABEL: func @test_erase
+//       CHECK-EN:   test.arg0
+//       CHECK-EN:   test.arg1
+//   CHECK-EN-NOT:   test.erase_op
 func.func @test_erase() {
   %0 = "test.arg0"() : () -> (i32)
   %1 = "test.arg1"() : () -> (i32)
   %erase = "test.erase_op"(%0, %1) : (i32, i32) -> (i32)
   return
 }
 
-// CHECK-LABEL: func @test_insert_same_op
-//       CHECK:   "test.insert_same_op"() {skip = true}
-//       CHECK:   "test.insert_same_op"() {skip = true}
+// -----
+
+// CHECK-EN-LABEL: func @test_insert_same_op
+//       CHECK-EN:   "test.insert_same_op"() {skip = true}
+//       CHECK-EN:   "test.insert_same_op"() {skip = true}
 func.func @test_insert_same_op() {
   %0 = "test.insert_same_op"() : () -> (i32)
   return
 }
 
-// CHECK-LABEL: func @test_replace_with_new_op
-//       CHECK:   %[[n:.*]] = "test.new_op"
-//       CHECK:   "test.dummy_user"(%[[n]])
-//       CHECK:   "test.dummy_user"(%[[n]])
+// -----
+
+// CHECK-EN-LABEL: func @test_replace_with_new_op
+//       CHECK-EN:   %[[n:.*]] = "test.new_op"
+//       CHECK-EN:   "test.dummy_user"(%[[n]])
+//       CHECK-EN:   "test.dummy_user"(%[[n]])
 func.func @test_replace_with_new_op() {
   %0 = "test.replace_with_new_op"() : () -> (i32)
   %1 = "test.dummy_user"(%0) : (i32) -> (i32)
   %2 = "test.dummy_user"(%0) : (i32) -> (i32)
   return
 }
 
-// CHECK-LABEL: func @test_replace_with_erase_op
-//   CHECK-NOT:   test.replace_with_new_op
-//   CHECK-NOT:   test.erase_op
+// -----
+
+// CHECK-EN-LABEL: func @test_replace_with_erase_op
+//   CHECK-EN-NOT:   test.replace_with_new_op
+//   CHECK-EN-NOT:   test.erase_op
+
+// CHECK-EX-LABEL: func @test_replace_with_erase_op
+//   CHECK-EX-NOT:   test.replace_with_new_op
+//       CHECK-EX:   test.erase_op
 func.func @test_replace_with_erase_op() {
   "test.replace_with_new_op"() {create_erase_op} : () -> ()
   return
 }

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), /*strict=*/true);
+  (void)applyOpPatternsAndFold(copyOps, std::move(patterns),
+                               GreedyRewriteStrictness::ExistingAndNewOps);
 }
 
 namespace mlir {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>();
 }
 } // namespace mlir

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

 };
 
 struct TestStrictPatternDriver
     : public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestStrictPatternDriver)
 
   TestStrictPatternDriver() = default;
-  TestStrictPatternDriver(const TestStrictPatternDriver &other) = default;
+  TestStrictPatternDriver(const TestStrictPatternDriver &other) {
+    strictMode = other.strictMode;
+  }
 
   StringRef getArgument() const final { return "test-strict-pattern-driver"; }
   StringRef getDescription() const final {
-    return "Run strict mode of pattern driver";
+    return "Test strict mode of pattern driver";
   }
 
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
     patterns.add<InsertSameOp, ReplaceWithNewOp, EraseOp>(&getContext());
     SmallVector<Operation *> ops;
     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;
+    if (strictMode == "AnyOp") {
+      mode = GreedyRewriteStrictness::AnyOp;
+    } else if (strictMode == "ExistingAndNewOps") {
+      mode = GreedyRewriteStrictness::ExistingAndNewOps;
+    } else if (strictMode == "ExistingOps") {
+      mode = 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.
-    (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns),
-                                 /*strict=*/true);
+    (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns), mode);
   }
 
+  Option<std::string> strictMode{
+      *this, "strictness",
+      llvm::cl::desc("Can be {AnyOp, ExistingAndNewOps, ExistingOps}"),
+      llvm::cl::init("AnyOp")};
+
 private:
   // New inserted operation is valid for further transformation.
   class InsertSameOp : public RewritePattern {
   public:
     InsertSameOp(MLIRContext *context)
         : RewritePattern("test.insert_same_op", /*benefit=*/1, context) {}
 
     LogicalResult matchAndRewrite(Operation *op,