[mlir] GreedyPatternRewriteDriver: Keep track of surviving ops

[springerm]

Summary

This change adds allErased to the applyOpPatternsAndFold(ArrayRef<Operation *>, ...) overload. This overload now supports all functionality that is also supported by applyOpPatternsAndFold(Operation *, ...) and can be used as a replacement.

This change has no performance implications when allErased = nullptr.

The single-operation overload is removed in a subsequent NFC change.

Depends On: D141904

[springerm]

The purpose of this change is to remove OpPatternRewriteDriver to avoid duplicate code.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

609

Since an operation can be erased once, we should be able to use a vector here?

I guess the issue is that vector does not help later if the intent is checking if an op has been erased.

That said this all seems a bit "expensive" overall here, and I'm not sure that the extra bool *allOpsErased justified this (in itself it seems already like a bit "ad-hoc" in the API).

[springerm]

address comments

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

609

The purpose of this revision is to support all functionality of OpPatternRewriteDriver, so that it can be removed and we don't have to maintain two simplifyLocally implementations (and we get the existing debug log messages for free).

How about instead of storing all erased ops, we maintain a set of "surviving ops". It would initially be populated with the list of ops that is passed to simplifyLocally. That way there is an upper bound on the size of the set and it will only shrink, never grow. It has the same overhead as strictModeFilteredOps. Also, the set is maintained only if allErased is specified. So if allErased == nullptr, there is no performance implication at all.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

624

Can you clarify that these are ops from the initially provided list of ops?

721

Can you use RAII ( llvm/include/llvm/ADT/ScopeExit.h ) for this?

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 ///   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.
+/// `allOpsErased` is set to true if all ops in `ops` were erased.
 LogicalResult applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                                      const FrozenRewritePatternSet &patterns,
                                      GreedyRewriteStrictness strictMode,
-                                     bool *changed = nullptr);
+                                     bool *changed = nullptr,
+                                     bool *allErased = nullptr);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
         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);
+  ///
+  /// All `ops` that survived the rewrite are stored in `surviving`.
+  LogicalResult
+  simplifyLocally(ArrayRef<Operation *> ops, bool *changed = nullptr,
+                  llvm::SmallDenseSet<Operation *, 4> *surviving = nullptr);
 
   void addToWorklist(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
   /// 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;
+
+  /// An optional set of ops that survived the rewrite.
+  llvm::SmallDenseSet<Operation *, 4> *survivingOps = nullptr;
 };
 
 } // namespace
 
-LogicalResult
-MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
-                                             bool *changed) {
+LogicalResult MultiOpPatternRewriteDriver::simplifyLocally(
+    ArrayRef<Operation *> ops, bool *changed,
+    llvm::SmallDenseSet<Operation *, 4> *surviving) {
+  if (surviving) {
+    survivingOps = surviving;
+    survivingOps->clear();
+    survivingOps->insert(ops.begin(), ops.end());
+  }
+
   if (strictMode != GreedyRewriteStrictness::AnyOp) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
   if (changed)
     *changed = false;
   worklist.clear();
while (!worklist.empty() &&
     // here.
     if (succeeded(matcher.matchAndRewrite(op, *this))) {
       if (changed)
         *changed = true;
       ++numRewrites;
     }
   }
 
+  surviving = nullptr;
   return success(worklist.empty());
 }
 
 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,
-    GreedyRewriteStrictness strictMode, bool *changed) {
+    GreedyRewriteStrictness strictMode, bool *changed, bool *allErased) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
+    if (allErased)
+      *allErased = true;
     return success();
   }
 
   // Start the pattern driver.
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
                                      strictMode);
-  return driver.simplifyLocally(ops, changed);
+  llvm::SmallDenseSet<Operation *, 4> surviving;
+  LogicalResult converged =
+      driver.simplifyLocally(ops, changed, allErased ? &surviving : nullptr);
+  if (allErased)
+    *allErased = surviving.empty();
+  return converged;
 }

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

 // RUN: mlir-opt -allow-unregistered-dialect \
 // RUN:     -test-strict-pattern-driver="strictness=ExistingAndNewOps" \
 // RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-EN
 
 // 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-SAME:     pattern_driver_all_erased = true, pattern_driver_changed = true}
 //       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-EN-LABEL: func @test_insert_same_op
+//  CHECK-EN-SAME:     {pattern_driver_all_erased = false, pattern_driver_changed = true}
 //       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-EN-LABEL: func @test_replace_with_new_op
+//  CHECK-EN-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //       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-EN-LABEL: func @test_replace_with_erase_op
+//  CHECK-EN-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //   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-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //   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/Test/TestPatterns.cpp

public:
   }
 
   StringRef getArgument() const final { return "test-strict-pattern-driver"; }
   StringRef getDescription() const final {
     return "Test strict mode of pattern driver";
   }
 
   void runOnOperation() override {
-    mlir::RewritePatternSet patterns(&getContext());
-    patterns.add<InsertSameOp, ReplaceWithNewOp, EraseOp>(&getContext());
+    MLIRContext *ctx = &getContext();
+    mlir::RewritePatternSet patterns(ctx);
+    patterns.add<InsertSameOp, ReplaceWithNewOp, EraseOp>(ctx);
     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), mode);
+    bool changed = false;
+    bool allErased = false;
+    (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns), mode,
+                                 &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",
       llvm::cl::desc("Can be {AnyOp, ExistingAndNewOps, ExistingOps}"),
       llvm::cl::init("AnyOp")};
 
 private: