[mlir][NFC] GreedyPatternRewriteDriver: Make classes single-use

[springerm]

Summary

Less mutable state, more const. This is to address a concern about complexity of state in D140304.

Depends On: D141934

[springerm]

rebase

[springerm]

rebase

[springerm]

update

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

530–531

Is this comment update intended here or is it outdated?

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

44

Actually I'm not sure about this pattern, this is a bit unusual to have to see something like std::move(driver).simplifyLocally(ops, changed); to be able to use the API.

[springerm]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

44

I haven't seen this before either, but I found this pattern online when looking for C++ ways to indicate that an object is single use.
E.g.: https://www.foonathan.net/2018/03/rvalue-references-api-guidelines/ (4. Mark One Time Operations that Destroy the Objects)

The rvalue reference forces the caller to std::move(driver). That is an indicator for the C++ programmer that driver can no longer be used afterwards. It was already "moved" and is in an invalid state now.

Or from another website (https://crascit.com/2015/11/29/member-function-overloading-choices-you-didnt-know-you-had/):

This means the object can be modified [...] and the state of the object at the end of the call does not have to be well defined other than it must be safe to destroy the object.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 // GreedyPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 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);
+  explicit GreedyPatternRewriteDriver(
+      MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
+      const GreedyRewriteConfig &config,
+      const std::optional<DenseSet<Region *>> &scope);
 
   /// Simplify the operations within the given regions.
-  bool simplify(MutableArrayRef<Region> regions);
+  bool simplify(MutableArrayRef<Region> regions) &&;
 
   /// 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:
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
 protected:
   /// Configuration information for how to simplify.
-  GreedyRewriteConfig config;
+  const GreedyRewriteConfig config;
 
 private:
-  /// Only ops within this scope are simplified. This is set at the beginning
-  /// of `simplify()` to the current scope the rewriter operates on.
-  std::optional<DenseSet<Region *>> scope = {};
+  /// Only ops within this scope are simplified.
+  const std::optional<DenseSet<Region *>> scope = {};
 
 #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)
-    : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
+    const GreedyRewriteConfig &config,
+    const std::optional<DenseSet<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) {
-  DenseSet<Region *> &s = scope.emplace();
-  for (Region &r : regions)
-    s.insert(&r);
-
+bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) && {
 #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;
mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
   // the region containing 'op'.
   auto regionIsIsolated = [](Region &region) {
     return region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>();
   };
   (void)regionIsIsolated;
   assert(llvm::all_of(regions, regionIsIsolated) &&
          "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);
-  bool converged = driver.simplify(regions);
+  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config,
+                                    scope);
+  bool converged = std::move(driver).simplify(regions);
   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,
-                                       GreedyRewriteStrictness strictMode)
-      : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
-        strictMode(strictMode) {}
+  explicit MultiOpPatternRewriteDriver(
+      MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
+      GreedyRewriteStrictness strictMode,
+      DenseSet<Operation *> *erasedOps = nullptr)
+      : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig(),
+                                   /*scope=*/{}),
+        strictMode(strictMode), erasedOps(erasedOps) {}
 
   /// 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.
   ///
   /// All erased ops are stored in `erased`.
   LogicalResult simplifyLocally(ArrayRef<Operation *> op,
-                                bool *changed = nullptr,
-                                DenseSet<Operation *> *erased = nullptr);
+                                bool *changed = nullptr) &&;
 
 protected:
   void addSingleOpToWorklist(Operation *op) override {
     if (strictMode == GreedyRewriteStrictness::AnyOp ||
         strictModeFilteredOps.contains(op))
       GreedyPatternRewriteDriver::addSingleOpToWorklist(op);
   }
 
void notifyOperationRemoved(Operation *op) override {
     if (erasedOps)
       erasedOps->insert(op);
     if (strictMode != GreedyRewriteStrictness::AnyOp)
       strictModeFilteredOps.erase(op);
   }
 
   /// `strictMode` control which ops are added to the worklist during
   /// simplificiation.
-  GreedyRewriteStrictness strictMode = GreedyRewriteStrictness::AnyOp;
+  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 were erased.
-  DenseSet<Operation *> *erasedOps = nullptr;
+  DenseSet<Operation *> *const erasedOps = nullptr;
 };
 
 } // namespace
 
-LogicalResult MultiOpPatternRewriteDriver::simplifyLocally(
-    ArrayRef<Operation *> ops, bool *changed, DenseSet<Operation *> *erased) {
-  erasedOps = erased;
-
+LogicalResult
+MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
+                                             bool *changed) && {
   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;
     }
   }
 
-  erased = nullptr;
   return success(worklist.empty());
 }
 
 LogicalResult mlir::applyOpPatternsAndFold(
     ArrayRef<Operation *> ops, const FrozenRewritePatternSet &patterns,
     GreedyRewriteStrictness strictMode, bool *changed, bool *allOpsErased) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
     if (allOpsErased)
       *allOpsErased = true;
     return success();
   }
 
   // Start the pattern driver.
-  MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
-                                     strictMode);
   DenseSet<Operation *> erased;
-  LogicalResult converged =
-      driver.simplifyLocally(ops, changed, allOpsErased ? &erased : nullptr);
+  MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
+                                     strictMode,
+                                     allOpsErased ? &erased : nullptr);
+  LogicalResult converged = std::move(driver).simplifyLocally(ops, changed);
   if (allOpsErased)
     *allOpsErased =
         llvm::all_of(ops, [&](Operation *op) { return erased.contains(op); });
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite did not converge after "
                  << GreedyRewriteConfig().maxNumRewrites << " rewrites";
   });
   return converged;
 }