[mlir] Refactor RewritePatternMatcher into a new PatternApplicator class.

[rriddle]

Summary

This class enables for abstracting more of the details for the rewrite process, and will allow for clients to apply specific cost models to the pattern list. This allows for DialectConversion and the GreedyPatternRewriter to share the same underlying matcher implementation. This also simplifies the plumbing necessary to support dynamic patterns.

[ftynse]

mlir/include/mlir/IR/PatternMatch.h

492

Do we expect null patterns? Otherwise, we may want to consider (const) RewritePattern & as function argument.

503

Let's also mention that this gets called if onSuccess return failure, because technically the pattern did match in that case.

mlir/lib/IR/PatternMatch.cpp

185

Would it be reasonable, performance-wise, to populate this map within the second for (auto &it : patterns) loop below for each root kind separately, then clear it at the end of the loop? We still get the benefit of only computing the benefit once since the same pattern cannot have different root kinds, and get the benefit of smaller memory footprint in this map.

196–197

This will permanently remove the patterns from the applicator. If want to apply a different cost model later, under which these patterns will no longer be impossible to match, we won't be able to. I'm not sure this is desirable.

[bondhugula]

PatternApplicator sounds like a weird name. PatternApplier / PatterCostApplier / PatterApplyManager ?

[rriddle]

Address comments

[rriddle]

Thanks for the review Alex.

mlir/include/mlir/IR/PatternMatch.h

503

Changed to not call onFailure if onSuccess fails. This should keep onFailure specific to pattern match failures.

mlir/lib/IR/PatternMatch.cpp

185

Switched to SmallDenseMap. Also added a fast path for size() == 1, which by and large the most common case.

196–197

Fixed, thanks for catching. This won't affect any existing usages, but at some point in the future we may apply more than one cost model to a specific applicator.

[rriddle]

In D81985#2098568, @bondhugula wrote:

PatternApplicator sounds like a weird name. PatternApplier / PatterCostApplier / PatterApplyManager ?

Applicator seems to like the correct english word to use though?

[rriddle]

Rebase

[ftynse]

Great, thanks River!

[nicolasvasilache]

Thanks for pushing on this @rriddle !

One question that stems from what I see here is whether you have recommendations on debugging.
I can imagine pitfalls related to applying this extra level of indirection and understanding what happens in what order.
Do you have plans for extra debugging tooling?

[bondhugula]

In D81985#2098781, @rriddle wrote:

In D81985#2098568, @bondhugula wrote:

PatternApplicator sounds like a weird name. PatternApplier / PatterCostApplier / PatterApplyManager ?

Applicator seems to like the correct english word to use though?

The dictionary says it's specifically used for a device that does such and all usage I saw on the web was in that device context; never heard it spoken though.

[rriddle]

In D81985#2100514, @nicolasvasilache wrote:

Thanks for pushing on this @rriddle !

One question that stems from what I see here is whether you have recommendations on debugging.
I can imagine pitfalls related to applying this extra level of indirection and understanding what happens in what order.
Do you have plans for extra debugging tooling?

I do have plans for some additional debug tooling, but I see that as mostly orthogonal to this change. For the end user, this change doesn't really affect the debugging experience realistically. Aside from this, there are several problems with providing pattern debug tooling:

• We need to pipe in the debug options from somewhere, or provide a global options cache(similar to the context, printer, etc.). This is kind of annoying because the only alternative is to push the options up to every conversion/pattern pass, which would be a lot of unnecessary additional boiler plate that most users don't want/need IMO.
• The internal workings of our two main drivers are vastly different, meaning that the debug output/format/experience is different for each driver. This makes it difficult to have generalized tooling that integrates nicely, leading to some amount of duplication.

That being said, there are a few things that I *want* to do that the above make more annoying. I currently have plans to add:

• print-after-all, print-before-all, etc. related functionality for printing in-between patterns.
• disable-patterns=* that allows for disabling specific patterns
• I'd also love to have an interactive command line driver for patterns (and passes too really) that users can interact with(i.e., decide if a pattern should be applied, print the current IR, etc.), but I'm not entirely sure how I'd pipe that through ATM.

• River

mlir/include/mlir/IR/PatternMatch.h

 ///
 /// This also has a sentinel representation that can be used for patterns that
 /// fail to match.
 ///
 class PatternBenefit {
   enum { ImpossibleToMatchSentinel = 65535 };
 
 public:
-  /*implicit*/ PatternBenefit(unsigned benefit);
+  PatternBenefit() : representation(ImpossibleToMatchSentinel) {}
+  PatternBenefit(unsigned benefit);
   PatternBenefit(const PatternBenefit &) = default;
   PatternBenefit &operator=(const PatternBenefit &) = default;
 
   static PatternBenefit impossibleToMatch() { return PatternBenefit(); }
   bool isImpossibleToMatch() const { return *this == impossibleToMatch(); }
 
   /// If the corresponding pattern can match, return its benefit.  If the
   // corresponding pattern isImpossibleToMatch() then this aborts.
   unsigned short getBenefit() const;
 
   bool operator==(const PatternBenefit &rhs) const {
     return representation == rhs.representation;
   }
   bool operator!=(const PatternBenefit &rhs) const { return !(*this == rhs); }
   bool operator<(const PatternBenefit &rhs) const {
     return representation < rhs.representation;
   }
+  bool operator>(const PatternBenefit &rhs) const { return rhs < *this; }
+  bool operator<=(const PatternBenefit &rhs) const { return !(*this > rhs); }
+  bool operator>=(const PatternBenefit &rhs) const { return !(*this < rhs); }
 
 private:
-  PatternBenefit() : representation(ImpossibleToMatchSentinel) {}
   unsigned short representation;
 };
 
 //===----------------------------------------------------------------------===//
 // Pattern class
 //===----------------------------------------------------------------------===//
 
 /// Instances of Pattern can be matched against SSA IR.  These matches get used
