[Canonicalize] Reapply logic to suppport operating top-down on the initial pass, NFC.

[lattner]

Summary

This adds all the support for top-down traversals and constant prefolding that was
reverted in this patch:
https://reviews.llvm.org/D99329

This has two changes from the previous patch:

1. It disables the 'processExistingConstants' prefolding call when doing a bottom-up traversal.
2. It defaults applyPatternsAndFoldGreedily to bottom-up traversal instead of top-down.

The consequence of these two changes is that there is *no behavior change* from this
patch. We can stage in the new behavior on a case by case basis over time.

[lattner]

FWIW, the change to the testsuite isn't necessary. I am not sure how to retroactively drop them from the patch, but if someone wants that and can help me with the right arc/git command then i'd be happy to do so!

> Done.

[lattner]

Microoptimize patch to reduce out a few comment deltas

[lattner]

Remove the testcase changes, further tweeze out spurious comment change.

[mehdi_amini]

There are two changes here: the option to change the (initial) traversal order and the constant materialization change, can you split the constant materialization in a follow up patch?

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

48

Can you update the doc?

[lattner]

Yes, I can split it if there is a good reason to. This patch was approved before, why does it need to be split now? This seems like needless work.

I'll update PatternRewriter.md, thanks for remembering that!

[lattner]

Document the new argument to applyPatternsAndFoldGreedily in both comments and markdown dox.

[jpienaar]

mlir/include/mlir/Transforms/FoldUtils.h

37

Wouldn't DialectFoldInterface::shouldMaterializeInto affect where these get moved to rather than always to entry?

[lattner]

I have no idea what shouldMaterializeInto is about. I see the comment in FoldInterfaces.h but it doesn't make much sense to me, and there are no clients of this in tree.

Can you give a concrete example of when this gets used?

[mehdi_amini]

In D102117#2749750, @lattner wrote:

I have no idea what shouldMaterializeInto is about. I see the comment in FoldInterfaces.h but it doesn't make much sense to me, and there are no clients of this in tree.

Can you give a concrete example of when this gets used?

It is used in the test dialect as an example.
In general this allows to decouple the "IsolatedFromAbove" property from the point of insertions for constants.
For example you may have something like:

func @foo() {
  ...
  // not isolated from above
  remote.dispatch() {
     ...
     op.fold() ... // Folding this op, we want the constant to be created in the `remote.dispatch` region and not in the parent region.
    ..
  }
  ...
}

The logic is in lib/Transforms/Utils/FoldUtils.cpp getInsertionRegion(). When folding op.fold() above, it'll walk the parent ops until it finds an "IsolatedFromAbove" operation or one that implement this interface.

[bondhugula]

mlir/docs/PatternRewriter.md

255–257

As far as the hierarchy goes, the traversal is actually pre order in both cases. (Parents are visited before children in the existing scheme and in the one you are adding.) Within a block, the old one is doing bottom up and the one being added is doing top down. This description needs an update? Perhaps you could just drop "post-order" and instead mention that parent ops are processed before child ops in both cases?

[lattner]

mlir/include/mlir/Transforms/FoldUtils.h

37

This implementation already handles DialectFoldInterface::shouldMaterializeInto correctly because it uses getInsertionRegion which handles it. I'll update the commend and add a testcase to show this.

[lattner]

Update the comment to make it more clear that the logic obeys the
'shouldMaterializeInto' hook.

Add a command line flag for canonicalize that enables the new behavior

Add a testcase showing that we don't hoist constants out of a
test.one_region_op, which implements the shouldMaterializeInto hook.

[lattner]

Ok, I think this is really really ready to land as a staging patch. As a reminder, this does not change behavior of any passes in tree or out of tree. This is intended to get the logic into the MLIR repo, so we can gently adopt it in the ecosystem over time. Can I get a LGTM please?

[lattner]

Clarify the documention as Uday suggests, thanks!

[bondhugula]

Looks good to me in general, but several comments and some potential minor improvements.

mlir/docs/PatternRewriter.md

255

It's "pre-order over the region tree" because the old scheme is doing post-order + reverse which would be pre-order since this is a directed tree. With the new scheme, you are doing pre-order + reverse + reverse, which is again pre-order. The original traversal is post order but since the rewriter is popping from the end, we might refer to the combined thing here?

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

42–44

Update here.

-> walks the operations in a block top-down

-> When set to false, operations in a block are processed bottom-up, which may match ....

?

mlir/include/mlir/Transforms/Passes.td

367

