Merge with mainline.

[rriddle]

Accidental change? Was this supposed to be landed as D102117?

[lattner]

Yeah I'm not sure what was going on here, I fumbled a git rebase. I'm sorry for the noise!

[bondhugula]

In D102636#2764142, @lattner wrote:

Yeah I'm not sure what was going on here, I fumbled a git rebase. I'm sorry for the noise!

In the git history, I see this commit appearing as:

commit 648f34a2840b75f4081884052f2ccb11f62f8209
Author: Chris Lattner <[e-mail]>
Date:   Mon May 17 09:51:45 2021 -0700

    Merge with mainline.
    
    Differential Revision: https://reviews.llvm.org/D102636

Revert and push with the right commit message if it's worth it while we have the chance? (since no one has likely touched those files)

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 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
+}
+