[mlir] Allow for patterns to match any root kind.

[rriddle]

Summary

Traditionally patterns have always had the root operation kind hardcoded to a specific operation name. This has worked well for quite some time, but it has certain limitations that make it undesirable. For example, some lowering have the same implementation for many different operations types with a few lowering entire dialects using the same pattern implementation. This problem has led to several "solutions":
a) Provide a template implementation to the user so that they can instantiate it for each operation combination, generally requiring the inclusion of the auto-generated operation definition file.
b) Use a non-templated pattern that allows for providing the name of the operation to match

• No one ever does this, because enumerating operation names can be cumbersome and so this quickly devolves into solution a.

This revision removes the restriction that patterns have a hardcoded root type, and allows for a class patterns that could match "any" operation type. The major downside of root-agnostic patterns is that they make certain pattern analyses more difficult, so it is still very highly encouraged that an operation specific pattern be used whenever possible.

Depends On D81985

[ftynse]

I wanted this for long time, thanks!

One minor suggestion: if we want to make people aware they create non-op-specific patterns, we can require to pass in a tag type (struct MatchAnyOperation{};) instead of the operation name in the Pattern constructor. Arguably, if it is explicit and doesn't save keystrokes, there will be less temptation to use "the constructor with shorter argument list" as the first choice.

mlir/include/mlir/IR/PatternMatch.h

455

Nit: I don't see ConstructorArg here

mlir/lib/Transforms/DialectConversion.cpp

1558

Ultra-nit: capitalize Legalizer in anyOplegalizerPatterns

[nicolasvasilache]

Thanks for adding this @rriddle .

mlir/include/mlir/IR/PatternMatch.h

524

Is the failure case expected to be more commonly specialized by the user which justifies onFailure is before onSuccess in the API ?

mlir/lib/IR/PatternMatch.cpp

253

typo: then

mlir/lib/Transforms/DialectConversion.cpp

1528

Seeing the extra complexity that appear when using this with legalization I am wondering whether it would be better to conservatively enforce a root in such cases?
Do you foresee that we are going to more generally use this in conjunction with A->B->C legalizations?
I can definitely see the benefit for individual/greedy pattern application but the A->B->C behavior looks iffy.

I can't claim I understand all the ramifications here, do you see this as a "nah it's fine don't worry" or as opening potential cans of worms ?

[rriddle]

Address comments

[rriddle]

In D82066#2100243, @ftynse wrote:

I wanted this for long time, thanks!

One minor suggestion: if we want to make people aware they create non-op-specific patterns, we can require to pass in a tag type (struct MatchAnyOperation{};) instead of the operation name in the Pattern constructor. Arguably, if it is explicit and doesn't save keystrokes, there will be less temptation to use "the constructor with shorter argument list" as the first choice.

Sounds great, added.

mlir/include/mlir/IR/PatternMatch.h

524

I couldn't say really. There are only two pattern drivers in core MLIR, and not any interesting ones internal to google. There are no use cases right now that would only specialize onSuccess, but not onFailure(or vice-versa). If many come up, we could swap but it is premature at this point given that there are relatively few drivers.

mlir/lib/Transforms/DialectConversion.cpp

1528

There are two parts of the A->B->C that are affected:

• This function, which mostly just does filtering of patterns that will "never" generate a legal result.

Negatively affecting this function isn't something I'm terribly concerned about ATM. Its only an optimization, and I'm of the opinion that we can do much better if we rewrote this and/or exposed at a higher level.

• The cost model

This one is a bit more concerning, but given the current set of usages that I've seen (expect) I'm not terribly concerned at the moment. I'm taking the same "wait and see" approach that I took when first writing the cost model. Even though the patterns *could* match any operation, in general I don't think these patterns will largely affect the cost model of others so I've kept it relatively simple. If this ever becomes a problem and we need a better cost analysis, we can rewrite and adapt. Until then, if everything runs relatively smoothly I don't see a need to over-engineer.

[schweitz]

Do you have any thoughts on matching roots based on their properties such as a supported interface? I could imagine patterns that match on any op that has memory effects, for example.

[rriddle]

In D82066#2103745, @schweitz wrote:

