[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

106–107

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?

339

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

106–107

Changed it to how it was before.

339

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

654

What if scope is nullptr here?

760

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

763

I'm not sure about the complexity of this?

Could there be some memoization?

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

106–107

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

[springerm]

address comments

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

106–107

Done.

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

654

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

760

You are right. Rewrote the entire thing with memoization.

763

Ops are no longer checked multiple times.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

759

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

772

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

772

// 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

762

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

772

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

772

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

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;

 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.
+  /// of `simplify()` and `simplifyLocally()` 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();

 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:
   /// of ops.
   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) {
   auto cleanup = llvm::make_scope_exit([&]() { survivingOps = nullptr; });
   if (surviving) {
     survivingOps = surviving;
     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)
-    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");
+  // Fast path in case there is only one op.
+  if (ops.size() == 1)
+    return ops.front()->getParentRegion();
+
+  Region *region = ops.front()->getParentRegion();
+  ops = ops.drop_front();
+  int sz = ops.size();
+  llvm::BitVector remainingOps(sz, true);
+  do {
+    int pos = -1;
+    // Iterate over all remaining ops.
+    while ((pos = remainingOps.find_first_in(pos + 1, sz)) != -1) {
+      // Is this op contained in `region`?
+      if (region->findAncestorOpInRegion(*ops[pos]))
+        remainingOps.reset(pos);
+    }
+    if (remainingOps.none())
+      break;
+  } while ((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;
 }