[mlir] GreedyPatternRewriteDriver: Enqueue ancestors in MultiOpPatternRewriteDriver

[springerm]

Summary

The GreedyPatternRewriteDriver was extended to enqueue ancestors in D140304. With this change, MultiOpPatternRewriteDriver behaves the same way.

Note: MultiOpPatternRewriteDriver now also has a scope that limits how far we go when checking ancestors. By default, this is the first common region of all given ops.

Depends On: D141921

[springerm]

update

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

104

The comment should be updated, now the invariant is less clear: what is the use-case when this is not set?
Also when is the state reset?

336

Not clear to me that it is safe to do this without a scope, basically that means going all the way up to the module all the time?

[springerm]

address comments

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

104

Changed it to how it was before.

336

Actually, good point. I think it would make sense to add a scope parameter to the MultiOpPatternRewriter. That way the two rewriters (Region-based and Multi-op-based) behave exactly the same.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

648

What if scope is nullptr here?

754

I can't understand the condition here, are you sure it is correct? Shouldn't it be instead: while (llvm::any_of(...``?

757

I'm not sure about the complexity of this?

Could there be some memoization?

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

104

scope seems to be set also in MultiOpPatternRewriteDriver::simplifyLocally?

[springerm]

address comments

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

104

Done.

Note: It's getting difficult to understand when certain fields are set. This is addressed in D141949.

648

It cannot. The caller (applyOpPatternsAndFold) must provide a valid scope. Added an assertion.

754

You are right. Rewrote the entire thing with memoization.

757

Ops are no longer checked multiple times.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

753

Should be able to write this as SmallVector<Operation *> remainingOps = to_vector(ops.drop_front())

766

Interesting memoization!

I was thinking about memoizing the region traversal, I would think the following to be close to optimal?

DenseMap<Region*, int count> visitedRegions;
for (Operation *op : ops) {
  Region *parent = op.getParentRegion();
  while (parent) {
    auto it = visitedRegions.insert({parent, 1});
    if (!it.second) {
      // This parent region was already visited, bump counter and stop traversal
      ++it.first->second;
      break;
    }
    parent = regrion->getParentRegion();
  }
}
// From any op, traverse again the parent region: the enclosing region is the one
// that has a count equal to the number of ops.
Region *region = op.front.getParentRegion();
do {
  assert(visitedRegions.count(region));
  if (visitedRegions[region] == ops.size())
    break;
  region = region->getParentRegion();
}
return region;

[springerm]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

766

// This parent region was already visited, bump counter and stop traversal

I think the traversal has to continue here. With that the original algorithm may be faster in practice. What do you think?

Previous algo stops iterating parent regions when an ancestor for all ops was found. So I expect to iterate one, maybe two parent regions in most cases. (Assuming that often ops in the same "local area" are rewritten.) In your algorithm, we iterate all the way up to the top-level region for each op.

Previous algorithm was O(num_ops^2 + num_ops * region_depth), but the num_ops^2 is just a lookup in an array. We can optimize it by using a bit vector.

Your algorithm is O(num_ops * region_depth), so asymptotically it is better.

There was a bug where the region was set to the parent once too many in my previous code. Fixed that and using a bit vector now.

[springerm]

address comments

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

756

What about a ops = ops.drop_front() after that? (and sz = ops.size())

766

I think the traversal has to continue here.

Mmmm right, I was thinking about a "union-find" kind of algorithm but I over-simplified the data structure here...
Having to traverse to the top for each op isn't what I intended, that defeats the memoization!

In terms of complexity I didn't get how you computed it? I was seeing it more as O(num_ops^2 * region_depth) isn't it?
The findAncestorOpInRegion call is itself O(region_depth) and is in the inner loop, and it'll traverse the entire hierarchy on every miss.
(this is the call that bothered me a bit in the inside loop that I was trying to avoid)

We can optimize it by using a bit vector.

Is the bit vector changing the complexity? I'm not sure that for small bit vector it'll be more efficient that a simple vector actually, do you know?

[mehdi_amini]

I'm not convinced we have the most optimal findCommonAncestor but I guess we can revisit this later if/when it shows up on a profile!

[springerm]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

766

Right, with the current loop structure: outer loop iterations = region_depth (or num_ops), inner loop iterations = num_ops, findAncestorOpInRegion: region_depth
The bit vector doesn't change the complexity, but maybe it could be beneficial for caching effects. Only one bit needed instead of 64 to mark an op as "remaining". Or maybe it doesn't do anything...

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 /// 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.
 ///
+/// 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,
                                      bool *changed = nullptr,
-                                     bool *allErased = nullptr);
+                                     bool *allErased = nullptr,
+                                     Region *scope = nullptr);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
                                       const GreedyRewriteConfig &config);
 
   /// Simplify the operations within the given regions.
   bool simplify(MutableArrayRef<Region> regions);
 