The convention for cmd-line option names isn't camel case I think - just hyphen separated. "enable" can be dropped as well.

enableTopDown -> top-down ?

369

-> `Process operations in a block top-down" ?

mlir/lib/Transforms/Utils/FoldUtils.cpp

97

auto -> ConstantMap.

102

You can easily make this return true when the op is erased - and let the subsequent walk go ahead if the op isn't erased. It is rare though that a constant op has a region, but no harm in that walk continuing.

104

Assert on single result for op?

123

This can be simplified into a single comparison. You don't need the constantIterator == insertBlock.end() part. If constantIterator is at the end, then it won't be equal to op. So,

if (constantIterator != Block::iterator(op))

?

129

Comment here please.

135

Side comment: even if you didn't delete it, a constant op will probably never have regions!

You can make processConstant return true/false and let the walk continue?

139

Nit: This comment doesn't seem to match logic below.

is isolated -> has a different region to insert constants into

The op may or may not be isolated.

344–346

auto -> Operation please.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

140

-> ... for operations in a block.

mlir/test/Transforms/canonicalize-td.mlir

6

TD-LABEL:
BU-LABEL:

10

BU -> BU-NEXT?

25

Same here. This isn't really being looked at by FileCheck?

27

BU -> BU-NEXT?

30

TD -> TD-NEXT?

[lattner]

Thank you for the feedback Uday, incorporated. At this point, this patch has been thoroughly nit-picked. I plan to land it tomorrow and we can continue to iterate in tree, thanks.

mlir/docs/PatternRewriter.md

255

It is worse than this. The original order is extremely hard to describe because of this code that it is using:

template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
          typename RetT = detail::walkResultType<FnT>>
typename std::enable_if<std::is_same<RetT, void>::value, RetT>::type
walk(FnT &&callback) {
  for (auto &block : *this)
    block.walk<Order>(callback);
}

It is calling it with the default post-order traversal, but notice that it walks the blocks in the region in-order! I will generalize the writing because it isn't the point of this to go into all the details, this is an overview doc.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

42–44

tweaked again, generalizing the claim. The old walk isn't principled as mentioned above.

mlir/include/mlir/Transforms/Passes.td

367

Right, thanks

369

Made more pedantically correct.

mlir/lib/Transforms/Utils/FoldUtils.cpp

97

No. This is following the convention of the file.

102

I think this approach is ok.

123

Nice improvement, yes.

139

I only just recently learned about this distinction. I will fix.

344–346

I'm changing this, but this has nothing to do with my patch. I don't think it is good form to try to impose your opinion of auto on patches touching general code like this.

mlir/test/Transforms/canonicalize-td.mlir

6

nice catch thx

25

No this is being ignored.

[lattner]

Incorporate feedback from Uday

[bondhugula]

mlir/docs/PatternRewriter.md

259

Typo: "is may"

mlir/lib/Transforms/Utils/FoldUtils.cpp

344–346

All your changes to this method have nothing to do with this patch - in that spirit, I just pointed this out since I felt it improved readability further.

[lattner]

One more typo

[lattner]

Thank you for the review Uday!

[lattner]

Landed in 648f34a2840b thanks all

mlir/docs/PatternRewriter.md

 between, and within dialects using a concept of "legality". This framework
 allows for transforming illegal operations to those supported by a provided
 conversion target, via a set of pattern-based operation rewriting patterns. This
 framework also provides support for type conversions. More information on this
 driver can be found [here](DialectConversion.md).
 
 ### Greedy Pattern Rewrite Driver
 
-This driver performs a post order traversal over the provided operations and
-greedily applies the patterns that locally have the most benefit. The benefit of
+This driver walks the provided operations and greedily applies the patterns that
+locally have the most benefit. The benefit of
 a pattern is decided solely by the benefit specified on the pattern, and the
 relative order of the pattern within the pattern list (when two patterns have
 the same local benefit). Patterns are iteratively applied to operations until a
 fixed point is reached, at which point the driver finishes. This driver may be
 used via the following: `applyPatternsAndFoldGreedily` and
 `applyOpPatternsAndFold`. The latter of which only applies patterns to the
 provided operation, and will not traverse the IR.
 
+The driver is configurable and supports two modes: 1) you may opt-in to a
+"top-down" traversal, which seeds the worklist with each operation top down and
+in a pre-order over the region tree.  This is generally more efficient in
+compile time.  2) the default is a "bottom up" traversal, which builds the
+initial worklist with a postorder traversal of the region tree.  This may
+match larger patterns with ambiguous pattern sets.
+
 Note: This driver is the one used by the [canonicalization](Canonicalization.md)
 [pass](Passes.md#-canonicalize-canonicalize-operations) in MLIR.

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

 
 //===----------------------------------------------------------------------===//
 // 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);
+                             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.
-LogicalResult
-applyPatternsAndFoldGreedily(Operation *op,
-                             const FrozenRewritePatternSet &patterns,
-                             unsigned maxIterations);
+LogicalResult applyPatternsAndFoldGreedily(
+    Operation *op, const FrozenRewritePatternSet &patterns,
+    unsigned maxIterations, bool useTopDownTraversal = false);
 
 /// Rewrite the given regions, which must be isolated from above.
 LogicalResult
 applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
-                             const FrozenRewritePatternSet &patterns);
+                             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);
+LogicalResult applyPatternsAndFoldGreedily(
+    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
+    unsigned maxIterations, bool useTopDownTraversal = false);
 
 /// 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/include/mlir/Transforms/Passes.td

def Canonicalizer : Pass<"canonicalize"> {
   let summary = "Canonicalize operations";
   let description = [{
     This pass performs various types of canonicalizations over a set of
     operations. See [Operation Canonicalization](Canonicalization.md) for more
     details.
   }];
   let constructor = "mlir::createCanonicalizerPass()";
   let dependentDialects = ["memref::MemRefDialect"];
+  let options = [
+    Option<"topDownProcessingEnabled", "top-down", "bool",
+           /*default=*/"false",
+           "Seed the worklist in general top-down order">
+  ];
 }
 
 def CSE : Pass<"cse"> {
   let summary = "Eliminate common sub-expressions";
   let description = [{
     This pass implements a generalized algorithm for common sub-expression
     elimination. This pass relies on information provided by the
     `Memory SideEffect` interface to identify when it is safe to eliminate

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);
+    (void)applyPatternsAndFoldGreedily(
+        getOperation()->getRegions(), patterns,
+        /*maxIterations=*/10, /*useTopDownTraversal=*/
+        topDownProcessingEnabled);
   }
 
   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.

 
 /// Try to get or create a new constant entry. On success this returns the
 /// constant operation value, nullptr otherwise.
 Operation *OperationFolder::tryGetOrCreateConstant(
     ConstantMap &uniquedConstants, Dialect *dialect, OpBuilder &builder,
     Attribute value, Type type, Location loc) {
   // Check if an existing mapping already exists.
   auto constKey = std::make_tuple(dialect, value, type);
-  auto *&constInst = uniquedConstants[constKey];
-  if (constInst)
-    return constInst;
+  Operation *&constOp = uniquedConstants[constKey];
+  if (constOp)
+    return constOp;
 
   // If one doesn't exist, try to materialize one.
-  if (!(constInst = materializeConstant(dialect, builder, value, type, loc)))
+  if (!(constOp = materializeConstant(dialect, builder, value, type, loc)))
     return nullptr;
 
   // Check to see if the generated constant is in the expected dialect.
-  auto *newDialect = constInst->getDialect();
+  auto *newDialect = constOp->getDialect();
   if (newDialect == dialect) {
-    referencedDialects[constInst].push_back(dialect);
-    return constInst;
+    referencedDialects[constOp].push_back(dialect);
+    return constOp;
   }
 
   // If it isn't, then we also need to make sure that the mapping for the new
   // dialect is valid.
   auto newKey = std::make_tuple(newDialect, value, type);
 
   // If an existing operation in the new dialect already exists, delete the
   // materialized operation in favor of the existing one.
   if (auto *existingOp = uniquedConstants.lookup(newKey)) {
-    constInst->erase();
+    constOp->erase();
     referencedDialects[existingOp].push_back(dialect);
-    return constInst = existingOp;
+    return constOp = existingOp;
   }
 
   // Otherwise, update the new dialect to the materialized operation.
-  referencedDialects[constInst].assign({dialect, newDialect});
-  auto newIt = uniquedConstants.insert({newKey, constInst});
+  referencedDialects[constOp].assign({dialect, newDialect});
+  auto newIt = uniquedConstants.insert({newKey, constOp});
   return newIt.first->second;
 }

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)
-      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
+                                      const FrozenRewritePatternSet &patterns,
+                                      bool useTopDownTraversal)
+      : PatternRewriter(ctx), matcher(patterns), folder(ctx),
+        useTopDownTraversal(useTopDownTraversal) {
     worklist.reserve(64);
 
     // Apply a simple cost model based solely on pattern benefit.
     matcher.applyDefaultCostModel();
   }
 
   bool simplify(MutableArrayRef<Region> regions, int maxIterations);
 
private:
   /// need to be revisited, plus their index in the worklist.  This allows us to
   /// 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;
 };
 } // 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) {
-  // Add the given operation to the worklist.
-  auto collectOps = [this](Operation *op) { addToWorklist(op); };
+  // 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 changed = false;
-  int i = 0;
+  int iteration = 0;
   do {
-    // Add all nested operations to the worklist.
+    worklist.clear();
+    worklistMap.clear();
+
+    // Add all nested operations to the worklist in preorder.
     for (auto &region : regions)
-      region.walk(collectOps);
+      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.
     SmallVector<Value, 8> originalOperands, resultValues;
 
     changed = false;
     while (!worklist.empty()) {
       auto *op = popFromWorklist();
 
while (!worklist.empty()) {
         // to revisit them.
         for (auto result : op->getResults())
           for (auto *userOp : result.getUsers())
             addToWorklist(userOp);
 
         notifyOperationRemoved(op);
       };
 
+      // Add the given operation to the worklist.
+      auto collectOps = [this](Operation *op) { addToWorklist(op); };
+
       // Try to fold this op.
       bool inPlaceUpdate;
       if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
                                       &inPlaceUpdate)))) {
         changed = true;
         if (!inPlaceUpdate)
           continue;
       }
 
       // Try to match one of the patterns. The rewriter is automatically
-      // notified of any necessary changes, so there is nothing else to do here.
+      // 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.
+    // After applying patterns, make sure that the CFG of each of the regions
+    // is kept up to date.
     changed |= succeeded(simplifyRegions(*this, regions));
-  } while (changed && ++i < maxIterations);
+  } while (changed && ++iteration < 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) {
-  return applyPatternsAndFoldGreedily(op, patterns, maxPatternMatchIterations);
-}
-LogicalResult
-mlir::applyPatternsAndFoldGreedily(Operation *op,
                                    const FrozenRewritePatternSet &patterns,
-                                   unsigned maxIterations) {
-  return applyPatternsAndFoldGreedily(op->getRegions(), patterns,
-                                      maxIterations);
+                                   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) {
-  return applyPatternsAndFoldGreedily(regions, patterns,
-                                      maxPatternMatchIterations);
-}
-LogicalResult
-mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
                                    const FrozenRewritePatternSet &patterns,
-                                   unsigned maxIterations) {
+                                   bool useTopDownTraversal) {
+  return applyPatternsAndFoldGreedily(
+      regions, patterns, maxPatternMatchIterations, useTopDownTraversal);
+}
+LogicalResult mlir::applyPatternsAndFoldGreedily(
+    MutableArrayRef<Region> regions, const FrozenRewritePatternSet &patterns,
+    unsigned maxIterations, bool useTopDownTraversal) {
   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);
+  GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns,
+                                    useTopDownTraversal);
   bool converged = driver.simplify(regions, maxIterations);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << maxIterations << " times\n";
   });
   return success(converged);
 }
 

mlir/test/Transforms/canonicalize-td.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(canonicalize{top-down=true})' | FileCheck %s --check-prefix=TD
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(canonicalize)' | FileCheck %s --check-prefix=BU
+
+
+// BU-LABEL: func @default_insertion_position
+// TD-LABEL: func @default_insertion_position
+func @default_insertion_position(%cond: i1) {
+  // Constant should be folded into the entry block.
+
+  // BU: constant 2
+  // BU-NEXT: scf.if
+
+  // TD: constant 2
+  // TD-NEXT: scf.if
+  scf.if %cond {
+    %0 = constant 1 : i32
+    %2 = addi %0, %0 : i32
+    "foo.yield"(%2) : (i32) -> ()
+  }
+  return
+}
+
+// This shows that we don't pull the constant out of the region because it
+// wants to be the insertion point for the constant.
+// BU-LABEL: func @custom_insertion_position
+// TD-LABEL: func @custom_insertion_position
+func @custom_insertion_position() {
+  // BU: test.one_region_op
+  // BU-NEXT: constant 2
+
+  // TD: test.one_region_op
+  // TD-NEXT: constant 2
+  "test.one_region_op"() ({
+
+    %0 = constant 1 : i32
+    %2 = addi %0, %0 : i32
+    "foo.yield"(%2) : (i32) -> ()
+  }) : () -> ()
+  return
+}
+