[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

525–526

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

 
 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 GreedyRewriteConfig &config,
+                                      const 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;
 
-  /// Only ops within this scope are simplified. This is set at the beginning
-  /// of `simplify()` and `simplifyLocally()` to the current scope the rewriter
-  /// operates on.
-  DenseSet<Region *> scope;
+  /// Only ops within this scope are simplified.
+  const DenseSet<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)
-    : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
+    const GreedyRewriteConfig &config, const 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) {
-  scope.clear();
-  for (Region &r : regions)
-    scope.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,
+      const DenseSet<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`.
   ///
   /// 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 `ops` that survived the rewrite are stored in `surviving`.
-  LogicalResult
-  simplifyLocally(ArrayRef<Operation *> ops, bool *changed = nullptr,
-                  llvm::SmallDenseSet<Operation *, 4> *surviving = nullptr,
-                  Region *scope = nullptr);
+  /// 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);
   }
 
void notifyOperationRemoved(Operation *op) override {
     if (survivingOps)
       survivingOps->erase(op);
     if (strictMode != GreedyRewriteStrictness::AnyOp)
       strictModeFilteredOps.erase(op);
   }
 
   /// `strictMode` control which ops are added to the worklist during
   /// simplification.
-  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 survived the rewrite.
-  llvm::SmallDenseSet<Operation *, 4> *survivingOps = nullptr;
+  /// An optional set of ops that were erased.
+  llvm::SmallDenseSet<Operation *, 4> *const survivingOps = nullptr;
 };
 
 } // namespace
 
-LogicalResult MultiOpPatternRewriteDriver::simplifyLocally(
-    ArrayRef<Operation *> ops, bool *changed,
-    llvm::SmallDenseSet<Operation *, 4> *surviving, Region *scope) {
-  if (surviving) {
-    survivingOps = surviving;
+LogicalResult
+MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
+                                             bool *changed) && {
+  if (survivingOps) {
     survivingOps->clear();
     survivingOps->insert(ops.begin(), ops.end());
   }
 
   if (strictMode != GreedyRewriteStrictness::AnyOp) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
-  assert(scope && "scope is mandatory");
-  this->scope.clear();
-  this->scope.insert(scope);
-
   if (changed)
     *changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
     addSingleOpToWorklist(op);
 
   // These are scratch vectors used in the folding loop below.
while (!worklist.empty() &&
     // here.
     if (succeeded(matcher.matchAndRewrite(op, *this))) {
       if (changed)
         *changed = true;
       ++numRewrites;
     }
   }
 
-  surviving = nullptr;
   return success(worklist.empty());
 }
 
 /// Find the region that is the closest common ancestor of all given ops.
 static Region *findCommonAncestor(ArrayRef<Operation *> ops) {
   assert(!ops.empty() && "expected at least one op");
   Region *region = ops.front()->getParentRegion();
   SmallVector<Operation *> remainingOps(ops.drop_front().begin(),
#ifndef NDEBUG
     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.
-  MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
-                                     strictMode);
   llvm::SmallDenseSet<Operation *, 4> surviving;
-  LogicalResult converged = driver.simplifyLocally(
-      ops, changed, allErased ? &surviving : nullptr, /*scope=*/scope);
+  DenseSet<Region *> scopeSet;
+  scopeSet.insert(findCommonAncestor(ops));
+  MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
+                                     scopeSet, 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;
 }