[GreedyPatternRewriter] Introduce a config object that allows controlling internal parameters. NFC.

[lattner]

Summary

This exposes the iterations and top-down processing as flags, and also
allows controlling whether region simplification is desirable for a client.
This allows deleting some duplicated entrypoints to
applyPatternsAndFoldGreedily.

This also deletes the Constant Preprocessing pass, which isn't worth it
on balance.

All defaults are all kept the same, so no one should see a behavior change.

[mehdi_amini]

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

26

bottom-up?

Or "reverse-post-order" to be exact (we populate the worklist in post-order, but then pop-back from the worklist).

[lattner]

Thank you for the review!

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

26

This is the problem - it isn't RPO or bottom up. It is RPO on the region tree but it is bottom up on the operations within each block. This hybrid approach isn't actually good for the "bottom up" case either. Clarified the comment though. Thanks!

[lattner]

Update comment, silence a warning.

[mehdi_amini]

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

26

If you consider blocks part of the "walk tree", where the list of operations in the block are just "child", isn't it just RPO?
(inside a block, what is the difference between bottom up and RPO?)

[rriddle]

LGTM

[lattner]

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

26

No, it isn't just RPO. RPO on a linear sequence is top-down. Consider that basic blocks in a CFG is just an optimization: it is also possible to model every operation as having successors (most operations would have one). in that case, RPO would be top down.

[jpienaar]

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

27

Note: this also changes more than just affecting ambigious patterns. We have cases such as:

%381 = X.opA (%arg1161 : !X.T) {
  ...
  X.yield %11842 : tensor<2304x3x96xf32>
}

%1108 = X.opA (%arg1161 : !X.T) {
  ...
  %11842 = X.opB 2 %381[%arg1161] : (tensor<2304x24x96xf32>, !X.T) -> tensor<2304x24x12xf32>
  ...
  X.yield %11846 : tensor<1x1024x24x12xf32>
}

In bottom up %1108 is seen as unused and so removed (along with all the ops inside its region) and so when %381 is encountered it is known dead, but in top down only the operands of the dead op %1108 is added back to worklist and %381 has already been processed so it doesn't get added back to worklist. This results in the bottom up finishing in 1 iteration, while top-down takes 157 iterations in canonicalizer (as the above form ends up being quite long chains of this pattern). So for removal of deadcode, bottom up is performing much better.

To have top down and bottom up behave the same here, one would need to enqueue all the operands of ops nested inside dead op that processes a value referenced directly (rather than passed in as argument) to worklist, as these ops can become dead now too.