Do you have any thoughts on matching roots based on their properties such as a supported interface? I could imagine patterns that match on any op that has memory effects, for example.

That is what this patch allows. You can add whatever kind of matching you want.

mlir/include/mlir/IR/PatternMatch.h

 public:
   /// Return the benefit (the inverse of "cost") of matching this pattern.  The
   /// benefit of a Pattern is always static - rewrites that may have dynamic
   /// benefit can be instantiated multiple times (different Pattern instances)
   /// for each benefit that they may return, and be guarded by different match
   /// condition predicates.
   PatternBenefit getBenefit() const { return benefit; }
 
-  /// Return the root node that this pattern matches.  Patterns that can
-  /// match multiple root types are instantiated once per root.
-  OperationName getRootKind() const { return rootKind; }
+  /// Return the root node that this pattern matches. Patterns that can match
+  /// multiple root types return None.
+  Optional<OperationName> getRootKind() const { return rootKind; }
 
   //===--------------------------------------------------------------------===//
   // Implementation hooks for patterns to implement.
   //===--------------------------------------------------------------------===//
 
   /// Attempt to match against code rooted at the specified operation,
   /// which is the same operation code as getRootKind().
   virtual LogicalResult match(Operation *op) const = 0;
 
   virtual ~Pattern() {}
 
 protected:
-  /// Patterns must specify the root operation name they match against, and can
-  /// also specify the benefit of the pattern matching.
+  /// This constructor is used for patterns that match against a specific
+  /// operation type. The `benefit` is the expected benefit of matching this
+  /// pattern.
   Pattern(StringRef rootName, PatternBenefit benefit, MLIRContext *context);
 
+  /// This contructor is used when a pattern may match against multiple
+  /// different types of operations. The `benefit` is the expected benefit of
+  /// matching this pattern.
+  Pattern(PatternBenefit benefit);
+
 private:
-  const OperationName rootKind;
+  /// The root operation of the pattern. If the pattern matches a specific
+  /// operation, this contains the name of that operation. Contains None
+  /// otherwise.
+  Optional<OperationName> rootKind;
+
+  /// The expected benefit of matching this pattern.
   const PatternBenefit benefit;
 
   virtual void anchor();
 };
 
 /// RewritePattern is the common base class for all DAG to DAG replacements.
 /// There are two possible usages of this class:
 ///   * Multi-step RewritePattern with "match" and "rewrite"

 
 protected:
   /// Patterns must specify the root operation name they match against, and can
   /// also specify the benefit of the pattern matching.
   RewritePattern(StringRef rootName, PatternBenefit benefit,
                  MLIRContext *context)
       : Pattern(rootName, benefit, context) {}
   /// Patterns must specify the root operation name they match against, and can
+  /// also specify the benefit of the pattern matching.
+  RewritePattern(PatternBenefit benefit) : Pattern(benefit) {}
+  /// Patterns must specify the root operation name they match against, and can
   /// also specify the benefit of the pattern matching. They can also specify
   /// the names of operations that may be generated during a successful rewrite.
   RewritePattern(StringRef rootName, ArrayRef<StringRef> generatedNames,
                  PatternBenefit benefit, MLIRContext *context);
+  /// Patterns must specify the root operation name they match against, and can
+  /// also specify the benefit of the pattern matching. They can also specify
+  /// the names of operations that may be generated during a successful rewrite.
+  RewritePattern(ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
+                 MLIRContext *context);
 
   /// A list of the potential operations that may be generated when rewriting
   /// an op with this pattern.
   SmallVector<OperationName, 2> generatedOps;
 };
 
 /// OpRewritePattern is a wrapper around RewritePattern that allows for
 /// matching and rewriting against an instance of a derived operation class as
OwningRewritePatternList &insert(ConstructorArg &&arg,
     // 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'.
     (void)std::initializer_list<int>{
         0, (patterns.emplace_back(std::make_unique<Ts>(arg, args...)), 0)...};
     return *this;
   }
 