private:
   /// uses of op with uses of newOp.
   void replaceOpWithResultsOfAnotherOp(Operation *op, Operation *newOp);
 };
 
 //===----------------------------------------------------------------------===//
 // Pattern-driven rewriters
 //===----------------------------------------------------------------------===//
 
+//===----------------------------------------------------------------------===//
+// OwningRewritePatternList
+
 class OwningRewritePatternList {
   using PatternListT = std::vector<std::unique_ptr<RewritePattern>>;
 
 public:
   OwningRewritePatternList() = default;
 
   /// Construct a OwningRewritePatternList populated with the pattern `t` of
   /// type `T`.
   template <typename T>
   OwningRewritePatternList(T &&t) {
     patterns.emplace_back(std::make_unique<T>(std::forward<T>(t)));
   }
 
   PatternListT::iterator begin() { return patterns.begin(); }
   PatternListT::iterator end() { return patterns.end(); }
   PatternListT::const_iterator begin() const { return patterns.begin(); }
   PatternListT::const_iterator end() const { return patterns.end(); }
+  PatternListT::size_type size() const { return patterns.size(); }
   void clear() { patterns.clear(); }
 
   //===--------------------------------------------------------------------===//
   // Pattern Insertion
   //===--------------------------------------------------------------------===//
 
   /// Add an instance of each of the pattern types 'Ts' to the pattern list with
   /// the given arguments. Return a reference to `this` for chaining insertions.
   /// Note: ConstructorArg is necessary here to separate the two variadic lists.
   template <typename... Ts, typename ConstructorArg,
             typename... ConstructorArgs,
             typename = std::enable_if_t<sizeof...(Ts) != 0>>
   OwningRewritePatternList &insert(ConstructorArg &&arg,
                                    ConstructorArgs &&... args) {
     // The following expands a call to emplace_back for each of the pattern
     // types 'Ts'. This magic is necessary due to a limitation in the places
     // that a parameter pack can be expanded in c++11.
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
-    using dummy = int[];
-    (void)dummy{
+    (void)std::initializer_list<int>{
         0, (patterns.emplace_back(std::make_unique<Ts>(arg, args...)), 0)...};
     return *this;
   }
 
+  /// Add the given pattern to the pattern list.
+  void insert(std::unique_ptr<RewritePattern> pattern) {
+    patterns.emplace_back(std::move(pattern));
+  }
+
 private:
   PatternListT patterns;
 };
 
-/// This class manages optimization and execution of a group of rewrite
-/// patterns, providing an API for finding and applying, the best match against
-/// a given node.
-///
-class RewritePatternMatcher {
+//===----------------------------------------------------------------------===//
+// PatternApplicator
+
+/// This class manages the application of a group of rewrite patterns, with a
+/// user-provided cost model.
+class PatternApplicator {
 public:
-  /// Create a RewritePatternMatcher with the specified set of patterns.
-  explicit RewritePatternMatcher(const OwningRewritePatternList &patterns);
+  /// The cost model dynamically assigns a PatternBenefit to a particular
+  /// pattern. Users can query contained patterns and pass analysis results to
+  /// applyCostModel. Patterns to be discarded should have a benefit of
+  /// `impossibleToMatch`.
+  using CostModel = function_ref<PatternBenefit(const RewritePattern &)>;
+
+  explicit PatternApplicator(const OwningRewritePatternList &owningPatternList)
+      : owningPatternList(owningPatternList) {}
+
+  /// Attempt to match and rewrite the given op with any pattern, allowing a
+  /// predicate to decide if a pattern can be applied or not, and hooks for if
+  /// the pattern match was a success or failure.
+  ///
+  /// canApply:  called before each match and rewrite attempt; return false to
+  ///            skip pattern.
+  /// onFailure: called when a pattern fails to match to perform cleanup.
+  /// onSuccess: called when a pattern match succeeds; return failure() to
+  ///            invalidate the match and try another pattern.
+  LogicalResult matchAndRewrite(
+      Operation *op, PatternRewriter &rewriter,
+      function_ref<bool(const RewritePattern &)> canApply = {},
+      function_ref<void(const RewritePattern &)> onFailure = {},
+      function_ref<LogicalResult(const RewritePattern &)> onSuccess = {});
+
+  /// Apply a cost model to the patterns within this applicator.
+  void applyCostModel(CostModel model);
+
+  /// Apply the default cost model that solely uses the pattern's static
+  /// benefit.
+  void applyDefaultCostModel() {
+    applyCostModel(
+        [](const RewritePattern &pattern) { return pattern.getBenefit(); });
+  }
 
-  /// Try to match the given operation to a pattern and rewrite it. Return
-  /// true if any pattern matches.
-  bool matchAndRewrite(Operation *op, PatternRewriter &rewriter);
+  /// Walk all of the rewrite patterns within the applicator.
+  void walkAllPatterns(function_ref<void(const RewritePattern &)> walk);
 
 private:
-  RewritePatternMatcher(const RewritePatternMatcher &) = delete;
-  void operator=(const RewritePatternMatcher &) = delete;
+  /// The list that owns the patterns used within this applicator.
+  const OwningRewritePatternList &owningPatternList;
 
-  /// The group of patterns that are matched for optimization through this
-  /// matcher.
-  std::vector<RewritePattern *> patterns;
+  /// The set of patterns to match for each operation, stable sorted by benefit.
+  DenseMap<OperationName, SmallVector<RewritePattern *, 2>> patterns;
 };
 
+//===----------------------------------------------------------------------===//
+// 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. Return true if no more patterns can be matched in
-/// the result operation regions.
-/// Note: This does not apply patterns to the top-level operation itself.
-/// Note: These methods also perform folding and simple dead-code elimination
+/// 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. Note:
+///       These methods also perform folding and simple dead-code elimination
 ///       before attempting to match any of the provided patterns.
 ///
-bool applyPatternsAndFoldGreedily(Operation *op,
-                                  const OwningRewritePatternList &patterns);
+LogicalResult
+applyPatternsAndFoldGreedily(Operation *op,
 /// Rewrite the given regions, which must be isolated from above.
-bool applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
-                                  const OwningRewritePatternList &patterns);
+LogicalResult
+applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
 
 /// Applies the specified patterns on `op` alone while also trying to fold it,
-/// by selecting the highest benefits patterns in a greedy manner. Returns true
-/// 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
+/// 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`.
-bool applyOpPatternsAndFold(Operation *op,
+LogicalResult applyOpPatternsAndFold(Operation *op,
                                      const OwningRewritePatternList &patterns,
                                      bool *erased = nullptr);
 } // end namespace mlir
 
 #endif // MLIR_PATTERN_MATCH_H

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

LogicalResult mlir::linalg::applyStagedPatterns(
     Operation *op, ArrayRef<OwningRewritePatternList> stage1Patterns,
     const OwningRewritePatternList &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda) {
   unsigned iteration = 0;
   (void)iteration;
   for (const auto &patterns : stage1Patterns) {
     LLVM_DEBUG(DBGS() << "Before 1st stage, iter: " << ++iteration << "\n"
                       << *op);
-    if (!applyPatternsAndFoldGreedily(op, patterns)) {
+    if (failed(applyPatternsAndFoldGreedily(op, patterns))) {
       LLVM_DEBUG(DBGS() << "Underlying first stage rewrite did not converge");
       return failure();
     }
     LLVM_DEBUG(DBGS() << "After 1st stage, iter: " << ++iteration << "\n"
                       << *op);
-    if (!applyPatternsAndFoldGreedily(op, stage2Patterns)) {
+    if (failed(applyPatternsAndFoldGreedily(op, stage2Patterns))) {
       LLVM_DEBUG(DBGS() << "Underlying 2nd stage rewrite did not converge");
       return failure();
     }
     LLVM_DEBUG(DBGS() << "After 2nd stage, iter : " << iteration << "\n"
                       << *op);
     if (stage3Lambda) {
       if (failed(stage3Lambda(op)))
         return failure();
       LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
                         << *op);
     }
   }
   return success();
 }

mlir/lib/IR/PatternMatch.cpp

 using namespace mlir;
 
 PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {
   assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&
          "This pattern match benefit is too large to represent");
 }
 
 unsigned short PatternBenefit::getBenefit() const {
-  assert(representation != ImpossibleToMatchSentinel &&
-         "Pattern doesn't match");
+  assert(!isImpossibleToMatch() && "Pattern doesn't match");
   return representation;
 }
 
 //===----------------------------------------------------------------------===//
 // Pattern implementation
 //===----------------------------------------------------------------------===//
 
 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,

 void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {
   cloneRegionBefore(region, *before->getParent(), before->getIterator());
 }
 
 //===----------------------------------------------------------------------===//
 // PatternMatcher implementation
 //===----------------------------------------------------------------------===//
 
-RewritePatternMatcher::RewritePatternMatcher(
-    const OwningRewritePatternList &patterns) {
-  for (auto &pattern : patterns)
-    this->patterns.push_back(pattern.get());
-
-  // Sort the patterns by benefit to simplify the matching logic.
-  std::stable_sort(this->patterns.begin(), this->patterns.end(),
-                   [](RewritePattern *l, RewritePattern *r) {
-                     return r->getBenefit() < l->getBenefit();
-                   });
+void PatternApplicator::applyCostModel(CostModel model) {
+  // Separate patterns by root kind to simplify lookup later on.
+  patterns.clear();
+  for (const auto &pat : owningPatternList)
+    patterns[pat->getRootKind()].push_back(pat.get());
+
+  // Sort the patterns using the provided cost model.
+  llvm::SmallDenseMap<RewritePattern *, PatternBenefit> benefits;
+  auto cmp = [&benefits](RewritePattern *lhs, RewritePattern *rhs) {
+    return benefits[lhs] > benefits[rhs];
+  };
+  for (auto &it : patterns) {
+    SmallVectorImpl<RewritePattern *> &list = it.second;
+
+    // Special case for one pattern in the list, which is the most common case.
+    if (list.size() == 1) {
+      if (model(*list.front()).isImpossibleToMatch())
+        list.clear();
+      continue;
+    }
+
+    // Collect the dynamic benefits for the current pattern list.
+    benefits.clear();
+    for (RewritePattern *pat : list)
+      benefits.try_emplace(pat, model(*pat));
+
+    // Sort patterns with highest benefit first, and remove those that are
+    // impossible to match.
+    std::stable_sort(list.begin(), list.end(), cmp);
+    while (!list.empty() && benefits[list.back()].isImpossibleToMatch())
+      list.pop_back();
+  }
+}
+
+void PatternApplicator::walkAllPatterns(
+    function_ref<void(const RewritePattern &)> walk) {
+  for (auto &it : owningPatternList)
+    walk(*it);
 }
 
 /// Try to match the given operation to a pattern and rewrite it.
-bool RewritePatternMatcher::matchAndRewrite(Operation *op,
-                                            PatternRewriter &rewriter) {
-  for (auto *pattern : patterns) {
-    // Ignore patterns that are for the wrong root or are impossible to match.
-    if (pattern->getRootKind() != op->getName() ||
-        pattern->getBenefit().isImpossibleToMatch())
+LogicalResult PatternApplicator::matchAndRewrite(
+    Operation *op, PatternRewriter &rewriter,
+    function_ref<bool(const RewritePattern &)> canApply,
+    function_ref<void(const RewritePattern &)> onFailure,
+    function_ref<LogicalResult(const RewritePattern &)> onSuccess) {
+  auto patternIt = patterns.find(op->getName());
+  if (patternIt == patterns.end())
+    return failure();
+
+  for (auto *pattern : patternIt->second) {
+    // Check that the pattern can be applied.
+    if (canApply && !canApply(*pattern))
       continue;
 
     // Try to match and rewrite this pattern. The patterns are sorted by
-    // benefit, so if we match we can immediately rewrite and return.
-    if (succeeded(pattern->matchAndRewrite(op, rewriter)))
-      return true;
+    // benefit, so if we match we can immediately rewrite.
+    rewriter.setInsertionPoint(op);
+    if (succeeded(pattern->matchAndRewrite(op, rewriter))) {
+      if (!onSuccess || succeeded(onSuccess(*pattern)))
+        return success();
+      continue;
+    }
+
+    if (onFailure)
+      onFailure(*pattern);
   }
-  return false;
+  return failure();
 }

mlir/lib/Transforms/DialectConversion.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // OperationLegalizer
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A set of rewrite patterns that can be used to legalize a given operation.
-using LegalizationPatterns = SmallVector<RewritePattern *, 1>;
+using LegalizationPatterns = SmallVector<const RewritePattern *, 1>;
 
 /// This class defines a recursive operation legalizer.
 class OperationLegalizer {
 public:
   using LegalizationAction = ConversionTarget::LegalizationAction;
 
   OperationLegalizer(ConversionTarget &targetInfo,
-                     const OwningRewritePatternList &patterns)
-      : target(targetInfo) {
-    buildLegalizationGraph(patterns);
-    computeLegalizationGraphBenefit();
-  }
+                     const OwningRewritePatternList &patterns);
 
   /// Returns if the given operation is known to be illegal on the target.
   bool isIllegal(Operation *op) const;
 
   /// Attempt to legalize the given operation. Returns success if the operation
   /// was legalized, failure otherwise.
   LogicalResult legalize(Operation *op, ConversionPatternRewriter &rewriter);
 
   /// Returns the conversion target in use by the legalizer.
   ConversionTarget &getTarget() { return target; }
 
 private:
   /// Attempt to legalize the given operation by folding it.
   LogicalResult legalizeWithFold(Operation *op,
                                  ConversionPatternRewriter &rewriter);
 
-  /// Attempt to legalize the given operation by applying the provided pattern.
-  /// Returns success if the operation was legalized, failure otherwise.
-  LogicalResult legalizePattern(Operation *op, RewritePattern *pattern,
-                                ConversionPatternRewriter &rewriter);
+  /// Attempt to legalize the given operation by applying a pattern. Returns
+  /// success if the operation was legalized, failure otherwise.
+  LogicalResult legalizeWithPattern(Operation *op,
+                                    ConversionPatternRewriter &rewriter);
+
+  /// Return true if the given pattern may be applied to the given operation,
+  /// false otherwise.
+  bool canApplyPattern(Operation *op, const RewritePattern &pattern,
 
+  /// Legalize the resultant IR after successfully applying the given pattern.
+  LogicalResult legalizePatternResult(Operation *op,
+                                      const RewritePattern &pattern,
+                                      ConversionPatternRewriter &rewriter,
+                                      RewriterState &curState);
+
   /// Build an optimistic legalization graph given the provided patterns. This
   /// function populates 'legalizerPatterns' with the operations that are not
   /// directly legal, but may be transitively legal for the current target given
   /// the provided patterns.
-  void buildLegalizationGraph(const OwningRewritePatternList &patterns);
+  void buildLegalizationGraph(
+      DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
 
   /// Compute the benefit of each node within the computed legalization graph.
   /// This orders the patterns within 'legalizerPatterns' based upon two
   /// criteria:
   ///  1) Prefer patterns that have the lowest legalization depth, i.e.
   ///     represent the more direct mapping to the target.
   ///  2) When comparing patterns with the same legalization depth, prefer the
   ///     pattern with the highest PatternBenefit. This allows for users to
   ///     prefer specific legalizations over others.
-  void computeLegalizationGraphBenefit();
+  void computeLegalizationGraphBenefit(
+      DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
 
   /// The current set of patterns that have been applied.
-  SmallPtrSet<RewritePattern *, 8> appliedPatterns;
-
-  /// The set of legality information for operations transitively supported by
-  /// the target.
-  DenseMap<OperationName, LegalizationPatterns> legalizerPatterns;
+  SmallPtrSet<const RewritePattern *, 8> appliedPatterns;
 
   /// The legalization information provided by the target.
   ConversionTarget &target;
+
+  /// The pattern applicator to use for conversions.
+  PatternApplicator applicator;
 };
 } // namespace
 
+OperationLegalizer::OperationLegalizer(ConversionTarget &targetInfo,
+                                       const OwningRewritePatternList &patterns)
+    : target(targetInfo), applicator(patterns) {
+  // The set of legality information for operations transitively supported by
+  // the target.
+  DenseMap<OperationName, LegalizationPatterns> legalizerPatterns;
+
+  buildLegalizationGraph(legalizerPatterns);
+  computeLegalizationGraphBenefit(legalizerPatterns);
+}
+
 bool OperationLegalizer::isIllegal(Operation *op) const {
   // Check if the target explicitly marked this operation as illegal.
   return target.getOpAction(op->getName()) == LegalizationAction::Illegal;
 }
 
 LogicalResult
 OperationLegalizer::legalize(Operation *op,
                              ConversionPatternRewriter &rewriter) {
if (succeeded(legalizeWithFold(op, rewriter))) {
     LLVM_DEBUG({
       logSuccess(rewriterImpl.logger, "operation was folded");
       rewriterImpl.logger.startLine() << logLineComment;
     });
     return success();
   }
 
   // Otherwise, we need to apply a legalization pattern to this operation.
-  auto it = legalizerPatterns.find(op->getName());
-  if (it == legalizerPatterns.end()) {
-    LLVM_DEBUG({
-      logFailure(rewriterImpl.logger, "no known legalization path");
-      rewriterImpl.logger.startLine() << logLineComment;
-    });
-    return failure();
-  }
-
-  // The patterns are sorted by expected benefit, so try to apply each in-order.
-  for (auto *pattern : it->second) {
-    if (succeeded(legalizePattern(op, pattern, rewriter))) {
-      LLVM_DEBUG({
-        logSuccess(rewriterImpl.logger, "");
-        rewriterImpl.logger.startLine() << logLineComment;
-      });
-      return success();
-    }
-  }
+  if (succeeded(legalizeWithPattern(op, rewriter))) {
 
   LLVM_DEBUG({
     logFailure(rewriterImpl.logger, "no matched legalization pattern");
     rewriterImpl.logger.startLine() << logLineComment;
   });
   return failure();
 }
 
for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
     }
   }
 
   LLVM_DEBUG(logSuccess(rewriterImpl.logger, ""));
   return success();
 }
 
 LogicalResult
-OperationLegalizer::legalizePattern(Operation *op, RewritePattern *pattern,
+OperationLegalizer::legalizeWithPattern(Operation *op,
                                         ConversionPatternRewriter &rewriter) {
   auto &rewriterImpl = rewriter.getImpl();
+
+  // Functor that returns if the given pattern may be applied.
+  auto canApply = [&](const RewritePattern &pattern) {
+    return canApplyPattern(op, pattern, rewriter);
+  };
+
+  // Functor that cleans up the rewriter state after a pattern failed to match.
+  RewriterState curState = rewriterImpl.getCurrentState();
+  auto onFailure = [&](const RewritePattern &pattern) {
+    LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern failed to match"));
+    rewriterImpl.resetState(curState);
+    appliedPatterns.erase(&pattern);
+  };
+
+  // Functor that performs additional legalization when a pattern is
+  // successfully applied.
+  auto onSuccess = [&](const RewritePattern &pattern) {
+    auto result = legalizePatternResult(op, pattern, rewriter, curState);
+    appliedPatterns.erase(&pattern);
+    if (failed(result))
+      rewriterImpl.resetState(curState);
+    return result;
+  };
+
+  // Try to match and rewrite a pattern on this operation.
+  return applicator.matchAndRewrite(op, rewriter, canApply, onFailure,
+                                    onSuccess);
+}
+
+bool OperationLegalizer::canApplyPattern(Operation *op,
+                                         const RewritePattern &pattern,
+                                         ConversionPatternRewriter &rewriter) {
   LLVM_DEBUG({
-    auto &os = rewriterImpl.logger;
+    auto &os = rewriter.getImpl().logger;
     os.getOStream() << "\n";
-    os.startLine() << "* Pattern : '" << pattern->getRootKind() << " -> (";
-    llvm::interleaveComma(pattern->getGeneratedOps(), llvm::dbgs());
+    os.startLine() << "* Pattern : '" << pattern.getRootKind() << " -> (";
+    llvm::interleaveComma(pattern.getGeneratedOps(), llvm::dbgs());
     os.getOStream() << ")' {\n";
     os.indent();
   });
 
   // Ensure that we don't cycle by not allowing the same pattern to be
   // applied twice in the same recursion stack if it is not known to be safe.
-  if (!pattern->hasBoundedRewriteRecursion() &&
-      !appliedPatterns.insert(pattern).second) {
-    LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern was already applied"));
-    return failure();
-  }
+  if (!pattern.hasBoundedRewriteRecursion() &&
+      !appliedPatterns.insert(&pattern).second) {
+    LLVM_DEBUG(
+        logFailure(rewriter.getImpl().logger, "pattern was already applied"));
+    return false;
+  }
+  return true;
 
-  RewriterState curState = rewriterImpl.getCurrentState();
-  auto cleanupFailure = [&] {
-    // Reset the rewriter state and pop this pattern.
-    rewriterImpl.resetState(curState);
-    appliedPatterns.erase(pattern);
-    return failure();
-  };
+LogicalResult OperationLegalizer::legalizePatternResult(
+    Operation *op, const RewritePattern &pattern,
+    ConversionPatternRewriter &rewriter, RewriterState &curState) {
+  auto &rewriterImpl = rewriter.getImpl();
 
-  // Try to rewrite with the given pattern.
-  rewriter.setInsertionPoint(op);
-  LogicalResult matchedPattern = pattern->matchAndRewrite(op, rewriter);
 #ifndef NDEBUG
   assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates");
 #endif
 
-  if (failed(matchedPattern)) {
-    LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern failed to match"));
-    return cleanupFailure();
-  }
-
   // If the pattern moved or created any blocks, try to legalize their types.
   // This ensures that the types of the block arguments are legal for the region
   // they were moved into.
   for (unsigned i = curState.numBlockActions,
                 e = rewriterImpl.blockActions.size();
        i != e; ++i) {
     auto &action = rewriterImpl.blockActions[i];
     if (action.kind ==
             ConversionPatternRewriterImpl::BlockActionKind::TypeConversion ||
         action.kind == ConversionPatternRewriterImpl::BlockActionKind::Erase)
       continue;
 
     // Convert the block signature.
     if (failed(rewriterImpl.convertBlockSignature(action.block))) {
       LLVM_DEBUG(logFailure(rewriterImpl.logger,
                             "failed to convert types of moved block"));
-      return cleanupFailure();
+      return failure();
     }
   }
 
   // Check all of the replacements to ensure that the pattern actually replaced
   // the root operation. We also mark any other replaced ops as 'dead' so that
   // we don't try to legalize them later.
   bool replacedRoot = false;
   for (unsigned i = curState.numReplacements,
#endif
   for (unsigned i = curState.numRootUpdates,
                 e = rewriterImpl.rootUpdates.size();
        i != e; ++i) {
     auto &state = rewriterImpl.rootUpdates[i];
     if (failed(legalize(state.getOperation(), rewriter))) {
       LLVM_DEBUG(logFailure(rewriterImpl.logger,
                             "operation updated in-place '{0}' was illegal",
                             op->getName()));
-      return cleanupFailure();
+      return failure();
     }
   }
 
   // Recursively legalize each of the new operations.
   for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
        i != e; ++i) {
     Operation *op = rewriterImpl.createdOps[i];
     if (failed(legalize(op, rewriter))) {
       LLVM_DEBUG(logFailure(rewriterImpl.logger,
                             "generated operation '{0}'({1}) was illegal",
                             op->getName(), op));
-      return cleanupFailure();
+      return failure();
     }
   }
 
   LLVM_DEBUG(logSuccess(rewriterImpl.logger, "pattern applied successfully"));
-  appliedPatterns.erase(pattern);
   return success();
 }
 
 void OperationLegalizer::buildLegalizationGraph(
-    const OwningRewritePatternList &patterns) {
+    DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
   // A mapping between an operation and a set of operations that can be used to
   // generate it.
   DenseMap<OperationName, SmallPtrSet<OperationName, 2>> parentOps;
   // A mapping between an operation and any currently invalid patterns it has.
-  DenseMap<OperationName, SmallPtrSet<RewritePattern *, 2>> invalidPatterns;
+  DenseMap<OperationName, SmallPtrSet<const RewritePattern *, 2>>
+      invalidPatterns;
   // A worklist of patterns to consider for legality.
-  llvm::SetVector<RewritePattern *> patternWorklist;
+  llvm::SetVector<const RewritePattern *> patternWorklist;
 
   // Build the mapping from operations to the parent ops that may generate them.
-  for (auto &pattern : patterns) {
-    auto root = pattern->getRootKind();
+  applicator.walkAllPatterns([&](const RewritePattern &pattern) {
+    OperationName root = pattern.getRootKind();
 
     // Skip operations that are always known to be legal.
     if (target.getOpAction(root) == LegalizationAction::Legal)
-      continue;
+      return;
 
     // Add this pattern to the invalid set for the root op and record this root
     // as a parent for any generated operations.
-    invalidPatterns[root].insert(pattern.get());
-    for (auto op : pattern->getGeneratedOps())
+    invalidPatterns[root].insert(&pattern);
+    for (auto op : pattern.getGeneratedOps())
       parentOps[op].insert(root);
 
     // Add this pattern to the worklist.
-    patternWorklist.insert(pattern.get());
-  }
+    patternWorklist.insert(&pattern);
+  });
 
   while (!patternWorklist.empty()) {
     auto *pattern = patternWorklist.pop_back_val();
 
     // Check to see if any of the generated operations are invalid.
     if (llvm::any_of(pattern->getGeneratedOps(), [&](OperationName op) {
           Optional<LegalizationAction> action = target.getOpAction(op);
           return !legalizerPatterns.count(op) &&
                  (!action || action == LegalizationAction::Illegal);
         }))
       continue;
 
     // Otherwise, if all of the generated operation are valid, this op is now
     // legal so add all of the child patterns to the worklist.
     legalizerPatterns[pattern->getRootKind()].push_back(pattern);
     invalidPatterns[pattern->getRootKind()].erase(pattern);
 
     // Add any invalid patterns of the parent operations to see if they have now
     // become legal.
     for (auto op : parentOps[pattern->getRootKind()])
       patternWorklist.set_union(invalidPatterns[op]);
   }
 }
 
-void OperationLegalizer::computeLegalizationGraphBenefit() {
+void OperationLegalizer::computeLegalizationGraphBenefit(
+    DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
   // The smallest pattern depth, when legalizing an operation.
   DenseMap<OperationName, unsigned> minPatternDepth;
 
   // Compute the minimum legalization depth for a given operation.
   std::function<unsigned(OperationName)> computeDepth = [&](OperationName op) {
     // Check for existing depth.
     auto depthIt = minPatternDepth.find(op);
     if (depthIt != minPatternDepth.end())
       return depthIt->second;
 
     // If a mapping for this operation does not exist, then this operation
     // is always legal. Return 0 as the depth for a directly legal operation.
     auto opPatternsIt = legalizerPatterns.find(op);
     if (opPatternsIt == legalizerPatterns.end() || opPatternsIt->second.empty())
       return 0u;
 
     // Initialize the depth to the maximum value.
     unsigned minDepth = std::numeric_limits<unsigned>::max();
 
     // Record this initial depth in case we encounter this op again when
     // recursively computing the depth.
     minPatternDepth.try_emplace(op, minDepth);
 
     // Compute the depth for each pattern used to legalize this operation.
-    SmallVector<std::pair<RewritePattern *, unsigned>, 4> patternsByDepth;
+    SmallVector<std::pair<const RewritePattern *, unsigned>, 4> patternsByDepth;
     patternsByDepth.reserve(opPatternsIt->second.size());
-    for (RewritePattern *pattern : opPatternsIt->second) {
+    for (const RewritePattern *pattern : opPatternsIt->second) {
       unsigned depth = 0;
       for (auto generatedOp : pattern->getGeneratedOps())
         depth = std::max(depth, computeDepth(generatedOp) + 1);
       patternsByDepth.emplace_back(pattern, depth);
 
       // Update the min depth for this operation.
       minDepth = std::min(minDepth, depth);
     }
 
     // Update the pattern depth.
     minPatternDepth[op] = minDepth;
 
     // If the operation only has one legalization pattern, there is no need to
     // sort them.
     if (patternsByDepth.size() == 1)
       return minDepth;
 
     // Sort the patterns by those likely to be the most beneficial.
     llvm::array_pod_sort(
         patternsByDepth.begin(), patternsByDepth.end(),
-        [](const std::pair<RewritePattern *, unsigned> *lhs,
-           const std::pair<RewritePattern *, unsigned> *rhs) {
+        [](const std::pair<const RewritePattern *, unsigned> *lhs,
+           const std::pair<const RewritePattern *, unsigned> *rhs) {
           // First sort by the smaller pattern legalization depth.
           if (lhs->second != rhs->second)
             return llvm::array_pod_sort_comparator<unsigned>(&lhs->second,
                                                              &rhs->second);
 
           // Then sort by the larger pattern benefit.
           auto lhsBenefit = lhs->first->getBenefit();
           auto rhsBenefit = rhs->first->getBenefit();
           return llvm::array_pod_sort_comparator<PatternBenefit>(&rhsBenefit,
                                                                  &lhsBenefit);
         });
 
     // Update the legalization pattern to use the new sorted list.
     opPatternsIt->second.clear();
     for (auto &patternIt : patternsByDepth)
       opPatternsIt->second.push_back(patternIt.first);
 
     return minDepth;
   };
 
   // For each operation that is transitively legal, compute a cost for it.
   for (auto &opIt : legalizerPatterns)
     if (!minPatternDepth.count(opIt.first))
       computeDepth(opIt.first);
+
+  // Apply a cost model to the pattern applicator. We order patterns first by
+  // depth then benefit. `legalizerPatterns` contains per-op patterns by
+  // decreasing benefit.
+  applicator.applyCostModel([&](const RewritePattern &p) {
+    auto &list = legalizerPatterns[p.getRootKind()];
+
+    // If the pattern is not found, then it was removed and cannot be matched.
+    LegalizationPatterns::iterator it = llvm::find(list, &p);
+    if (it == list.end())
+      return PatternBenefit::impossibleToMatch();
+
+    // Patterns found earlier in the list have higher benefit.
+    return PatternBenefit(std::distance(it, list.end()));
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // OperationConverter
 //===----------------------------------------------------------------------===//
 namespace {
 enum OpConversionMode {
   // In this mode, the conversion will ignore failed conversions to allow

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 /// 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 OwningRewritePatternList &patterns)
       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
     worklist.reserve(64);
+
+    // Apply a simple cost model based solely on pattern benefit.
+    matcher.applyDefaultCostModel();
   }
 
   bool simplify(MutableArrayRef<Region> regions, int maxIterations);
 
   void addToWorklist(Operation *op) {
     // Check to see if the worklist already contains this op.
     if (worklistMap.count(op))
       return;
for (auto result : op->getResults())
         addToWorklist(user);
   }
 
 private:
   // 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
   // simplifications.
-  template <typename Operands>
-  void addToWorklist(Operands &&operands) {
+  template <typename Operands> void addToWorklist(Operands &&operands) {
     for (Value operand : operands) {
       // If the use count of this operand is now < 2, we re-add the defining
       // operation to the worklist.
       // TODO(riverriddle) This is based on the fact that zero use operations
       // may be deleted, and that single use values often have more
       // canonicalization opportunities.
       if (!operand.use_empty() && !operand.hasOneUse())
         continue;
       if (auto *defInst = operand.getDefiningOp())
         addToWorklist(defInst);
     }
   }
 
-  /// The low-level pattern matcher.
-  RewritePatternMatcher matcher;
+  /// The low-level pattern applicator.
+  PatternApplicator matcher;
 
   /// The worklist for this transformation keeps track of the operations that
   /// 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;
 
while (!worklist.empty()) {
       bool inPlaceUpdate;
       if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
                                       &inPlaceUpdate)))) {
         changed = true;
         if (!inPlaceUpdate)
           continue;
       }
 
-      // Make sure that any new operations are inserted at this point.
-      setInsertionPoint(op);
-
       // 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 |= matcher.matchAndRewrite(op, *this);
+      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 (succeeded(simplifyRegions(regions))) {
       folder.clear();
       changed = true;
     }
   } while (changed && ++i < 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 true if no more patterns can be matched in
-/// the result operation regions.
-/// Note: This does not apply patterns to the top-level operation itself.
+/// 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.
 ///
-bool mlir::applyPatternsAndFoldGreedily(
-    Operation *op, const OwningRewritePatternList &patterns) {
+LogicalResult
+mlir::applyPatternsAndFoldGreedily(Operation *op,
+                                   const OwningRewritePatternList &patterns) {
   return applyPatternsAndFoldGreedily(op->getRegions(), patterns);
 }
-
 /// Rewrite the given regions, which must be isolated from above.
-bool mlir::applyPatternsAndFoldGreedily(
-    MutableArrayRef<Region> regions, const OwningRewritePatternList &patterns) {
+LogicalResult
+mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
+                                   const OwningRewritePatternList &patterns) {
   if (regions.empty())
-    return true;
+    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()->isKnownIsolatedFromAbove();
   };
   (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);
   bool converged = driver.simplify(regions, maxPatternMatchIterations);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << maxPatternMatchIterations << " times";
   });
-  return converged;
+  return success(converged);
 }
 
 //===----------------------------------------------------------------------===//
 // OpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is a simple driver for the PatternMatcher to apply patterns and perform
 /// folding on a single op. It repeatedly applies locally optimal patterns.
 class OpPatternRewriteDriver : public PatternRewriter {
 public:
   explicit OpPatternRewriteDriver(MLIRContext *ctx,
                                   const OwningRewritePatternList &patterns)
-      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {}
+      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
+    // Apply a simple cost model based solely on pattern benefit.
+    matcher.applyDefaultCostModel();
+  }
 
-  bool simplifyLocally(Operation *op, int maxIterations, bool &erased);
+  /// Performs the rewrites and folding only on `op`. The simplification
+  /// converges if the op is erased as a result of being folded, replaced, or
+  /// dead, or no more changes happen in an iteration. Returns success if the
+  /// rewrite converges in `maxIterations`. `erased` is set to true if `op` gets
+  /// erased.
+  LogicalResult simplifyLocally(Operation *op, int maxIterations, bool &erased);
 
   // These are hooks implemented for PatternRewriter.
 protected:
   /// If an operation is about to be removed, mark it so that we can let clients
   /// know.
   void notifyOperationRemoved(Operation *op) override {
     opErasedViaPatternRewrites = true;
   }
 
   // When a root is going to be replaced, its removal will be notified as well.
   // So there is nothing to do here.
   void notifyRootReplaced(Operation *op) override {}
 
 private:
-  /// The low-level pattern matcher.
-  RewritePatternMatcher matcher;
+  /// The low-level pattern applicator.
+  PatternApplicator matcher;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
   /// Set to true if the operation has been erased via pattern rewrites.
   bool opErasedViaPatternRewrites = false;
 };
 
 } // anonymous namespace
 
-/// Performs the rewrites and folding only on `op`. The simplification converges
-/// if the op is erased as a result of being folded, replaced, or dead, or no
-/// more changes happen in an iteration. Returns true if the rewrite converges
-/// in `maxIterations`. `erased` is set to true if `op` gets erased.
-bool OpPatternRewriteDriver::simplifyLocally(Operation *op, int maxIterations,
+LogicalResult OpPatternRewriteDriver::simplifyLocally(Operation *op,
+                                                      int maxIterations,
                                                       bool &erased) {
   bool changed = false;
   erased = false;
   opErasedViaPatternRewrites = false;
   int i = 0;
   // Iterate until convergence or until maxIterations. Deletion of the op as
   // a result of being dead or folded is convergence.
   do {
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       op->erase();
       erased = true;
-      return true;
+      return success();
     }
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
                                    /*preReplaceAction=*/nullptr,
                                    &inPlaceUpdate))) {
       changed = true;
       if (!inPlaceUpdate) {
         erased = true;
-        return true;
+        return success();
       }
     }
 
-    // Make sure that any new operations are inserted at this point.
-    setInsertionPoint(op);
-
     // 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 |= matcher.matchAndRewrite(op, *this);
+    changed |= succeeded(matcher.matchAndRewrite(op, *this));
     if ((erased = opErasedViaPatternRewrites))
-      return true;
+      return success();
   } while (changed && ++i < maxIterations);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
-  return !changed;
+  return failure(changed);
 }
 
 /// 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.
-bool mlir::applyOpPatternsAndFold(Operation *op,
-                                  const OwningRewritePatternList &patterns,
-                                  bool *erased) {
+LogicalResult mlir::applyOpPatternsAndFold(
+    Operation *op, const OwningRewritePatternList &patterns, bool *erased) {
   // Start the pattern driver.
   OpPatternRewriteDriver driver(op->getContext(), patterns);
   bool opErased;
-  bool converged =
+  LogicalResult converged =
       driver.simplifyLocally(op, maxPatternMatchIterations, opErased);
   if (erased)
     *erased = opErased;
-  LLVM_DEBUG(if (!converged) {
+  LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << maxPatternMatchIterations << " times";
   });
   return converged;
 }