mlir/include/mlir/Transforms/FoldUtils.h

 //===--------------------------------------------------------------------===//
 
 /// A utility class for folding operations, and unifying duplicated constants
 /// generated along the way.
 class OperationFolder {
 public:
   OperationFolder(MLIRContext *ctx) : interfaces(ctx) {}
 
-  /// Scan the specified region for constants that can be used in folding,
-  /// moving them to the entry block (or any custom insertion location specified
-  /// by shouldMaterializeInto), and add them to our known-constants table.
-  void processExistingConstants(Region &region);
-
   /// Tries to perform folding on the given `op`, including unifying
   /// deduplicated constants. If successful, replaces `op`'s uses with
   /// folded results, and returns success. `preReplaceAction` is invoked on `op`
   /// before it is replaced. 'processGeneratedConstants' is invoked for any new
   /// operations generated when folding. If the op was completely folded it is
   /// erased. If it is just updated in place, `inPlaceUpdate` is set to true.
   LogicalResult
   tryToFold(Operation *op,

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 
 #ifndef MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_
 #define MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_
 
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 
 namespace mlir {
 
+/// This struct allows control over how the GreedyPatternRewriteDriver works.
+struct GreedyRewriteConfig {
+  /// This specifies the order of initial traversal that populates the rewriters
+  /// worklist.  When set to true, it walks the operations top-down, which is
+  /// generally more efficient in compile time.  When set to false, its initial
+  /// traversal of the region tree is bottom up on each block, which may match
+  /// larger patterns when given an ambiguous pattern set.
+  bool useTopDownTraversal = false;
+
+  // Perform control flow optimizations to the region tree after applying all
+  // patterns.
+  bool enableRegionSimplification = true;
+
+  /// This specifies the maximum number of times the rewriter will iterate
+  /// between applying patterns and simplifying regions.
+  unsigned maxIterations = 10;
+};
+
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner.
 ///
 /// This variant may stop after a predefined number of iterations, see the
 /// alternative below to provide a specific number of iterations before stopping
 /// in absence of convergence.
 ///
 /// Return success if the iterative process converged and no more patterns can
 /// be matched in the result operation regions.
 ///
 /// Note: This does not apply patterns to the top-level operation itself.
 ///       These methods also perform folding and simple dead-code elimination
 ///       before attempting to match any of the provided patterns.
 ///
-/// You may choose the order of initial traversal with the `useTopDownTraversal`
-/// boolean.  When set to true, it walks the operations top-down, which is
-/// generally more efficient in compile time.  When set to false, its initial
-/// traversal of the region tree is post-order, which may match larger patterns
-/// when given an ambiguous pattern set.
-LogicalResult
-applyPatternsAndFoldGreedily(Operation *op,
-                             const FrozenRewritePatternSet &patterns,
-                             bool useTopDownTraversal = false);
-
-/// Rewrite the regions of the specified operation, with a user-provided limit
-/// on iterations to attempt before reaching convergence.
+/// You may configure several aspects of this with GreedyRewriteConfig.
 LogicalResult applyPatternsAndFoldGreedily(
-    Operation *op, const FrozenRewritePatternSet &patterns,
-    unsigned maxIterations, bool useTopDownTraversal = false);
+    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
+    GreedyRewriteConfig config = GreedyRewriteConfig());
 
 /// Rewrite the given regions, which must be isolated from above.
-LogicalResult
-applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
-                             const FrozenRewritePatternSet &patterns,
-                             bool useTopDownTraversal = false);
-
-/// Rewrite the given regions, with a user-provided limit on iterations to
-/// attempt before reaching convergence.
-LogicalResult applyPatternsAndFoldGreedily(
-    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
-    unsigned maxIterations, bool useTopDownTraversal = false);
+inline LogicalResult applyPatternsAndFoldGreedily(
+    Operation *op, const FrozenRewritePatternSet &patterns,
+    GreedyRewriteConfig config = GreedyRewriteConfig()) {
+  return applyPatternsAndFoldGreedily(op->getRegions(), patterns, config);
+}
 
 /// Applies the specified patterns on `op` alone while also trying to fold it,
 /// by selecting the highest benefits patterns in a greedy manner. Returns
 /// success if no more patterns can be matched. `erased` is set to true if `op`
 /// was folded away or erased as a result of becoming dead. Note: This does not
 /// apply any patterns recursively to the regions of `op`.
 LogicalResult applyOpPatternsAndFold(Operation *op,
                                      const FrozenRewritePatternSet &patterns,
                                      bool *erased = nullptr);
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Transforms/Canonicalizer.cpp

struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
   LogicalResult initialize(MLIRContext *context) override {
     RewritePatternSet owningPatterns(context);
     for (auto *op : context->getRegisteredOperations())
       op->getCanonicalizationPatterns(owningPatterns, context);
     patterns = std::move(owningPatterns);
     return success();
   }
   void runOnOperation() override {
-    (void)applyPatternsAndFoldGreedily(
-        getOperation()->getRegions(), patterns,
-        /*maxIterations=*/10, /*useTopDownTraversal=*/
-        topDownProcessingEnabled);
+    GreedyRewriteConfig config;
+    config.useTopDownTraversal = topDownProcessingEnabled;
+    (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns,
+                                       config);
   }
 
   FrozenRewritePatternSet patterns;
 };
 } // end anonymous namespace
 
 /// Create a Canonicalizer pass.
 std::unique_ptr<Pass> mlir::createCanonicalizerPass() {
   return std::make_unique<Canonicalizer>();
 }

mlir/lib/Transforms/Utils/FoldUtils.cpp

if (ConstantOp::isBuildableWith(value, type))
     return builder.create<ConstantOp>(loc, type, value);
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // OperationFolder
 //===----------------------------------------------------------------------===//
 