+  /// Add an instance of each of the pattern types 'Ts'. Return a reference to
+  /// `this` for chaining insertions.
+  /// Note: ConstructorArg is necessary here to separate the two variadic lists.
+  template <typename... Ts, typename = std::enable_if_t<sizeof...(Ts) != 0>>
+  OwningRewritePatternList &insert() {
+    (void)std::initializer_list<int>{
+        0, (patterns.emplace_back(std::make_unique<Ts>()), 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;
 };
void applyDefaultCostModel() {
     applyCostModel(
         [](const RewritePattern &pattern) { return pattern.getBenefit(); });
   }
 
   /// Walk all of the rewrite patterns within the applicator.
   void walkAllPatterns(function_ref<void(const RewritePattern &)> walk);
 
 private:
+  /// Attempt to match and rewrite the given op with the given 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.
+  LogicalResult matchAndRewrite(
+      Operation *op, const RewritePattern &pattern, PatternRewriter &rewriter,
+      function_ref<bool(const RewritePattern &)> canApply,
+      function_ref<void(const RewritePattern &)> onFailure,
+      function_ref<LogicalResult(const RewritePattern &)> onSuccess);
+
   /// The list that owns the patterns used within this applicator.
   const OwningRewritePatternList &owningPatternList;
 
   /// The set of patterns to match for each operation, stable sorted by benefit.
   DenseMap<OperationName, SmallVector<RewritePattern *, 2>> patterns;
+  /// The set of patterns that may match against any operation type, stable
+  /// sorted by benefit.
+  SmallVector<RewritePattern *, 1> anyOpPatterns;
 };
 
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy

mlir/lib/IR/PatternMatch.cpp

 
 //===----------------------------------------------------------------------===//
 // Pattern implementation
 //===----------------------------------------------------------------------===//
 
 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,
                  MLIRContext *context)
     : rootKind(OperationName(rootName, context)), benefit(benefit) {}
+Pattern::Pattern(PatternBenefit benefit) : benefit(benefit) {}
 
 // Out-of-line vtable anchor.
 void Pattern::anchor() {}
 
 //===----------------------------------------------------------------------===//
 // RewritePattern and PatternRewriter implementation
 //===----------------------------------------------------------------------===//
 

 
 //===----------------------------------------------------------------------===//
 // PatternMatcher implementation
 //===----------------------------------------------------------------------===//
 
 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());
+  anyOpPatterns.clear();
+  for (const auto &pat : owningPatternList) {
+    // If the pattern is always impossible to match, just ignore it.
+    if (pat->getBenefit().isImpossibleToMatch())
+      continue;
+    if (Optional<OperationName> opName = pat->getRootKind())
+      patterns[*opName].push_back(pat.get());
+    else
+      anyOpPatterns.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;
-
+  auto processPatternList = [&](SmallVectorImpl<RewritePattern *> &list) {
     // 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;
+      return;
     }
 
     // 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();
-  }
+  };
+  for (auto &it : patterns)
+    processPatternList(it.second);
+  processPatternList(anyOpPatterns);
 }
 
 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.
 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) {
+  // Check to see if there are patterns matching this specific operation type.
+  MutableArrayRef<RewritePattern *> opPatterns;
   auto patternIt = patterns.find(op->getName());
-  if (patternIt == patterns.end())
-    return failure();
+  if (patternIt != patterns.end())
+    opPatterns = patternIt->second;
+
+  // Process the patterns for that match the specific operation type, and any
+  // operation type in an interleaved fashion.
+  // FIXME: It'd be nice to just write an llvm::make_merge_range utility
+  // and pass in a comparison function. That would make this code trivial.
+  auto opIt = opPatterns.begin(), opE = opPatterns.end();
+  auto anyIt = anyOpPatterns.begin(), anyE = anyOpPatterns.end();
+  while (opIt != opE && anyIt != anyE) {
+    // Try to match the pattern providing the most benefit.
+    RewritePattern *pattern;
+    if ((*opIt)->getBenefit() >= (*anyIt)->getBenefit())
+      pattern = *(opIt++);
+    else
+      pattern = *(anyIt++);
+
+    // Otherwise, try to match the generic pattern.
+    if (succeeded(matchAndRewrite(op, *pattern, rewriter, canApply, onFailure,
+                                  onSuccess)))
+      return success();
+  }
+  // If we break from the loop, than only one of the ranges can still have
+  // elements. Loop over both without checking given that we don't need to
+  // interleave anymore.
+  for (RewritePattern *pattern : llvm::concat<RewritePattern *>(
+           llvm::make_range(opIt, opE), llvm::make_range(anyIt, anyE))) {
+    if (succeeded(matchAndRewrite(op, *pattern, rewriter, canApply, onFailure,
+                                  onSuccess)))
+      return success();
+  }
+}
 