-  /// Add the given operation to the worklist. Parent ops may or may not be
-  /// added to the worklist, depending on the type of rewrite driver. By
-  /// default, parent ops are added.
-  virtual void addToWorklist(Operation *op);
+  /// 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.
   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.
-  void addSingleOpToWorklist(Operation *op);
+  virtual 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.
   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.
   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) {
   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);
 
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
 
   /// A utility function to log a process result for the given reason.
while (Region *region = op->getParentRegion()) {
     if (!op)
       break;
     ancestors.push_back(op);
   }
 }
 
 void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *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();
if (this->op == op)
       opErasedViaPatternRewrites = true;
   }
 
   // When a root is going to be replaced, its removal will be notified as well.
   // So there is nothing to do here.
   void notifyRootReplaced(Operation *op, ValueRange replacement) override {}
 
 private:
-  /// Op that is being processed.
-  Operation *op = nullptr;
-
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
   /// Op that is being processed.
   Operation *op = nullptr;

 public:
   explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
                                        const FrozenRewritePatternSet &patterns,
                                        GreedyRewriteStrictness strictMode)
       : 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.
+  /// 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);
+                  llvm::SmallDenseSet<Operation *, 4> *surviving = nullptr,
+                  Region *scope = nullptr);
 
-  void addToWorklist(Operation *op) override {
+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)
private:
   /// An optional set of ops that survived the rewrite.
   llvm::SmallDenseSet<Operation *, 4> *survivingOps = nullptr;
 };
 
 } // namespace
 
 LogicalResult MultiOpPatternRewriteDriver::simplifyLocally(
     ArrayRef<Operation *> ops, bool *changed,
-    llvm::SmallDenseSet<Operation *, 4> *surviving) {
+    llvm::SmallDenseSet<Operation *, 4> *surviving, Region *scope) {
   if (surviving) {
     survivingOps = surviving;
     survivingOps->clear();
     survivingOps->insert(ops.begin(), ops.end());
   }
 
   if (strictMode != GreedyRewriteStrictness::AnyOp) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
+  this->scope.clear();
+  this->scope.insert(scope);
+
   if (changed)
     *changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
-    addToWorklist(op);
+    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();
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, bool *allErased) {
+/// 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();
+  while (!llvm::any_of(ops, [&](Operation *op) {
+    return !region->findAncestorOpInRegion(*op);
+  })) {
+    region = region->getParentRegion();
+    assert(region && "could not find common parent region");
+  }
+  return region;
+}
+
+LogicalResult
+mlir::applyOpPatternsAndFold(ArrayRef<Operation *> ops,
+                             const FrozenRewritePatternSet &patterns,
+                             GreedyRewriteStrictness strictMode, bool *changed,
+                             bool *allErased, Region *scope) {
   if (ops.empty()) {
     if (changed)
       *changed = false;
     if (allErased)
       *allErased = true;
     return success();
   }
 
+  if (!scope) {
+    // Compute scope if none was provided.
+    scope = findCommonAncestor(ops);
+  } else {
+    // If a scope was provided, make sure that all ops are in scope.
+#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);
+  LogicalResult converged = driver.simplifyLocally(
+      ops, changed, allErased ? &surviving : nullptr, /*scope=*/scope);
   if (allErased)
     *allErased = surviving.empty();
   return converged;
 }