-/// Scan the specified region for constants that can be used in folding,
-/// moving them to the entry block (or any custom insertion location specified
-/// by shouldMaterializeInto), and add them to our known-constants table.
-void OperationFolder::processExistingConstants(Region &region) {
-  if (region.empty())
-    return;
-
-  // March the constant insertion point forward, moving all constants to the
-  // top of the block, but keeping them in their order of discovery.
-  Region *insertRegion = getInsertionRegion(interfaces, &region.front());
-  auto &uniquedConstants = foldScopes[insertRegion];
-
-  Block &insertBlock = insertRegion->front();
-  Block::iterator constantIterator = insertBlock.begin();
-
-  // Process each constant that we discover in this region.
-  auto processConstant = [&](Operation *op, Attribute value) {
-    assert(op->getNumResults() == 1 && "constants have one result");
-    // Check to see if we already have an instance of this constant.
-    Operation *&constOp = uniquedConstants[std::make_tuple(
-        op->getDialect(), value, op->getResult(0).getType())];
-
-    // If we already have an instance of this constant, CSE/delete this one as
-    // we go.
-    if (constOp) {
-      if (constantIterator == Block::iterator(op))
-        ++constantIterator; // Don't invalidate our iterator when scanning.
-      op->getResult(0).replaceAllUsesWith(constOp->getResult(0));
-      op->erase();
-      return;
-    }
-
-    // Otherwise, remember that we have this constant.
-    constOp = op;
-    referencedDialects[op].push_back(op->getDialect());
-
-    // If the constant isn't already at the insertion point then move it up.
-    if (constantIterator != Block::iterator(op))
-      op->moveBefore(&insertBlock, constantIterator);
-    else
-      ++constantIterator; // It was pointing at the constant.
-  };
-
-  // Collect all the constants for this region of isolation or insertion (as
-  // specified by the shouldMaterializeInto hook).  Collect any subregions of
-  // isolation/constant insertion for subsequent processing.
-  SmallVector<Operation *> insertionSubregionOps;
-  region.walk<WalkOrder::PreOrder>([&](Operation *op) {
-    // If this is a constant, process it.
-    Attribute value;
-    if (matchPattern(op, m_Constant(&value))) {
-      processConstant(op, value);
-      // We may have deleted the operation, don't check it for regions.
-      return WalkResult::skip();
-    }
-
-    // If the operation has regions and is isolated, don't recurse into it.
-    if (op->getNumRegions() != 0) {
-      auto hasDifferentInsertRegion = [&](Region &region) {
-        return !region.empty() &&
-               getInsertionRegion(interfaces, &region.front()) != insertRegion;
-      };
-      if (llvm::any_of(op->getRegions(), hasDifferentInsertRegion)) {
-        insertionSubregionOps.push_back(op);
-        return WalkResult::skip();
-      }
-    }
-
-    // Otherwise keep going.
-    return WalkResult::advance();
-  });
-
-  // Process regions in any isolated ops separately.
-  for (Operation *subregionOps : insertionSubregionOps) {
-    for (Region &region : subregionOps->getRegions())
-      processExistingConstants(region);
-  }
-}
-
 LogicalResult OperationFolder::tryToFold(
     Operation *op, function_ref<void(Operation *)> processGeneratedConstants,
     function_ref<void(Operation *)> preReplaceAction, bool *inPlaceUpdate) {
   if (inPlaceUpdate)
     *inPlaceUpdate = false;
 
   // If this is a unique'd constant, return failure as we know that it has
   // already been folded.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "pattern-matcher"
 
-/// The max number of iterations scanning for pattern match.
-static unsigned maxPatternMatchIterations = 10;
-
 //===----------------------------------------------------------------------===//
 // 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,
-                                      bool useTopDownTraversal)
-      : PatternRewriter(ctx), matcher(patterns), folder(ctx),
-        useTopDownTraversal(useTopDownTraversal) {
+                                      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 simplify(MutableArrayRef<Region> regions, int maxIterations);
+  bool simplify(MutableArrayRef<Region> regions);
 
   void addToWorklist(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);
private:
   /// efficiently remove operations from the worklist when they are erased, even
   /// if they aren't the root of a pattern.
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
-  /// Whether to use a top-down or bottom-up traversal to seed the initial
-  /// worklist.
-  bool useTopDownTraversal;
+  /// Configuration information for how to simplify.
+  GreedyRewriteConfig config;
 };
 } // end anonymous namespace
 
 /// Performs the rewrites while folding and erasing any dead ops. Returns true
 /// if the rewrite converges in `maxIterations`.
-bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions,
-                                          int maxIterations) {
-  // For maximum compatibility with existing passes, do not process existing
-  // constants unless we're performing a top-down traversal.
-  // TODO: This is just for compatibility with older MLIR, remove this.
-  if (useTopDownTraversal) {
-    // Perform a prepass over the IR to discover constants.
-    for (auto &region : regions)
-      folder.processExistingConstants(region);
-  }
-
+bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
   bool changed = false;
-  int iteration = 0;
+  unsigned iteration = 0;
   do {
     worklist.clear();
     worklistMap.clear();
 
+    if (!config.useTopDownTraversal) {
+      // Add operations to the worklist in postorder.
+      for (auto &region : regions)
+        region.walk([this](Operation *op) { addToWorklist(op); });
+    } else {
       // Add all nested operations to the worklist in preorder.
       for (auto &region : regions)
-      if (useTopDownTraversal)
         region.walk<WalkOrder::PreOrder>(
             [this](Operation *op) { worklist.push_back(op); });
-      else
-        region.walk([this](Operation *op) { addToWorklist(op); });
 
-    if (useTopDownTraversal) {
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());
       // Remember the reverse index.
       for (size_t i = 0, e = worklist.size(); i != e; ++i)
         worklistMap[worklist[i]] = i;
     }
 
     // These are scratch vectors used in the folding loop below.
while (!worklist.empty()) {
       // Try to match one of the patterns. The rewriter is automatically
       // notified of any necessary changes, so there is nothing else to do
       // here.
       changed |= succeeded(matcher.matchAndRewrite(op, *this));
     }
 
     // After applying patterns, make sure that the CFG of each of the regions
     // is kept up to date.
+    if (config.enableRegionSimplification)
       changed |= succeeded(simplifyRegions(*this, regions));
-  } while (changed && ++iteration < maxIterations);
+  } while (changed && ++iteration < config.maxIterations);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return success if no more patterns can be matched
 /// in the result operation regions. Note: This does not apply patterns to the
 /// top-level operation itself.
 ///
 LogicalResult
-mlir::applyPatternsAndFoldGreedily(Operation *op,
-                                   const FrozenRewritePatternSet &patterns,
-                                   bool useTopDownTraversal) {
-  return applyPatternsAndFoldGreedily(op, patterns, maxPatternMatchIterations,
-                                      useTopDownTraversal);
-}
-LogicalResult mlir::applyPatternsAndFoldGreedily(
-    Operation *op, const FrozenRewritePatternSet &patterns,
-    unsigned maxIterations, bool useTopDownTraversal) {
-  return applyPatternsAndFoldGreedily(op->getRegions(), patterns, maxIterations,
-                                      useTopDownTraversal);
-}
-/// Rewrite the given regions, which must be isolated from above.
-LogicalResult
 mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
                                    const FrozenRewritePatternSet &patterns,
-                                   bool useTopDownTraversal) {
-  return applyPatternsAndFoldGreedily(
-      regions, patterns, maxPatternMatchIterations, useTopDownTraversal);
-}
-LogicalResult mlir::applyPatternsAndFoldGreedily(
-    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
-    unsigned maxIterations, bool useTopDownTraversal) {
+                                   GreedyRewriteConfig config) {
   if (regions.empty())
     return success();
 
   // The top-level operation must be known to be isolated from above to
   // prevent performing canonicalizations on operations defined at or above
   // 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");
 
   // Start the pattern driver.
-  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns,
-                                    useTopDownTraversal);
-  bool converged = driver.simplify(regions, maxIterations);
+  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config);
+  bool converged = driver.simplify(regions);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
-                 << maxIterations << " times\n";
+                 << config.maxIterations << " times\n";
   });
   return success(converged);
 }
 
 //===----------------------------------------------------------------------===//
 // OpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 

 }
 
 /// Rewrites only `op` using the supplied canonicalization patterns and
 /// folding. `erased` is set to true if the op is erased as a result of being
 /// folded, replaced, or dead.
 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, maxPatternMatchIterations, opErased);
+      driver.simplifyLocally(op, config.maxIterations, opErased);
   if (erased)
     *erased = opErased;
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
-                 << maxPatternMatchIterations << " times";
+                 << config.maxIterations << " times";
   });
   return converged;
 }