-  for (auto *pattern : patternIt->second) {
-    // Check that the pattern can be applied.
-    if (canApply && !canApply(*pattern))
-      continue;
+LogicalResult PatternApplicator::matchAndRewrite(
+    Operation *op, const RewritePattern &pattern, PatternRewriter &rewriter,
+    function_ref<bool(const RewritePattern &)> canApply,
+    function_ref<void(const RewritePattern &)> onFailure,
+    function_ref<LogicalResult(const RewritePattern &)> onSuccess) {
+  if (canApply && !canApply(pattern))
+    return failure();
 
-    // Try to match and rewrite this pattern. The patterns are sorted by
-    // 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 (succeeded(pattern.matchAndRewrite(op, rewriter)))
+    return success(!onSuccess || succeeded(onSuccess(pattern)));
 
-    if (onFailure)
-      onFailure(*pattern);
-  }
+    onFailure(pattern);
   return failure();
 }

mlir/lib/Transforms/DialectConversion.cpp

private:
 
   /// 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.
+  /// function populates 'anyOpLegalizerPatterns' and 'legalizerPatterns' with
+  /// patterns for operations that are not directly legal, but may be
+  /// transitively legal for the current target given the provided patterns.
   void buildLegalizationGraph(
+      LegalizationPatterns &anyOpLegalizerPatterns,
       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(
+      LegalizationPatterns &anyOpLegalizerPatterns,
+      DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
+
+  /// Compute the legalization depth when legalizing an operation of the given
+  /// type.
+  unsigned computeOpLegalizationDepth(
+      OperationName op, DenseMap<OperationName, unsigned> &minOpPatternDepth,
+      DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
+
+  /// Apply the conversion cost model to the given set of patterns, and return
+  /// the smallest legalization depth of any of the patterns. See
+  /// `computeLegalizationGraphBenefit` for the breakdown of the cost model.
+  unsigned applyCostModelToPatterns(
+      LegalizationPatterns &patterns,
+      DenseMap<OperationName, unsigned> &minOpPatternDepth,
       DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns);
 
   /// The current set of patterns that have been applied.
   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.
+  // The set of patterns that can be applied to illegal operations to transform
+  // them into legal ones.
   DenseMap<OperationName, LegalizationPatterns> legalizerPatterns;
+  LegalizationPatterns anyOpLegalizerPatterns;
 
-  buildLegalizationGraph(legalizerPatterns);
-  computeLegalizationGraphBenefit(legalizerPatterns);
+  buildLegalizationGraph(anyOpLegalizerPatterns, legalizerPatterns);
+  computeLegalizationGraphBenefit(anyOpLegalizerPatterns, 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

 }
 
 bool OperationLegalizer::canApplyPattern(Operation *op,
                                          const RewritePattern &pattern,
                                          ConversionPatternRewriter &rewriter) {
   LLVM_DEBUG({
     auto &os = rewriter.getImpl().logger;
     os.getOStream() << "\n";
-    os.startLine() << "* Pattern : '" << pattern.getRootKind() << " -> (";
+    os.startLine() << "* Pattern : '" << op->getName() << " -> (";
     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() &&
for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
     }
   }
 
   LLVM_DEBUG(logSuccess(rewriterImpl.logger, "pattern applied successfully"));
   return success();
 }
 
 void OperationLegalizer::buildLegalizationGraph(
+    LegalizationPatterns &anyOplegalizerPatterns,
     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<const RewritePattern *, 2>>
       invalidPatterns;
   // A worklist of patterns to consider for legality.
   llvm::SetVector<const RewritePattern *> patternWorklist;
 
   // Build the mapping from operations to the parent ops that may generate them.
   applicator.walkAllPatterns([&](const RewritePattern &pattern) {
-    OperationName root = pattern.getRootKind();
+    Optional<OperationName> root = pattern.getRootKind();
+
+    // If the pattern has no specific root, we can't analyze the relationship
+    // between the root op and generated operations. Given that, add all such
+    // patterns to the legalization set.
+    if (!root) {
+      anyOplegalizerPatterns.push_back(&pattern);
+      return;
+    }
 
     // Skip operations that are always known to be legal.
-    if (target.getOpAction(root) == LegalizationAction::Legal)
+    if (target.getOpAction(*root) == LegalizationAction::Legal)
       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);
+    invalidPatterns[*root].insert(&pattern);
     for (auto op : pattern.getGeneratedOps())
-      parentOps[op].insert(root);
+      parentOps[op].insert(*root);
 
     // Add this pattern to the worklist.
     patternWorklist.insert(&pattern);
   });
 
+  // If there are any patterns that don't have a specific root kind, we can't
+  // make direct assumptions about what operations will never be legalized.
+  // Note: Technically, we could but it would require an analysis that may
+  // recurse into itself. It would be better to perform this kind of filtering
+  // at a higher level than here anyways.
+  if (!anyOplegalizerPatterns.empty()) {
+    for (const RewritePattern *pattern : patternWorklist)
+      legalizerPatterns[*pattern->getRootKind()].push_back(pattern);
+    return;
+  }
+
   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);
+    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()])
+    for (auto op : parentOps[*pattern->getRootKind()])
       patternWorklist.set_union(invalidPatterns[op]);
   }
 }
 
 void OperationLegalizer::computeLegalizationGraphBenefit(
+    LegalizationPatterns &anyOplegalizerPatterns,
     DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
   // The smallest pattern depth, when legalizing an operation.
-  DenseMap<OperationName, unsigned> minPatternDepth;
+  DenseMap<OperationName, unsigned> minOpPatternDepth;
 
-  // 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;
+  // For each operation that is transitively legal, compute a cost for it.
+  for (auto &opIt : legalizerPatterns)
+    if (!minOpPatternDepth.count(opIt.first))
+      computeOpLegalizationDepth(opIt.first, minOpPatternDepth,
+                                 legalizerPatterns);
+
+  // Apply the cost model to the patterns that can match any operation. Those
+  // with a specific operation type are already resolved when computing the op
+  // legalization depth.
+  if (!anyOplegalizerPatterns.empty())
+    applyCostModelToPatterns(anyOplegalizerPatterns, minOpPatternDepth,
+                             legalizerPatterns);
+
+  // 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) {
+    ArrayRef<const RewritePattern *> orderedPatternList;
+    if (Optional<OperationName> rootName = p.getRootKind())
+      orderedPatternList = legalizerPatterns[*rootName];
+    else
+      orderedPatternList = anyOplegalizerPatterns;
+
+    // If the pattern is not found, then it was removed and cannot be matched.
+    auto it = llvm::find(orderedPatternList, &p);
+    if (it == orderedPatternList.end())
+      return PatternBenefit::impossibleToMatch();
+
+    // Patterns found earlier in the list have higher benefit.
+    return PatternBenefit(std::distance(it, orderedPatternList.end()));
+  });
+}
+
+unsigned OperationLegalizer::computeOpLegalizationDepth(
+    OperationName op, DenseMap<OperationName, unsigned> &minOpPatternDepth,
+    DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
+  auto depthIt = minOpPatternDepth.find(op);
+  if (depthIt != minOpPatternDepth.end())
 
-    // 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);
+  minOpPatternDepth.try_emplace(op, std::numeric_limits<unsigned>::max());
 
-    // Compute the depth for each pattern used to legalize this operation.
-    SmallVector<std::pair<const RewritePattern *, unsigned>, 4> patternsByDepth;
-    patternsByDepth.reserve(opPatternsIt->second.size());
-    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);
+  // Apply the cost model to the operation patterns, and update the minimum
+  // depth.
+  unsigned minDepth = applyCostModelToPatterns(
+      opPatternsIt->second, minOpPatternDepth, legalizerPatterns);
+  minOpPatternDepth[op] = minDepth;
+  return minDepth;
+}
+
+unsigned OperationLegalizer::applyCostModelToPatterns(
+    LegalizationPatterns &patterns,
+    DenseMap<OperationName, unsigned> &minOpPatternDepth,
+    DenseMap<OperationName, LegalizationPatterns> &legalizerPatterns) {
+  unsigned minDepth = std::numeric_limits<unsigned>::max();
+
+  // Compute the depth for each pattern within the set.
+  patternsByDepth.reserve(patterns.size());
+  for (const RewritePattern *pattern : patterns) {
+    for (auto generatedOp : pattern->getGeneratedOps()) {
+      unsigned generatedOpDepth = computeOpLegalizationDepth(
+          generatedOp, minOpPatternDepth, legalizerPatterns);
+      depth = std::max(depth, generatedOpDepth + 1);
+    }
 
-      // Update the min depth for this operation.
-      minDepth = std::min(minDepth, depth);
-    }
+    // Update the minimum depth of the pattern list.
 
-    // 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<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()));
-  });
+  patterns.clear();
+    patterns.push_back(patternIt.first);
 }
 
 //===----------------------------------------------------------------------===//
 // OperationConverter
 //===----------------------------------------------------------------------===//
 namespace {
 enum OpConversionMode {
   // In this mode, the conversion will ignore failed conversions to allow

mlir/test/Transforms/test-legalize-unknown-root.mlir

+// RUN: mlir-opt %s -test-legalize-unknown-root-patterns | FileCheck %s
+
+// Test that all `test` dialect operations are removed.
+// CHECK-LABEL: func @remove_all_ops
+func @remove_all_ops(%arg0: i32) {
+  // CHECK-NEXT: return
+  %0 = "test.illegal_op_a"() : () -> i32
+  %1 = "test.illegal_op_b"() : () -> i32
+  %2 = "test.illegal_op_c"() : () -> i32
+  %3 = "test.illegal_op_d"() : () -> i32
+  %4 = "test.illegal_op_e"() : () -> i32
+  return
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

void runOnFunction() override {
 
     if (failed(mlir::applyFullConversion(getFunction(), target, patterns))) {
       signalPassFailure();
     }
   }
 };
 } // end anonymous namespace
 
+//===----------------------------------------------------------------------===//
+// Test patterns without a specific root operation kind
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This pattern matches and removes any operation in the test dialect.
+struct RemoveTestDialectOps : public RewritePattern {
+  RemoveTestDialectOps() : RewritePattern(/*benefit=*/1) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (!isa<TestDialect>(op->getDialect()))
+      return failure();
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct TestUnknownRootOpDriver
+    : public mlir::PassWrapper<TestUnknownRootOpDriver, FunctionPass> {
+  void runOnFunction() override {
+    mlir::OwningRewritePatternList patterns;
+    patterns.insert<RemoveTestDialectOps>();
+
+    mlir::ConversionTarget target(getContext());
+    target.addIllegalDialect<TestDialect>();
+    if (failed(applyPartialConversion(getFunction(), target, patterns)))
+      signalPassFailure();
+  }
+};
+} // end anonymous namespace
+
 namespace mlir {
 void registerPatternsTestPass() {
-  mlir::PassRegistration<TestReturnTypeDriver>("test-return-type",
-                                               "Run return type functions");
+  PassRegistration<TestReturnTypeDriver>("test-return-type",
 
-  mlir::PassRegistration<TestDerivedAttributeDriver>(
-      "test-derived-attr", "Run test derived attributes");
+  PassRegistration<TestDerivedAttributeDriver>("test-derived-attr",
+                                               "Run test derived attributes");
 
-  mlir::PassRegistration<TestPatternDriver>("test-patterns",
-                                            "Run test dialect patterns");
+  PassRegistration<TestPatternDriver>("test-patterns",
 
-  mlir::PassRegistration<TestLegalizePatternDriver>(
+  PassRegistration<TestLegalizePatternDriver>(
       "test-legalize-patterns", "Run test dialect legalization patterns", [] {
         return std::make_unique<TestLegalizePatternDriver>(
             legalizerConversionMode);
       });
 
   PassRegistration<TestRemappedValue>(
       "test-remapped-value",
       "Test public remapped value mechanism in ConversionPatternRewriter");
+
+  PassRegistration<TestUnknownRootOpDriver>(
+      "test-legalize-unknown-root-patterns",
+      "Test public remapped value mechanism in ConversionPatternRewriter");
 }
 } // namespace mlir