[mlir][Pattern] Add better support for using interfaces/traits to match root operations in rewrite patterns

[rriddle]

Summary

To match an interface or trait, users currently have to use the MatchAny tag. This tag can be quite problematic for compile time for things like the canonicalizer, as the MatchAny patterns may get applied to *every* operation. This revision adds better support by bucketing interface/trait patterns based on which registered operations have them registered. This means that moving forward we will only attempt to match these patterns to operations that have this interface registered. Two simplify defining patterns that match traits and interfaces, two new utility classes have been added: OpTraitRewritePattern and OpInterfaceRewritePattern.

[mehdi_amini]

mlir/include/mlir/IR/PatternMatch.h

193

Wouldn't RootKind fit in 2-bits? Could this be the low-bits of the rootValue?

[ftynse]

Nice!

Please fix the linter messages and consider folding RootKind into some pointer.

mlir/include/mlir/IR/PatternMatch.h

73

Nit: why uint16_t specifically?

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

193

We can't use a PointerIntPair here because OperationName already uses a bit (internally it is a PointerUnion of AbstractionOperation/Identifier).

[mehdi_amini]

mlir/include/mlir/IR/PatternMatch.h

193

Ah! Maybe worth a comment?

Also would it fit within contextAndHasBoundedRecursion? (I don't know if there are 2 bits left there).

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

193

I considered fitting it into the context, but given that there is already padding for the benefit(it is only 2 bytes) I just slotted it in next to that.

[rriddle]

update

[rriddle]

rebase

[jpienaar]

Nice, any performance tests you did?

[jpienaar]

Nice, thanks

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

667

Does the filter not have a context?

[rriddle]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

667

Not that I could find.

flang/lib/Optimizer/Dialect/FIROps.cpp

 static bool isZero(mlir::Value v) { return checkIsIntegerConstant(v, 0); }
 static bool isOne(mlir::Value v) { return checkIsIntegerConstant(v, 1); }
 
 // Undo some complex patterns created in the front-end and turn them back into
 // complex ops.
 template <typename FltOp, typename CpxOp>
 struct UndoComplexPattern : public mlir::RewritePattern {
   UndoComplexPattern(mlir::MLIRContext *ctx)
-      : mlir::RewritePattern("fir.insert_value", {}, 2, ctx) {}
+      : mlir::RewritePattern("fir.insert_value", 2, ctx) {}
 
   mlir::LogicalResult
   matchAndRewrite(mlir::Operation *op,
                   mlir::PatternRewriter &rewriter) const override {
     auto insval = dyn_cast_or_null<fir::InsertValueOp>(op);
     if (!insval || !insval.getType().isa<fir::ComplexType>())
       return mlir::failure();
     auto insval2 =

mlir/include/mlir/Conversion/LinalgToStandard/LinalgToStandard.h

 // These patterns are exposed individually because they are expected to be
 // typically used individually.
 
 // Create a new call to the type-canonicalized `LinalgOp::getLibraryCallName()`
 // function. The implementation of the function can be either in the same module
 // or in an externally linked library.
 // This is a generic entry point for all LinalgOp, except for CopyOp and
 // IndexedGenericOp, for which omre specialized patterns are provided.
-class LinalgOpToLibraryCallRewrite : public RewritePattern {
+class LinalgOpToLibraryCallRewrite
+    : public OpInterfaceRewritePattern<LinalgOp> {
 public:
-  LinalgOpToLibraryCallRewrite()
-      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override;
 };
 
 /// Rewrite pattern specialization for CopyOp, kicks in when both input and
 /// output permutations are left unspecified or are the identity.
 class CopyOpToLibraryCallRewrite : public OpRewritePattern<CopyOp> {
 public:
   using OpRewritePattern<CopyOp>::OpRewritePattern;

mlir/include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h

 }
 
 template <typename PatternType, typename OpType, typename OptionsType>
 void enqueue(RewritePatternSet &patternList, OptionsType options,
              StringRef opName, linalg::LinalgTransformationFilter m) {
   if (!opName.empty())
     patternList.add<PatternType>(opName, patternList.getContext(), options, m);
   else
-    patternList.add<PatternType>(m.addOpFilter<OpType>(), options);
+    patternList.add<PatternType>(patternList.getContext(),
+                                 m.addOpFilter<OpType>(), options);
 }
 
 /// Promotion transformation enqueues a particular stage-1 pattern for
 /// `Tile<LinalgOpType>`with the appropriate `options`.
 template <typename LinalgOpType>
 struct Tile : public Transformation {
   explicit Tile(linalg::LinalgTilingOptions options,
                 linalg::LinalgTransformationFilter::FilterFunction f = nullptr)

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// Base pattern that applied the tiling transformation specified by `options`.
 /// Abort and return failure in 2 cases:
 ///   1. if the tiling specification is invalid and tiling fails to occur.
 ///   2. if tiling occurs but `options.paddingValueComputationFunction` is set
 ///      and some operand shape cannot be bounded statically.
 struct LinalgBaseTilingPattern : public RewritePattern {
   // Entry point to match any LinalgOp OpInterface.
   LinalgBaseTilingPattern(
-      LinalgTilingOptions options,
+      MLIRContext *context, LinalgTilingOptions options,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
   // Entry point to match a specific Linalg op.
   LinalgBaseTilingPattern(
       StringRef opName, MLIRContext *context, LinalgTilingOptions options,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
   LogicalResult matchAndRewriteBase(Operation *op, PatternRewriter &rewriter,

 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `vectorizeLinalgOp` for more details.
 
 /// Empty for now, used for SFINAE purposes only.
 struct LinalgVectorizationOptions {};
 
 struct LinalgBaseVectorizationPattern : public RewritePattern {
   /// MatchAnyOpTag-based constructor with a mandatory `filter`.
-  LinalgBaseVectorizationPattern(LinalgTransformationFilter filter,
+  LinalgBaseVectorizationPattern(MLIRContext *context,
+                                 LinalgTransformationFilter filter,
                                  PatternBenefit benefit = 1);
   /// Name-based constructor with an optional `filter`.
   LinalgBaseVectorizationPattern(
       StringRef opName, MLIRContext *context,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
 };
 
 struct LinalgVectorizationPattern : public LinalgBaseVectorizationPattern {
   /// These constructors are available to anyone.
   /// MatchAnyOpTag-based constructor with a mandatory `filter`.
   LinalgVectorizationPattern(
-      LinalgTransformationFilter filter,
+      MLIRContext *context, LinalgTransformationFilter filter,
       LinalgVectorizationOptions options = LinalgVectorizationOptions(),
       PatternBenefit benefit = 1)
-      : LinalgBaseVectorizationPattern(filter, benefit) {}
+      : LinalgBaseVectorizationPattern(context, filter, benefit) {}
   /// Name-based constructor with an optional `filter`.
   LinalgVectorizationPattern(
       StringRef opName, MLIRContext *context,
       LinalgVectorizationOptions options = LinalgVectorizationOptions(),
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
       : LinalgBaseVectorizationPattern(opName, context, filter, benefit) {}
 };

 
 /// SFINAE helper for single C++ class without a `getOperationName` method (e.g.
 /// an OpInterface).
 template <typename OpType, typename = std::enable_if_t<
                                !detect_has_get_operation_name<OpType>::value>>
 void insertVectorizationPatternImpl(RewritePatternSet &patternList,
                                     linalg::LinalgVectorizationOptions options,
                                     linalg::LinalgTransformationFilter f) {
-  patternList.add<linalg::LinalgVectorizationPattern>(f.addOpFilter<OpType>(),
-                                                      options);
+  patternList.add<linalg::LinalgVectorizationPattern>(
+      patternList.getContext(), f.addOpFilter<OpType>(), options);
 }
 
 /// Variadic helper function to insert vectorization patterns for C++ ops.
 template <typename... OpTypes>
 void insertVectorizationPatterns(RewritePatternSet &patternList,
                                  linalg::LinalgVectorizationOptions options,
                                  linalg::LinalgTransformationFilter f =
                                      linalg::LinalgTransformationFilter()) {

 };
 
 template <typename OpTy>
 struct LinalgLoweringPattern : public RewritePattern {
   LinalgLoweringPattern(
       MLIRContext *context, LinalgLoweringType loweringType,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       ArrayRef<unsigned> interchangeVector = {}, PatternBenefit benefit = 1)
-      : RewritePattern(OpTy::getOperationName(), {}, benefit, context),
+      : RewritePattern(OpTy::getOperationName(), benefit, context),
         filter(filter), loweringType(loweringType),
         interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
   // TODO: Move implementation to .cpp once named ops are auto-generated.
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

UnrollVectorOptions &setNativeShape(ArrayRef<int64_t> shape) {
     return *this;
   }
 };
 /// Pattern to apply `unrollSingleResultVectorOp` to a `targetShape`
 /// declaratively.
 struct UnrollVectorPattern : public RewritePattern {
   using FilterConstraintType = std::function<LogicalResult(Operation *op)>;
   UnrollVectorPattern(MLIRContext *context, UnrollVectorOptions options)
-      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()), options(options) {}
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
+        options(options) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (options.filterConstraint && failed(options.filterConstraint(op)))
       return failure();
     if (!options.nativeShape) {
       return op->emitError("vector unrolling expects the native shape or native"
                            "shape call back function to be set");
     }
using FilterConstraintType =
       std::function<LogicalResult(VectorTransferOpInterface op)>;
 
   explicit VectorTransferFullPartialRewriter(
       MLIRContext *context,
       VectorTransformsOptions options = VectorTransformsOptions(),
       FilterConstraintType filter =
           [](VectorTransferOpInterface op) { return success(); },
       PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()), options(options),
+      : RewritePattern(MatchAnyOpTypeTag(), benefit, context), options(options),
         filter(filter) {}
 
   /// Performs the rewrite.
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   VectorTransformsOptions options;

mlir/include/mlir/IR/OpDefinition.h

 #ifndef NDEBUG
     if (op->getName().getStringRef() == ConcreteType::getOperationName())
       llvm::report_fatal_error(
           "classof on '" + ConcreteType::getOperationName() +
           "' failed due to the operation not being registered");
 #endif
     return false;
   }
+  /// Provide `classof` support for other OpBase derived classes, such as
+  /// Interfaces.
+  template <typename T>
+  static std::enable_if_t<std::is_base_of<OpState, T>::value, bool>
+  classof(const T *op) {
+    return classof(const_cast<T *>(op)->getOperation());
+  }
 
   /// Expose the type we are instantiated on to template machinery that may want
   /// to introspect traits on this operation.
   using ConcreteOpType = ConcreteType;
 
   /// This is a public constructor.  Any op can be initialized to null.
   explicit Op() : OpState(nullptr) {}
   Op(std::nullptr_t) : OpState(nullptr) {}

mlir/include/mlir/IR/OperationSupport.h

public:
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this operation, null otherwise. This should not be used
   /// directly.
   template <typename T>
   typename T::Concept *getInterface() const {
     return interfaceMap.lookup<T>();
   }
 
+  /// Returns true if this operation has the given interface registered to it.
+  bool hasInterface(TypeID interfaceID) const {
+    return interfaceMap.contains(interfaceID);
+  }
+
   /// Returns true if the operation has a particular trait.
   template <template <typename T> class Trait>
   bool hasTrait() const {
     return hasTraitFn(TypeID::get<Trait>());
   }
 
+  /// Returns true if the operation has a particular trait.
+  bool hasTrait(TypeID traitID) const { return hasTraitFn(traitID); }
+
   /// Look up the specified operation in the specified MLIRContext and return a
   /// pointer to it if present.  Otherwise, return a null pointer.
   static const AbstractOperation *lookup(StringRef opName,
                                          MLIRContext *context);
 
   /// This constructor is used by Dialect objects when they register the list of
   /// operations they contain.
   template <typename T>

mlir/include/mlir/IR/PatternMatch.h

 // Pattern
 //===----------------------------------------------------------------------===//
 
 /// This class contains all of the data related to a pattern, but does not
 /// contain any methods or logic for the actual matching. This class is solely
 /// used to interface with the metadata of a pattern, such as the benefit or
 /// root operation.
 class Pattern {
+  /// This enum represents the kind of value used to select the root operations
+  /// that match this pattern.
+  enum class RootKind {
+    /// The pattern root matches "any" operation.
+    Any,
+    /// The pattern root is matched using a concrete operation name.
+    OperationName,
+    /// The pattern root is matched using an interface ID.
+    InterfaceID,
+    /// The patter root is matched using a trait ID.
+    TraitID
+  };
+
 public:
   /// Return a list of operations that may be generated when rewriting an
   /// operation instance with this pattern.
   ArrayRef<OperationName> getGeneratedOps() const { return generatedOps; }
 
   /// Return the root node that this pattern matches. Patterns that can match
   /// multiple root types return None.
-  Optional<OperationName> getRootKind() const { return rootKind; }
+  Optional<OperationName> getRootKind() const {
+    if (rootKind == RootKind::OperationName)
+      return OperationName::getFromOpaquePointer(rootValue);
+    return llvm::None;
+  }
+
+  /// Return the interface ID used to match the root operation of this pattern.
+  /// If the pattern does not use an interface ID for deciding the root match,
+  /// this returns None.
+  Optional<TypeID> getRootInterfaceID() const {
+    if (rootKind == RootKind::InterfaceID)
+      return TypeID::getFromOpaquePointer(rootValue);
+    return llvm::None;
+  }
+
+  /// Return the trait ID used to match the root operation of this pattern.
+  /// If the pattern does not use a trait ID for deciding the root match, this
+  /// returns None.
+  Optional<TypeID> getRootTraitID() const {
+    if (rootKind == RootKind::TraitID)
+      return TypeID::getFromOpaquePointer(rootValue);
+    return llvm::None;
+  }
 
   /// 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; }
 
   /// Returns true if this pattern is known to result in recursive application,
   /// i.e. this pattern may generate IR that also matches this pattern, but is
   /// known to bound the recursion. This signals to a rewrite driver that it is
   /// safe to apply this pattern recursively to generated IR.
-  bool hasBoundedRewriteRecursion() const { return hasBoundedRecursion; }
+  bool hasBoundedRewriteRecursion() const {
+    return contextAndHasBoundedRecursion.getInt();
+  }
+
+  /// Return the MLIRContext used to create this pattern.
+  MLIRContext *getContext() const {
+    return contextAndHasBoundedRecursion.getPointer();
+  }
 
 protected:
   /// This class acts as a special tag that makes the desire to match "any"
   /// operation type explicit. This helps to avoid unnecessary usages of this
   /// feature, and ensures that the user is making a conscious decision.
   struct MatchAnyOpTypeTag {};
+  /// This class acts as a special tag that makes the desire to match any
+  /// operation that implements a given interface explicit. This helps to avoid
+  /// unnecessary usages of this feature, and ensures that the user is making a
+  /// conscious decision.
+  struct MatchInterfaceOpTypeTag {};
+  /// This class acts as a special tag that makes the desire to match any
+  /// operation that implements a given trait explicit. This helps to avoid
+  /// unnecessary usages of this feature, and ensures that the user is making a
+  /// conscious decision.
+  struct MatchTraitOpTypeTag {};
 
   /// Construct a pattern with a certain benefit that matches the operation
   /// with the given root name.
-  Pattern(StringRef rootName, PatternBenefit benefit, MLIRContext *context);
-  /// Construct a pattern with a certain benefit that matches any operation
-  /// type. `MatchAnyOpTypeTag` is just a tag to ensure that the "match any"
-  /// behavior is what the user actually desired, `MatchAnyOpTypeTag()` should
-  /// always be supplied here.
-  Pattern(PatternBenefit benefit, MatchAnyOpTypeTag tag);
-  /// Construct a pattern with a certain benefit that matches the operation with
-  /// the given root name. `generatedNames` contains the names of operations
-  /// that may be generated during a successful rewrite.
-  Pattern(StringRef rootName, ArrayRef<StringRef> generatedNames,
-          PatternBenefit benefit, MLIRContext *context);
+  Pattern(StringRef rootName, PatternBenefit benefit, MLIRContext *context,
+          ArrayRef<StringRef> generatedNames = {});
   /// Construct a pattern that may match any operation type. `generatedNames`
   /// contains the names of operations that may be generated during a successful
   /// rewrite. `MatchAnyOpTypeTag` is just a tag to ensure that the "match any"
   /// behavior is what the user actually desired, `MatchAnyOpTypeTag()` should
   /// always be supplied here.
-  Pattern(ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
-          MLIRContext *context, MatchAnyOpTypeTag tag);
+  Pattern(MatchAnyOpTypeTag tag, PatternBenefit benefit, MLIRContext *context,
+          ArrayRef<StringRef> generatedNames = {});
+  /// Construct a pattern that may match any operation that implements the
+  /// interface defined by the provided `interfaceID`. `generatedNames` contains
+  /// the names of operations that may be generated during a successful rewrite.
+  /// `MatchInterfaceOpTypeTag` is just a tag to ensure that the "match
+  /// interface" behavior is what the user actually desired,
+  /// `MatchInterfaceOpTypeTag()` should always be supplied here.
+  Pattern(MatchInterfaceOpTypeTag tag, TypeID interfaceID,
+          PatternBenefit benefit, MLIRContext *context,
+          ArrayRef<StringRef> generatedNames = {});
+  /// Construct a pattern that may match any operation that implements the
+  /// trait defined by the provided `traitID`. `generatedNames` contains the
+  /// names of operations that may be generated during a successful rewrite.
+  /// `MatchTraitOpTypeTag` is just a tag to ensure that the "match trait"
+  /// behavior is what the user actually desired, `MatchTraitOpTypeTag()` should
+  /// always be supplied here.
+  Pattern(MatchTraitOpTypeTag tag, TypeID traitID, PatternBenefit benefit,
+          MLIRContext *context, ArrayRef<StringRef> generatedNames = {});
 
   /// Set the flag detailing if this pattern has bounded rewrite recursion or
   /// not.
   void setHasBoundedRewriteRecursion(bool hasBoundedRecursionArg = true) {
-    hasBoundedRecursion = hasBoundedRecursionArg;
+    contextAndHasBoundedRecursion.setInt(hasBoundedRecursionArg);
   }
 
 private:
-  /// A list of the potential operations that may be generated when rewriting
-  /// an op with this pattern.
-  SmallVector<OperationName, 2> generatedOps;
-
-  /// 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;
+  Pattern(const void *rootValue, RootKind rootKind,
+          ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
+          MLIRContext *context);
+
+  /// The value used to match the root operation of the pattern.
+  const void *rootValue;
+  RootKind rootKind;
 
   /// The expected benefit of matching this pattern.
   const PatternBenefit benefit;
 
-  /// A boolean flag of whether this pattern has bounded recursion or not.
-  bool hasBoundedRecursion = false;
+  /// The context this pattern was created from, and a boolean flag indicating
+  /// whether this pattern has bounded recursion or not.
+  llvm::PointerIntPair<MLIRContext *, 1, bool> contextAndHasBoundedRecursion;
+
+  /// A list of the potential operations that may be generated when rewriting
+  /// an op with this pattern.
+  SmallVector<OperationName, 2> generatedOps;
 };
 
 //===----------------------------------------------------------------------===//
 // RewritePattern
 //===----------------------------------------------------------------------===//
 
 /// RewritePattern is the common base class for all DAG to DAG replacements.
 /// There are two possible usages of this class:

 protected:
   /// Inherit the base constructors from `Pattern`.
   using Pattern::Pattern;
 
   /// An anchor for the virtual table.
   virtual void anchor();
 };
 
-/// OpRewritePattern is a wrapper around RewritePattern that allows for
-/// matching and rewriting against an instance of a derived operation class as
-/// opposed to a raw Operation.
+namespace detail {
+/// OpOrInterfaceRewritePatternBase is a wrapper around RewritePattern that
+/// allows for matching and rewriting against an instance of a derived operation
+/// class or Interface.
 template <typename SourceOp>
-struct OpRewritePattern : public RewritePattern {
-  /// Patterns must specify the root operation name they match against, and can
-  /// also specify the benefit of the pattern matching.
-  OpRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
-      : RewritePattern(SourceOp::getOperationName(), benefit, context) {}
+struct OpOrInterfaceRewritePatternBase : public RewritePattern {
+  using RewritePattern::RewritePattern;
 
   /// Wrappers around the RewritePattern methods that pass the derived op type.
   void rewrite(Operation *op, PatternRewriter &rewriter) const final {
     rewrite(cast<SourceOp>(op), rewriter);
   }
   LogicalResult match(Operation *op) const final {
     return match(cast<SourceOp>(op));
   }
virtual LogicalResult matchAndRewrite(SourceOp op,
                                         PatternRewriter &rewriter) const {
     if (succeeded(match(op))) {
       rewrite(op, rewriter);
       return success();
     }
     return failure();
   }
 };
+} // namespace detail
+
+/// OpRewritePattern is a wrapper around RewritePattern that allows for
+/// matching and rewriting against an instance of a derived operation class as
+/// opposed to a raw Operation.
+template <typename SourceOp>
+struct OpRewritePattern
+    : public detail::OpOrInterfaceRewritePatternBase<SourceOp> {
+  /// Patterns must specify the root operation name they match against, and can
+  /// also specify the benefit of the pattern matching.
+  OpRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceRewritePatternBase<SourceOp>(
+            SourceOp::getOperationName(), benefit, context) {}
+};
+
+/// OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for
+/// matching and rewriting against an instance of an operation interface instead
+/// of a raw Operation.
+template <typename SourceOp>
+struct OpInterfaceRewritePattern
+    : public detail::OpOrInterfaceRewritePatternBase<SourceOp> {
+  OpInterfaceRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceRewritePatternBase<SourceOp>(
+            Pattern::MatchInterfaceOpTypeTag(), SourceOp::getInterfaceID(),
+            benefit, context) {}
+};
+
+/// OpTraitRewritePattern is a wrapper around RewritePattern that allows for
+/// matching and rewriting against instances of an operation that possess a
+/// given trait.
+template <template <typename> class TraitType>
+class OpTraitRewritePattern : public RewritePattern {
+public:
+  OpTraitRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : RewritePattern(Pattern::MatchTraitOpTypeTag(), TypeID::get<TraitType>(),
+                       benefit, context) {}
+};
 
 //===----------------------------------------------------------------------===//
 // PDLPatternModule
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // PDLValue
 

mlir/include/mlir/Rewrite/FrozenRewritePatternSet.h

 /// pattern applicator. This class is designed to enable caching pattern lists
 /// such that they need not be continuously recomputed. Note that all copies of
 /// this class share the same compiled pattern list, allowing for a reduction in
 /// the number of duplicated patterns that need to be created.
 class FrozenRewritePatternSet {
   using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
 
 public:
+  /// A map of operation specific native patterns.
+  using OpSpecificNativePatternListT =
+      DenseMap<OperationName, std::vector<RewritePattern *>>;
+
   /// Freeze the patterns held in `patterns`, and take ownership.
   FrozenRewritePatternSet();
   FrozenRewritePatternSet(RewritePatternSet &&patterns);
   FrozenRewritePatternSet(FrozenRewritePatternSet &&patterns) = default;
   FrozenRewritePatternSet(const FrozenRewritePatternSet &patterns) = default;
   FrozenRewritePatternSet &
   operator=(const FrozenRewritePatternSet &patterns) = default;
   FrozenRewritePatternSet &
   operator=(FrozenRewritePatternSet &&patterns) = default;
   ~FrozenRewritePatternSet();
 
-  /// Return the native patterns held by this list.
+  /// Return the op specific native patterns held by this list.
+  const OpSpecificNativePatternListT &getOpSpecificNativePatterns() const {
+    return impl->nativeOpSpecificPatternMap;
+  }
+
+  /// Return the "match any" native patterns held by this list.
   iterator_range<llvm::pointee_iterator<NativePatternListT::const_iterator>>
-  getNativePatterns() const {
-    const NativePatternListT &nativePatterns = impl->nativePatterns;
-    return llvm::make_pointee_range(nativePatterns);
+  getMatchAnyOpNativePatterns() const {
+    const NativePatternListT &nativeList = impl->nativeAnyOpPatterns;
+    return llvm::make_pointee_range(nativeList);
   }
 
   /// Return the compiled PDL bytecode held by this list. Returns null if
   /// there are no PDL patterns within the list.
   const detail::PDLByteCode *getPDLByteCode() const {
     return impl->pdlByteCode.get();
   }
 
 private:
   /// The internal implementation of the frozen pattern list.
   struct Impl {
-    /// The set of native C++ rewrite patterns.
-    NativePatternListT nativePatterns;
+    /// The set of native C++ rewrite patterns that are matched to specific
+    /// operation kinds.
+    OpSpecificNativePatternListT nativeOpSpecificPatternMap;
+
+    /// The full op-specific native rewrite list. This allows for the map above
+    /// to contain duplicate patterns, e.g. for interfaces and traits.
+    NativePatternListT nativeOpSpecificPatternList;
+
+    /// The set of native C++ rewrite patterns that are matched to "any"
+    /// operation.
+    NativePatternListT nativeAnyOpPatterns;
 
     /// The bytecode containing the compiled PDL patterns.
     std::unique_ptr<detail::PDLByteCode> pdlByteCode;
   };
 
   /// A pointer to the internal pattern list. This uses a shared_ptr to avoid
   /// the need to compile the same pattern list multiple times. For example,
   /// during multi-threaded pass execution, all copies of a pass can share the

mlir/include/mlir/Support/InterfaceSupport.h

public:
   }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this map, null otherwise.
   template <typename T> typename T::Concept *lookup() const {
     return reinterpret_cast<typename T::Concept *>(lookup(T::getInterfaceID()));
   }
 
+  /// Returns true if the interface map contains an interface for the given id.
+  bool contains(TypeID interfaceID) const { return lookup(interfaceID); }
+
 private:
   /// Compare two TypeID instances by comparing the underlying pointer.
   static bool compare(TypeID lhs, TypeID rhs) {
     return lhs.getAsOpaquePointer() < rhs.getAsOpaquePointer();
   }
 
   InterfaceMap() = default;
   InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)

mlir/include/mlir/Transforms/DialectConversion.h

public:
   getTypeConverter() const {
     return static_cast<ConverterTy *>(typeConverter);
   }
 
 protected:
   /// See `RewritePattern::RewritePattern` for information on the other
   /// available constructors.
   using RewritePattern::RewritePattern;
-  /// Construct a conversion pattern that matches an operation with the given
-  /// root name. This constructor allows for providing a type converter to use
-  /// within the pattern.
-  ConversionPattern(StringRef rootName, PatternBenefit benefit,
-                    TypeConverter &typeConverter, MLIRContext *ctx)
-      : RewritePattern(rootName, benefit, ctx), typeConverter(&typeConverter) {}
-  /// Construct a conversion pattern that matches any operation type. This
-  /// constructor allows for providing a type converter to use within the
-  /// pattern. `MatchAnyOpTypeTag` is just a tag to ensure that the "match any"
-  /// behavior is what the user actually desired, `MatchAnyOpTypeTag()` should
-  /// always be supplied here.
-  ConversionPattern(PatternBenefit benefit, TypeConverter &typeConverter,
-                    MatchAnyOpTypeTag tag)
-      : RewritePattern(benefit, tag), typeConverter(&typeConverter) {}
+  /// Construct a conversion pattern with the given converter, and forward the
+  /// remaining arguments to RewritePattern.
+  template <typename... Args>
+  ConversionPattern(TypeConverter &typeConverter, Args &&... args)
+      : RewritePattern(std::forward<Args>(args)...),
+        typeConverter(&typeConverter) {}
 
 protected:
   /// An optional type converter for use by this pattern.
   TypeConverter *typeConverter = nullptr;
 
 private:
   using RewritePattern::rewrite;
 };
 
-/// OpConversionPattern is a wrapper around ConversionPattern that allows for
-/// matching and rewriting against an instance of a derived operation class as
-/// opposed to a raw Operation.
+namespace detail {
+/// OpOrInterfaceConversionPatternBase is a wrapper around ConversionPattern
+/// that allows for matching and rewriting against an instance of a derived
+/// operation class or an Interface as opposed to a raw Operation.
 template <typename SourceOp>
-struct OpConversionPattern : public ConversionPattern {
-  OpConversionPattern(MLIRContext *context, PatternBenefit benefit = 1)
-      : ConversionPattern(SourceOp::getOperationName(), benefit, context) {}
-  OpConversionPattern(TypeConverter &typeConverter, MLIRContext *context,
-                      PatternBenefit benefit = 1)
-      : ConversionPattern(SourceOp::getOperationName(), benefit, typeConverter,
-                          context) {}
+struct OpOrInterfaceConversionPatternBase : public ConversionPattern {
+  using ConversionPattern::ConversionPattern;
 
   /// Wrappers around the ConversionPattern methods that pass the derived op
   /// type.
   void rewrite(Operation *op, ArrayRef<Value> operands,
                ConversionPatternRewriter &rewriter) const final {
     rewrite(cast<SourceOp>(op), operands, rewriter);
   }
   LogicalResult
if (failed(match(op)))
       return failure();
     rewrite(op, operands, rewriter);
     return success();
   }
 
 private:
   using ConversionPattern::matchAndRewrite;
 };
+} // namespace detail
+
+/// OpConversionPattern is a wrapper around ConversionPattern that allows for
+/// matching and rewriting against an instance of a derived operation class as
+/// opposed to a raw Operation.
+template <typename SourceOp>
+struct OpConversionPattern
+    : public detail::OpOrInterfaceConversionPatternBase<SourceOp> {
+  OpConversionPattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceConversionPatternBase<SourceOp>(
+            SourceOp::getOperationName(), benefit, context) {}
+  OpConversionPattern(TypeConverter &typeConverter, MLIRContext *context,
+                      PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceConversionPatternBase<SourceOp>(
+            typeConverter, SourceOp::getOperationName(), benefit, context) {}
+};
+
+/// OpInterfaceConversionPattern is a wrapper around ConversionPattern that
+/// allows for matching and rewriting against an instance of an OpInterface
+/// class as opposed to a raw Operation.
+template <typename SourceOp>
+struct OpInterfaceConversionPattern
+    : public detail::OpOrInterfaceConversionPatternBase<SourceOp> {
+  OpInterfaceConversionPattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceConversionPatternBase<SourceOp>(
+            Pattern::MatchInterfaceOpTypeTag(), SourceOp::getInterfaceID(),
+            benefit, context) {}
+  OpInterfaceConversionPattern(TypeConverter &typeConverter,
+                               MLIRContext *context, PatternBenefit benefit = 1)
+      : detail::OpOrInterfaceConversionPatternBase<SourceOp>(
+            typeConverter, Pattern::MatchInterfaceOpTypeTag(),
+            SourceOp::getInterfaceID(), benefit, context) {}
+};
 
 /// Add a pattern to the given pattern list to convert the signature of a
 /// FunctionLike op with the given type converter. This only supports
 /// FunctionLike ops which use FunctionType to represent their type.
 void populateFunctionLikeTypeConversionPattern(StringRef functionLikeOpName,
                                                RewritePatternSet &patterns,
                                                TypeConverter &converter);
 

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

for (auto op : operands) {
     Value cast =
         b.create<memref::CastOp>(loc, eraseStridedLayout(memrefType), op);
     res.push_back(cast);
   }
   return res;
 }
 
 LogicalResult mlir::linalg::LinalgOpToLibraryCallRewrite::matchAndRewrite(
-    Operation *op, PatternRewriter &rewriter) const {
+    LinalgOp op, PatternRewriter &rewriter) const {
   // Only LinalgOp for which there is no specialized pattern go through this.
-  if (!isa<LinalgOp>(op) || isa<CopyOp>(op) || isa<IndexedGenericOp>(op))
+  if (isa<CopyOp>(op) || isa<IndexedGenericOp>(op))
     return failure();
 
   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
   if (!libraryCallName)
     return failure();
 
   rewriter.replaceOpWithNewOp<mlir::CallOp>(
       op, libraryCallName.getValue(), TypeRange(),

 void mlir::linalg::populateLinalgToStandardConversionPatterns(
     RewritePatternSet &patterns) {
   // TODO: ConvOp conversion needs to export a descriptor with relevant
   // attribute values such as kernel striding and dilation.
   // clang-format off
   patterns.add<
       CopyOpToLibraryCallRewrite,
       CopyTransposeRewrite,
-      IndexedGenericOpToLibraryCallRewrite>(patterns.getContext());
-  patterns.add<LinalgOpToLibraryCallRewrite>();
+      IndexedGenericOpToLibraryCallRewrite,
+      LinalgOpToLibraryCallRewrite>(patterns.getContext());
   // clang-format on
 }
 
 namespace {
 struct ConvertLinalgToStandardPass
     : public ConvertLinalgToStandardBase<ConvertLinalgToStandardPass> {
   void runOnOperation() override;
 };

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

if (auto memrefTy = stdTypes[i].dyn_cast<MemRefType>())
       values[i] = MemRefDescriptor::fromStaticShape(rewriter, loc, *this,
                                                     memrefTy, values[i]);
 }
 
 ConvertToLLVMPattern::ConvertToLLVMPattern(StringRef rootOpName,
                                            MLIRContext *context,
                                            LLVMTypeConverter &typeConverter,
                                            PatternBenefit benefit)
-    : ConversionPattern(rootOpName, benefit, typeConverter, context) {}
+    : ConversionPattern(typeConverter, rootOpName, benefit, context) {}
 
 //===----------------------------------------------------------------------===//
 // StructBuilder implementation
 //===----------------------------------------------------------------------===//
 
 StructBuilder::StructBuilder(Value v) : value(v), structType(v.getType()) {
   assert(value != nullptr && "value cannot be null");
   assert(LLVM::isCompatibleType(structType) && "expected llvm type");

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

static LogicalResult verifyNamedStructuredOp(NamedStructuredOpType op) {
   return verifyGenericOp<NamedStructuredOpType>(op);
 }
 
 //===----------------------------------------------------------------------===//
 // Canonicalizers and Folders.
 //===----------------------------------------------------------------------===//
 
 namespace {
-struct EraseDeadLinalgOp : public RewritePattern {
-  EraseDeadLinalgOp(PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()) {}
+struct EraseDeadLinalgOp : public OpInterfaceRewritePattern<LinalgOp> {
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    auto linalgOp = dyn_cast<LinalgOp>(op);
-    if (!linalgOp)
-      return failure();
-    for (Value v : linalgOp.getShapedOperands()) {
+    for (Value v : op.getShapedOperands()) {
       // Linalg "inputs" may be either tensor or memref type.
       // tensor<0xelt_type> is a convention that may not always mean
       // "0 iterations". Only erase in cases we see memref<...x0x...>.
       auto mt = v.getType().dyn_cast<MemRefType>();
       if (!mt)
         continue;
       if (llvm::is_contained(mt.getShape(), 0)) {
-        rewriter.eraseOp(linalgOp);
+        rewriter.eraseOp(op);
         return success();
       }
     }
     return failure();
   }
 };
 
-struct FoldTensorCastOp : public RewritePattern {
-  FoldTensorCastOp(PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()) {}
+struct FoldTensorCastOp : public OpInterfaceRewritePattern<LinalgOp> {
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    auto linalgOp = dyn_cast<LinalgOp>(op);
-    if (!linalgOp)
-      return failure();
-
     // If no operand comes from a tensor::CastOp and can be folded then fail.
     bool hasTensorCastOperand =
-        llvm::any_of(linalgOp.getShapedOperands(), [&](Value v) {
+        llvm::any_of(op.getShapedOperands(), [&](Value v) {
           if (v.isa<BlockArgument>())
             return false;
           auto castOp = v.getDefiningOp<tensor::CastOp>();
           return castOp && canFoldIntoConsumerOp(castOp);
         });
     if (!hasTensorCastOperand)
       return failure();
 
     SmallVector<Type, 4> newResultTypes;
     newResultTypes.reserve(op->getNumResults());
     SmallVector<Value, 4> newOperands;
     newOperands.reserve(op->getNumOperands());
     // Inputs may fold.
-    for (Value v : linalgOp.getInputs()) {
+    for (Value v : op.getInputs()) {
       auto tensorCastOp = v.getDefiningOp<tensor::CastOp>();
       newOperands.push_back(
           canFoldIntoConsumerOp(tensorCastOp) ? tensorCastOp.source() : v);
     }
     // Init tensors may fold, in which case the resultType must also change.
-    for (Value v : linalgOp.getOutputs()) {
+    for (Value v : op.getOutputs()) {
       auto tensorCastOp = v.getDefiningOp<tensor::CastOp>();
       bool fold = canFoldIntoConsumerOp(tensorCastOp);
       newOperands.push_back(fold ? tensorCastOp.getOperand() : v);
       newResultTypes.push_back(newOperands.back().getType());
     }
-    auto extraOperands = linalgOp.getAssumedNonShapedOperands();
+    auto extraOperands = op.getAssumedNonShapedOperands();
     newOperands.append(extraOperands.begin(), extraOperands.end());
     // Clone op.
     Operation *newOp =
-        linalgOp.clone(rewriter, op->getLoc(), newResultTypes, newOperands);
+        op.clone(rewriter, op->getLoc(), newResultTypes, newOperands);
     SmallVector<Value, 4> replacements;
     replacements.reserve(newOp->getNumResults());
     for (auto result : llvm::zip(op->getResults(), newOp->getResults())) {
       Value oldResult = std::get<0>(result);
       Value newResult = std::get<1>(result);
       if (newResult.getType() != oldResult.getType()) {
         replacements.push_back(rewriter.create<tensor::CastOp>(
             op->getLoc(), oldResult.getType(), newResult));
struct ReplaceDimOfLinalgOpResult : public OpRewritePattern<memref::DimOp> {
   }
 };
 
 } // namespace
 
 namespace {
 // Deduplicate redundant args of a linalg op.
 // An arg is redundant if it has the same Value and indexing map as another.
-struct DeduplicateInputs : public RewritePattern {
-  DeduplicateInputs(PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()) {}
+struct DeduplicateInputs : public OpInterfaceRewritePattern<LinalgOp> {
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
     // This pattern reduces the number of arguments of an op, which breaks
     // the invariants of semantically charged named ops.
     if (!isa<GenericOp, IndexedGenericOp>(op))
       return failure();
-    auto linalgOp = cast<LinalgOp>(op);
 
     // Associate each input to an equivalent "canonical" input that has the same
     // Value and indexing map.
     //
     // In the non-duplicate case, input `i` will have canonical input `i`. But
     // in the case of duplicated inputs, the canonical input could be some other
     // input `< i`. That is, a later input will have some earlier input as its
     // canonical input.
     llvm::SmallDenseMap<std::pair<Value, AffineMap>, int> canonicalInput;
     // For later remapping tasks like deduplicating payload block arguments,
     // having a simple "inputIndex -> canonicalInputIndex" integer mapping is
     // convenient.
     SmallVector<int, 6> canonicalInputIndices;
-    for (int i = 0, e = linalgOp.getNumInputs(); i != e; i++) {
-      Value input = linalgOp.getInput(i);
-      AffineMap indexingMap = linalgOp.getInputIndexingMap(i);
+    for (int i = 0, e = op.getNumInputs(); i != e; i++) {
+      Value input = op.getInput(i);
+      AffineMap indexingMap = op.getInputIndexingMap(i);
       // STL-like maps have a convenient behavior for our use case here. In the
       // case of duplicate keys, the insertion is rejected, and the returned
       // iterator gives access to the value already in the map.
       auto pair = canonicalInput.insert({{input, indexingMap}, i});
       canonicalInputIndices.push_back(pair.first->second);
     }
 
     // If there are no duplicate args, then bail out.
-    if (canonicalInput.size() == linalgOp.getNumInputs())
+    if (canonicalInput.size() == op.getNumInputs())
       return failure();
 
     // The operands for the newly canonicalized op.
     SmallVector<Value, 6> newOperands;
-    for (auto v : llvm::enumerate(linalgOp.getInputs()))
+    for (auto v : llvm::enumerate(op.getInputs()))
       if (canonicalInputIndices[v.index()] == static_cast<int>(v.index()))
         newOperands.push_back(v.value());
-    llvm::append_range(newOperands, linalgOp.getOutputs());
-    llvm::append_range(newOperands, linalgOp.getAssumedNonShapedOperands());
+    llvm::append_range(newOperands, op.getOutputs());
+    llvm::append_range(newOperands, op.getAssumedNonShapedOperands());
 
     // Clone the old op with new operands.
-    Operation *newOp = linalgOp.clone(rewriter, op->getLoc(),
-                                      op->getResultTypes(), newOperands);
+    Operation *newOp =
+        op.clone(rewriter, op->getLoc(), op->getResultTypes(), newOperands);
     auto newLinalgOp = cast<LinalgOp>(newOp);
 
     // Repair the indexing maps by filtering out the ones that have been
     // eliminated.
     SmallVector<AffineMap, 6> newIndexingMaps;
     for (int i = 0, e = newLinalgOp.getNumInputs(); i != e; i++)
       if (canonicalInputIndices[i] == i)
         newIndexingMaps.push_back(newLinalgOp.getIndexingMap(i));
     for (int i = 0, e = newLinalgOp.getNumOutputs(); i != e; i++)
       newIndexingMaps.push_back(newLinalgOp.getOutputIndexingMap(i));
     newOp->setAttr("indexing_maps",
                    rewriter.getAffineMapArrayAttr(newIndexingMaps));
 
     // Set the number of inputs to the new value. The `clone` call above kept
     // the value from the original op.
     newLinalgOp.setNumInputs(canonicalInput.size());
 
     // linalg.indexed_generic payloads have additional arguments prepended to
     // the block arg list.
     int bbArgBaseOffset = newLinalgOp.getNumPayloadInductionVariables();
 
     // Repair the payload entry block by RAUW'ing redundant arguments and
     // erasing them.
     Block &payload = newOp->getRegion(0).front();
-    for (int i = 0, e = linalgOp.getNumInputs(); i < e; i++) {
+    for (int i = 0, e = op.getNumInputs(); i < e; i++) {
       // Iterate in reverse, so that we erase later args first, preventing the
       // argument list from shifting unexpectedly and invalidating all our
       // indices.
       int reversed = e - i - 1;
       int canonicalIndex = canonicalInputIndices[reversed];
       if (canonicalInputIndices[reversed] == reversed)
         continue;
       payload.getArgument(bbArgBaseOffset + reversed)
           .replaceAllUsesWith(
               payload.getArgument(bbArgBaseOffset + canonicalIndex));
       payload.eraseArgument(bbArgBaseOffset + reversed);
     }
 
     rewriter.replaceOp(op, newOp->getResults());
     return success();
   }
 };
 
 /// Remove generic/indexed_generic operations (on tensors) that are just copying
 /// the values from inputs to the results. Requirements are
 /// 1) All iterator types are parallel
 /// 2) The body contains just a yield operation with the yielded values being
 ///    the arguments corresponding to the operands.
-struct RemoveIdentityLinalgOps : public RewritePattern {
-  RemoveIdentityLinalgOps(PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()) {}
+struct RemoveIdentityLinalgOps : public OpInterfaceRewritePattern<LinalgOp> {
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    if (auto copyOp = dyn_cast<CopyOp>(op)) {
+    if (auto copyOp = dyn_cast<CopyOp>(*op)) {
       assert(copyOp.hasBufferSemantics());
       if (copyOp.input() == copyOp.output() &&
           copyOp.inputPermutation() == copyOp.outputPermutation()) {
         rewriter.eraseOp(op);
         return success();
       }
     }
 
     if (!isa<GenericOp, IndexedGenericOp>(op))
       return failure();
-    LinalgOp genericOp = cast<LinalgOp>(op);
-    if (!genericOp.hasTensorSemantics())
+    if (!op.hasTensorSemantics())
       return failure();
     // Check all indexing maps are identity.
-    if (llvm::any_of(genericOp.getIndexingMaps(),
+    if (llvm::any_of(op.getIndexingMaps(),
                      [](AffineMap map) { return !map.isIdentity(); }))
       return failure();
 
     // Check that the body of the linalg operation is just a linalg.yield
     // operation.
     Block &body = op->getRegion(0).front();
     if (!llvm::hasSingleElement(body))
       return failure();
     auto yieldOp = dyn_cast<linalg::YieldOp>(body.getTerminator());
     if (!yieldOp)
       return failure();
 
     // Get the argument number of the returned values. That is the operand
     // number to use for replacing uses of this operation.
-    unsigned numIndexArgs = genericOp.getNumPayloadInductionVariables();
+    unsigned numIndexArgs = op.getNumPayloadInductionVariables();
     SmallVector<Value, 4> returnedArgs;
     for (Value yieldVal : yieldOp.values()) {
       auto yieldArg = yieldVal.dyn_cast<BlockArgument>();
       if (!yieldArg || yieldArg.getOwner() != &body)
         return failure();
       unsigned argumentNumber = yieldArg.getArgNumber();
       if (argumentNumber < numIndexArgs)
         return failure();
       returnedArgs.push_back(op->getOperand(argumentNumber - numIndexArgs));
     }
-    if (returnedArgs.size() != genericOp.getOperation()->getNumResults())
+    if (returnedArgs.size() != op.getOperation()->getNumResults())
       return failure();
-    rewriter.replaceOp(genericOp, returnedArgs);
+    rewriter.replaceOp(op, returnedArgs);
     return success();
   }
 };
 } // namespace
 
 #define CANONICALIZERS_AND_FOLDERS(XXX)                                        \
   void XXX::getCanonicalizationPatterns(RewritePatternSet &results,            \
                                         MLIRContext *context) {                \
     results.add<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp,        \
-                RemoveIdentityLinalgOps>();                                    \
-    results.add<ReplaceDimOfLinalgOpResult>(context);                          \
+                RemoveIdentityLinalgOps, ReplaceDimOfLinalgOpResult>(context); \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
                           SmallVectorImpl<OpFoldResult> &) {                   \
     return foldMemRefCast(*this);                                              \
   }
 
 CANONICALIZERS_AND_FOLDERS(ConvOp)

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

matchAndRewrite(FillOp op, ArrayRef<Value> operands,
     rewriter.replaceOp(op, adaptor.output());
 
     return success();
   }
 };
 
 /// Generic conversion pattern that matches any LinalgOp. This avoids template
 /// instantiating one pattern for each LinalgOp.
-class BufferizeAnyLinalgOp : public ConversionPattern {
+class BufferizeAnyLinalgOp : public OpInterfaceConversionPattern<LinalgOp> {
 public:
-  BufferizeAnyLinalgOp(TypeConverter &typeConverter)
-      : ConversionPattern(/*benefit=*/1, typeConverter, MatchAnyOpTypeTag()) {}
+  using OpInterfaceConversionPattern<LinalgOp>::OpInterfaceConversionPattern;
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(LinalgOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
-
-    LinalgOp linalgOp = dyn_cast<linalg::LinalgOp>(op);
-    if (!linalgOp)
+    // GenericOpAdaptor below expects an `operand_segment_sizes` attribute.
+    if (!op->hasAttr("operand_segment_sizes"))
       return failure();
 
     // We abuse the GenericOpAdaptor here.
     // TODO: Manually create an Adaptor that captures inputs and outputs for all
     // linalg::LinalgOp interface ops.
     linalg::GenericOpAdaptor adaptor(operands, op->getAttrDictionary());
 
-    Location loc = linalgOp.getLoc();
+    Location loc = op.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
-    if (failed(allocateBuffersForResults(loc, linalgOp, adaptor.outputs(),
+    if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
                                          newOutputBuffers, rewriter))) {
-      linalgOp.emitOpError()
+      return op.emitOpError()
              << "Failed to allocate buffers for tensor results.";
-      return failure();
     }
 
     // Delegate to the linalg generic pattern.
-    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
+    if (auto genericOp = dyn_cast<linalg::GenericOp>(*op)) {
       finalizeBufferAllocationForGenericOp<GenericOp>(
           rewriter, genericOp, adaptor.inputs(), newOutputBuffers);
       return success();
     }
 
     // Delegate to the linalg indexed generic pattern.
-    if (auto genericOp = dyn_cast<linalg::IndexedGenericOp>(op)) {
+    if (auto genericOp = dyn_cast<linalg::IndexedGenericOp>(*op)) {
       finalizeBufferAllocationForGenericOp<IndexedGenericOp>(
           rewriter, genericOp, adaptor.inputs(), newOutputBuffers);
       return success();
     }
 
-    finalizeBufferAllocation(rewriter, linalgOp, adaptor.inputs(),
-                             newOutputBuffers);
+    finalizeBufferAllocation(rewriter, op, adaptor.inputs(), newOutputBuffers);
     return success();
   }
 };
 
 /// Convert `subtensor %t [offsets][sizes][strides] -> %st` to an alloc + copy
 /// pattern.
 /// ```
 ///   %a = alloc(sizes)

 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgBufferizePass() {
   return std::make_unique<LinalgBufferizePass>();
 }
 
 void mlir::linalg::populateLinalgBufferizePatterns(
     BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) {
-  patterns.add<BufferizeAnyLinalgOp>(typeConverter);
   // TODO: Drop this once tensor constants work in standard.
   // clang-format off
   patterns.add<
+      BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
       SubTensorOpConverter,
       SubTensorInsertOpConverter
     >(typeConverter, patterns.getContext());
   // clang-format on
 }

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

public:
   }
 };
 
 /// A conversion pattern for detensoring internal (non-entry) blocks within a
 /// function.
 struct FunctionNonEntryBlockConversion : public ConversionPattern {
   FunctionNonEntryBlockConversion(StringRef functionLikeOpName,
                                   MLIRContext *ctx, TypeConverter &converter)
-      : ConversionPattern(functionLikeOpName, /*benefit=*/1, converter, ctx) {}
+      : ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.startRootUpdate(op);
 
     if (failed(rewriter.convertNonEntryRegionTypes(
             &mlir::impl::getFunctionBody(op), *typeConverter))) {

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

for (Type t : rankedTensorTypes) {
     res.push_back(b.create<linalg::InitTensorOp>(
         loc, dynamicShape, staticShape, rankedTensorType.getElementType()));
   }
   return res;
 }
 
 namespace {
 struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
-  ConvertAnyElementwiseMappableOpOnRankedTensors()
-      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
+  ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const final {
     if (!isElementwiseMappableOpOnRankedTensors(op))
       return rewriter.notifyMatchFailure(
           op, "requires elementwise op on ranked tensors");
 
     auto rank = op->getResult(0).getType().cast<RankedTensorType>().getRank();
     SmallVector<AffineMap, 3> indexingMaps(
rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         });
     return success();
   }
 };
 } // namespace
 
 void mlir::populateElementwiseToLinalgConversionPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<ConvertAnyElementwiseMappableOpOnRankedTensors>();
+  patterns.add<ConvertAnyElementwiseMappableOpOnRankedTensors>(
+      patterns.getContext());
 }
 
 namespace {
 class ConvertElementwiseToLinalgPass
     : public ConvertElementwiseToLinalgBase<ConvertElementwiseToLinalgPass> {
 
   void runOnFunction() final {
     auto func = getOperation();

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

 };
 
 /// Catch-all pattern for converting all named ops with a region builder into
 /// linalg.generic.
 struct LinalgNamedOpGeneralizationPattern : RewritePattern {
   LinalgNamedOpGeneralizationPattern(MLIRContext *context,
                                      linalg::LinalgTransformationFilter marker,
                                      PatternBenefit benefit = 1)
-      : RewritePattern(benefit, MatchAnyOpTypeTag()),
+      : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
         marker(std::move(marker)) {}
 
   LogicalResult matchAndRewrite(Operation *rootOp,
                                 PatternRewriter &rewriter) const override {
     auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp);
     if (!linalgOp)
       return failure();
     if (failed(marker.checkAndNotify(rewriter, linalgOp)))

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

linalgOpToLoopsImpl(Operation *op, OpBuilder &builder,
   LinalgLoops loops(loopSet.begin(), loopSet.end());
   return loops;
 }
 
 namespace {
 template <typename LoopType>
 class LinalgRewritePattern : public RewritePattern {
 public:
-  LinalgRewritePattern(ArrayRef<unsigned> interchangeVector)
-      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()),
+  LinalgRewritePattern(MLIRContext *context,
+                       ArrayRef<unsigned> interchangeVector)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
         interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (!isa<LinalgOp>(op))
       return failure();
     if (!linalgOpToLoopsImpl<LoopType>(op, rewriter, interchangeVector))
       return failure();
     rewriter.eraseOp(op);
     return success();
   }
 
 private:
   SmallVector<unsigned, 4> interchangeVector;
 };
 
 struct FoldAffineOp;
 } // namespace
 
 template <typename LoopType>
 static void lowerLinalgToLoopsImpl(FuncOp funcOp,
                                    ArrayRef<unsigned> interchangeVector) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet patterns(context);
-  patterns.add<LinalgRewritePattern<LoopType>>(interchangeVector);
+  patterns.add<LinalgRewritePattern<LoopType>>(context, interchangeVector);
   memref::DimOp::getCanonicalizationPatterns(patterns, context);
   AffineApplyOp::getCanonicalizationPatterns(patterns, context);
   patterns.add<FoldAffineOp>(context);
   // Just apply the patterns greedily.
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 namespace {

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

static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
   res = TiledLinalgOp{paddedOp, res.loops, res.tensorResults};
   return success();
 }
 
 /// Linalg base tiling pattern.
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
-    : RewritePattern(opName, {}, benefit, context), filter(filter),
+    : RewritePattern(opName, benefit, context), filter(filter),
       options(options) {}
 
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
-    LinalgTilingOptions options, LinalgTransformationFilter filter,
-    PatternBenefit benefit)
-    : RewritePattern(benefit, MatchAnyOpTypeTag()), filter(filter),
+    MLIRContext *context, LinalgTilingOptions options,
+    LinalgTransformationFilter filter, PatternBenefit benefit)
+    : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
     Operation *op, PatternRewriter &rewriter, TiledLinalgOp &result) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))

 }
 
 mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
     StringRef opName, MLIRContext *context,
     const LinalgDependenceGraph &dependenceGraph,
     LinalgTilingOptions tilingOptions, LinalgFusionOptions fusionOptions,
     LinalgTransformationFilter filter, LinalgTransformationFilter fusedOpMarker,
     LinalgTransformationFilter originalOpMarker, PatternBenefit benefit)
-    : RewritePattern(opName, {}, benefit, context),
+    : RewritePattern(opName, benefit, context, {}),
       dependenceGraph(dependenceGraph), tilingOptions(tilingOptions),
       fusionOptions(fusionOptions), filter(filter),
       fusedOpMarker(fusedOpMarker), originalOpMarker(originalOpMarker) {}
 
 LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
   return success();
 }
 
 /// Linalg base interchange pattern.
 mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
     StringRef opName, MLIRContext *context,
     ArrayRef<unsigned> interchangeVector, LinalgTransformationFilter filter,
     PatternBenefit benefit)
-    : RewritePattern(opName, {}, benefit, context), filter(filter),
+    : RewritePattern(opName, benefit, context, {}), filter(filter),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
 LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
rewriter.updateRootInPlace(op, [&]() {
     filter.replaceLinalgTransformationFilter(rewriter, op);
   });
   return success();
 }
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
-    : RewritePattern(opName, {}, benefit, context), filter(filter),
+    : RewritePattern(opName, benefit, context, {}), filter(filter),
       options(options) {}
 
 LogicalResult mlir::linalg::LinalgBasePromotionPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   if (failed(filter.checkAndNotify(rewriter, op)))
     return failure();
   if (failed(promoteSubviewsPrecondition(op, options)))
     return failure();
if (!promotedOp) {
     return op->emitError("subview promotion failed");
   }
   rewriter.finalizeRootUpdate(op);
   filter.replaceLinalgTransformationFilter(rewriter, op);
   return success();
 }
 
 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
-    LinalgTransformationFilter filter, PatternBenefit benefit)
-    : RewritePattern(benefit, MatchAnyOpTypeTag()), filter(filter) {}
+    MLIRContext *context, LinalgTransformationFilter filter,
+    PatternBenefit benefit)
+    : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter) {}
 
 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
-    : RewritePattern(opName, {}, benefit, context), filter(filter) {}
+    : RewritePattern(opName, benefit, context, {}), filter(filter) {}
 
 LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

void mlir::populateCallOpTypeConversionPattern(RewritePatternSet &patterns,
                                                TypeConverter &converter) {
   patterns.add<CallOpSignatureConversion>(converter, patterns.getContext());
 }
 
 namespace {
 /// Only needed to support partial conversion of functions where this pattern
 /// ensures that the branch operation arguments matches up with the succesor
 /// block arguments.
-class BranchOpInterfaceTypeConversion : public ConversionPattern {
+class BranchOpInterfaceTypeConversion
+    : public OpInterfaceConversionPattern<BranchOpInterface> {
 public:
-  BranchOpInterfaceTypeConversion(TypeConverter &typeConverter,
-                                  MLIRContext *ctx)
-      : ConversionPattern(/*benefit=*/1, typeConverter, MatchAnyOpTypeTag()) {}
+  using OpInterfaceConversionPattern<
+      BranchOpInterface>::OpInterfaceConversionPattern;
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(BranchOpInterface op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
-    auto branchOp = dyn_cast<BranchOpInterface>(op);
-    if (!branchOp)
-      return failure();
-
     // For a branch operation, only some operands go to the target blocks, so
     // only rewrite those.
     SmallVector<Value, 4> newOperands(op->operand_begin(), op->operand_end());
     for (int succIdx = 0, succEnd = op->getBlock()->getNumSuccessors();
          succIdx < succEnd; ++succIdx) {
-      auto successorOperands = branchOp.getSuccessorOperands(succIdx);
+      auto successorOperands = op.getSuccessorOperands(succIdx);
       if (!successorOperands)
         continue;
       for (int idx = successorOperands->getBeginOperandIndex(),
                eidx = idx + successorOperands->size();
            idx < eidx; ++idx) {
         newOperands[idx] = operands[idx];
       }
     }

mlir/lib/IR/PatternMatch.cpp

unsigned short PatternBenefit::getBenefit() const {
   assert(!isImpossibleToMatch() && "Pattern doesn't match");
   return representation;
 }
 
 //===----------------------------------------------------------------------===//
 // Pattern
 //===----------------------------------------------------------------------===//
 
+//===----------------------------------------------------------------------===//
+// OperationName Root Constructors
+
 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,
+                 MLIRContext *context, ArrayRef<StringRef> generatedNames)
+    : Pattern(OperationName(rootName, context).getAsOpaquePointer(),
+              RootKind::OperationName, generatedNames, benefit, context) {}
+
+//===----------------------------------------------------------------------===//
+// MatchAnyOpTypeTag Root Constructors
+
+Pattern::Pattern(MatchAnyOpTypeTag tag, PatternBenefit benefit,
+                 MLIRContext *context, ArrayRef<StringRef> generatedNames)
+    : Pattern(nullptr, RootKind::Any, generatedNames, benefit, context) {}
+
+//===----------------------------------------------------------------------===//
+// MatchInterfaceOpTypeTag Root Constructors
+
+Pattern::Pattern(MatchInterfaceOpTypeTag tag, TypeID interfaceID,
+                 PatternBenefit benefit, MLIRContext *context,
+                 ArrayRef<StringRef> generatedNames)
+    : Pattern(interfaceID.getAsOpaquePointer(), RootKind::InterfaceID,
+              generatedNames, benefit, context) {}
+
+//===----------------------------------------------------------------------===//
+// MatchTraitOpTypeTag Root Constructors
+
+Pattern::Pattern(MatchTraitOpTypeTag tag, TypeID traitID,
+                 PatternBenefit benefit, MLIRContext *context,
+                 ArrayRef<StringRef> generatedNames)
+    : Pattern(traitID.getAsOpaquePointer(), RootKind::TraitID, generatedNames,
+              benefit, context) {}
+
+//===----------------------------------------------------------------------===//
+// General Constructors
+
+Pattern::Pattern(const void *rootValue, RootKind rootKind,
+                 ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
                  MLIRContext *context)
-    : rootKind(OperationName(rootName, context)), benefit(benefit) {}
-Pattern::Pattern(PatternBenefit benefit, MatchAnyOpTypeTag tag)
-    : benefit(benefit) {}
-Pattern::Pattern(StringRef rootName, ArrayRef<StringRef> generatedNames,
-                 PatternBenefit benefit, MLIRContext *context)
-    : Pattern(rootName, benefit, context) {
-  generatedOps.reserve(generatedNames.size());
-  std::transform(generatedNames.begin(), generatedNames.end(),
-                 std::back_inserter(generatedOps), [context](StringRef name) {
-                   return OperationName(name, context);
-                 });
-}
-Pattern::Pattern(ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
-                 MLIRContext *context, MatchAnyOpTypeTag tag)
-    : Pattern(benefit, tag) {
+    : rootValue(rootValue), rootKind(rootKind), benefit(benefit),
+      contextAndHasBoundedRecursion(context, false) {
+  if (generatedNames.empty())
+    return;
   generatedOps.reserve(generatedNames.size());
   std::transform(generatedNames.begin(), generatedNames.end(),
                  std::back_inserter(generatedOps), [context](StringRef name) {
                    return OperationName(name, context);
                  });
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Rewrite/ByteCode.cpp

if (ArrayAttr generatedOpsAttr = matchOp.generatedOpsAttr())
     generatedOps =
         llvm::to_vector<8>(generatedOpsAttr.getAsValueRange<StringAttr>());
 
   PatternBenefit benefit = matchOp.benefit();
   MLIRContext *ctx = matchOp.getContext();
 
   // Check to see if this is pattern matches a specific operation type.
   if (Optional<StringRef> rootKind = matchOp.rootKind())
-    return PDLByteCodePattern(rewriterAddr, *rootKind, generatedOps, benefit,
-                              ctx);
-  return PDLByteCodePattern(rewriterAddr, generatedOps, benefit, ctx,
-                            MatchAnyOpTypeTag());
+    return PDLByteCodePattern(rewriterAddr, *rootKind, benefit, ctx,
+                              generatedOps);
+  return PDLByteCodePattern(rewriterAddr, MatchAnyOpTypeTag(), benefit, ctx,
+                            generatedOps);
 }
 
 //===----------------------------------------------------------------------===//
 // PDLByteCodeMutableState
 //===----------------------------------------------------------------------===//
 
 /// Set the new benefit for a bytecode pattern. The `patternIndex` corresponds
 /// to the position of the pattern within the range returned by

mlir/lib/Rewrite/FrozenRewritePatternSet.cpp

 // FrozenRewritePatternSet
 //===----------------------------------------------------------------------===//
 
 FrozenRewritePatternSet::FrozenRewritePatternSet()
     : impl(std::make_shared<Impl>()) {}
 
 FrozenRewritePatternSet::FrozenRewritePatternSet(RewritePatternSet &&patterns)
     : impl(std::make_shared<Impl>()) {
-  impl->nativePatterns = std::move(patterns.getNativePatterns());
+  // Functor used to walk all of the operations registered in the context. This
+  // is useful for patterns that get applied to multiple operations, such as
+  // interface and trait based patterns.
+  std::vector<AbstractOperation *> abstractOps;
+  auto addToOpsWhen = [&](std::unique_ptr<RewritePattern> &pattern,
+                          function_ref<bool(AbstractOperation *)> callbackFn) {
+    if (abstractOps.empty())
+      abstractOps = pattern->getContext()->getRegisteredOperations();
+    for (AbstractOperation *absOp : abstractOps) {
+      if (callbackFn(absOp)) {
+        OperationName opName(absOp);
+        impl->nativeOpSpecificPatternMap[opName].push_back(pattern.get());
+      }
+    }
+    impl->nativeOpSpecificPatternList.push_back(std::move(pattern));
+  };
+
+  for (std::unique_ptr<RewritePattern> &pat : patterns.getNativePatterns()) {
+    if (Optional<OperationName> rootName = pat->getRootKind()) {
+      impl->nativeOpSpecificPatternMap[*rootName].push_back(pat.get());
+      impl->nativeOpSpecificPatternList.push_back(std::move(pat));
+      continue;
+    }
+    if (Optional<TypeID> interfaceID = pat->getRootInterfaceID()) {
+      addToOpsWhen(pat, [&](AbstractOperation *absOp) {
+        return absOp->hasInterface(*interfaceID);
+      });
+      continue;
+    }
+    if (Optional<TypeID> traitID = pat->getRootTraitID()) {
+      addToOpsWhen(pat, [&](AbstractOperation *absOp) {
+        return absOp->hasTrait(*traitID);
+      });
+      continue;
+    }
+    impl->nativeAnyOpPatterns.push_back(std::move(pat));
+  }
 
   // Generate the bytecode for the PDL patterns if any were provided.
   PDLPatternModule &pdlPatterns = patterns.getPDLPatterns();
   ModuleOp pdlModule = pdlPatterns.getModule();
   if (!pdlModule)
     return;
   if (failed(convertPDLToPDLInterp(pdlModule)))
     llvm::report_fatal_error(

mlir/lib/Rewrite/PatternApplicator.cpp

 // user defined cost model.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "ByteCode.h"
 #include "llvm/Support/Debug.h"
 
+#define DEBUG_TYPE "pattern-match"
+
 using namespace mlir;
 using namespace mlir::detail;
 
 PatternApplicator::PatternApplicator(
     const FrozenRewritePatternSet &frozenPatternList)
     : frozenPatternList(frozenPatternList) {
   if (const PDLByteCode *bytecode = frozenPatternList.getPDLByteCode()) {
     mutableByteCodeState = std::make_unique<PDLByteCodeMutableState>();
     bytecode->initializeMutableState(*mutableByteCodeState);
   }
 }
 PatternApplicator::~PatternApplicator() {}
 
-#define DEBUG_TYPE "pattern-match"
+/// Log a message for a pattern that is impossible to match.
+static void logImpossibleToMatch(const Pattern &pattern) {
+  LLVM_DEBUG({
+    llvm::dbgs() << "Ignoring pattern '" << pattern.getRootKind()
+                 << "' because it is impossible to match or cannot lead "
+                    "to legal IR (by cost model)\n";
+  });
+}
 
 void PatternApplicator::applyCostModel(CostModel model) {
   // Apply the cost model to the bytecode patterns first, and then the native
   // patterns.
   if (const PDLByteCode *bytecode = frozenPatternList.getPDLByteCode()) {
     for (auto it : llvm::enumerate(bytecode->getPatterns()))
       mutableByteCodeState->updatePatternBenefit(it.index(), model(it.value()));
   }
 
-  // Separate patterns by root kind to simplify lookup later on.
+  // Copy over the patterns so that we can sort by benefit based on the cost
+  // model. Patterns that are already impossible to match are ignored.
   patterns.clear();
-  anyOpPatterns.clear();
-  for (const auto &pat : frozenPatternList.getNativePatterns()) {
-    // If the pattern is always impossible to match, just ignore it.
-    if (pat.getBenefit().isImpossibleToMatch()) {
-      LLVM_DEBUG({
-        llvm::dbgs()
-            << "Ignoring pattern '" << pat.getRootKind()
-            << "' because it is impossible to match (by pattern benefit)\n";
-      });
-      continue;
-    }
-    if (Optional<OperationName> opName = pat.getRootKind())
-      patterns[*opName].push_back(&pat);
+  for (const auto &it : frozenPatternList.getOpSpecificNativePatterns()) {
+    for (const RewritePattern *pattern : it.second) {
+      if (pattern->getBenefit().isImpossibleToMatch())
+        logImpossibleToMatch(*pattern);
+      else
+        patterns[it.first].push_back(pattern);
+    }
+  anyOpPatterns.clear();
+  for (const RewritePattern &pattern :
+       frozenPatternList.getMatchAnyOpNativePatterns()) {
+    if (pattern.getBenefit().isImpossibleToMatch())
+      logImpossibleToMatch(pattern);
     else
-      anyOpPatterns.push_back(&pat);
+      anyOpPatterns.push_back(&pattern);
   }
 
   // Sort the patterns using the provided cost model.
   llvm::SmallDenseMap<const Pattern *, PatternBenefit> benefits;
   auto cmp = [&benefits](const Pattern *lhs, const Pattern *rhs) {
     return benefits[lhs] > benefits[rhs];
   };
   auto processPatternList = [&](SmallVectorImpl<const 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()) {
-        LLVM_DEBUG({
-          llvm::dbgs() << "Ignoring pattern '" << list.front()->getRootKind()
-                       << "' because it is impossible to match or cannot lead "
-                          "to legal IR (by cost model)\n";
-        });
+        logImpossibleToMatch(*list.front());
         list.clear();
       }
       return;
     }
 
     // Collect the dynamic benefits for the current pattern list.
     benefits.clear();
     for (const Pattern *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()) {
-      LLVM_DEBUG({
-        llvm::dbgs() << "Ignoring pattern '" << list.back()->getRootKind()
-                     << "' because it is impossible to match or cannot lead to "
-                        "legal IR (by cost model)\n";
-      });
-      list.pop_back();
-    }
+    while (!list.empty() && benefits[list.back()].isImpossibleToMatch())
+      logImpossibleToMatch(*list.pop_back_val());
   };
   for (auto &it : patterns)
     processPatternList(it.second);
   processPatternList(anyOpPatterns);
 }
 
 void PatternApplicator::walkAllPatterns(
     function_ref<void(const Pattern &)> walk) {
-  for (const Pattern &it : frozenPatternList.getNativePatterns())
+  for (const auto &it : frozenPatternList.getOpSpecificNativePatterns())
+    for (const auto &pattern : it.second)
+      walk(*pattern);
+  for (const Pattern &it : frozenPatternList.getMatchAnyOpNativePatterns())
     walk(it);
   if (const PDLByteCode *bytecode = frozenPatternList.getPDLByteCode()) {
     for (const Pattern &it : bytecode->getPatterns())
       walk(it);
   }
 }
 
 LogicalResult PatternApplicator::matchAndRewrite(

mlir/lib/Transforms/Utils/DialectConversion.cpp

 
 /// Create a default conversion pattern that rewrites the type signature of a
 /// FunctionLike op. This only supports FunctionLike ops which use FunctionType
 /// to represent their type.
 namespace {
 struct FunctionLikeSignatureConversion : public ConversionPattern {
   FunctionLikeSignatureConversion(StringRef functionLikeOpName,
                                   MLIRContext *ctx, TypeConverter &converter)
-      : ConversionPattern(functionLikeOpName, /*benefit=*/1, converter, ctx) {}
+      : ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx) {}
 
   /// Hook to implement combined matching and rewriting for FunctionLike ops.
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     FunctionType type = mlir::impl::getFunctionType(op);
 
     // Convert the original function types.

mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp

patterns.add<ConvertToAtomCmpExchangeWeak, ConvertToBitReverse,
                ConvertToGroupNonUniformBallot, ConvertToModule,
                ConvertToSubgroupBallot>(context);
 
   if (failed(applyPartialConversion(fn, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 ConvertToAtomCmpExchangeWeak::ConvertToAtomCmpExchangeWeak(MLIRContext *context)
-    : RewritePattern("test.convert_to_atomic_compare_exchange_weak_op",
-                     {"spv.AtomicCompareExchangeWeak"}, 1, context) {}
+    : RewritePattern("test.convert_to_atomic_compare_exchange_weak_op", 1,
+                     context, {"spv.AtomicCompareExchangeWeak"}) {}
 
 LogicalResult
 ConvertToAtomCmpExchangeWeak::matchAndRewrite(Operation *op,
                                               PatternRewriter &rewriter) const {
   Value ptr = op->getOperand(0);
   Value value = op->getOperand(1);
   Value comparator = op->getOperand(2);
 
   // Create a spv.AtomicCompareExchangeWeak op with AtomicCounterMemory bits in
   // memory semantics to additionally require AtomicStorage capability.
   rewriter.replaceOpWithNewOp<spirv::AtomicCompareExchangeWeakOp>(
       op, value.getType(), ptr, spirv::Scope::Workgroup,
       spirv::MemorySemantics::AcquireRelease |
           spirv::MemorySemantics::AtomicCounterMemory,
       spirv::MemorySemantics::Acquire, value, comparator);
   return success();
 }
 
 ConvertToBitReverse::ConvertToBitReverse(MLIRContext *context)
-    : RewritePattern("test.convert_to_bit_reverse_op", {"spv.BitReverse"}, 1,
-                     context) {}
+    : RewritePattern("test.convert_to_bit_reverse_op", 1, context,
+                     {"spv.BitReverse"}) {}
 
 LogicalResult
 ConvertToBitReverse::matchAndRewrite(Operation *op,
                                      PatternRewriter &rewriter) const {
   Value predicate = op->getOperand(0);
 
   rewriter.replaceOpWithNewOp<spirv::BitReverseOp>(
       op, op->getResult(0).getType(), predicate);
   return success();
 }
 
 ConvertToGroupNonUniformBallot::ConvertToGroupNonUniformBallot(
     MLIRContext *context)
-    : RewritePattern("test.convert_to_group_non_uniform_ballot_op",
-                     {"spv.GroupNonUniformBallot"}, 1, context) {}
+    : RewritePattern("test.convert_to_group_non_uniform_ballot_op", 1, context,
+                     {"spv.GroupNonUniformBallot"}) {}
 
 LogicalResult ConvertToGroupNonUniformBallot::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   Value predicate = op->getOperand(0);
 
   rewriter.replaceOpWithNewOp<spirv::GroupNonUniformBallotOp>(
       op, op->getResult(0).getType(), spirv::Scope::Workgroup, predicate);
   return success();
 }
 
 ConvertToModule::ConvertToModule(MLIRContext *context)
-    : RewritePattern("test.convert_to_module_op", {"spv.module"}, 1, context) {}
+    : RewritePattern("test.convert_to_module_op", 1, context, {"spv.module"}) {}
 
 LogicalResult
 ConvertToModule::matchAndRewrite(Operation *op,
                                  PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<spirv::ModuleOp>(
       op, spirv::AddressingModel::PhysicalStorageBuffer64,
       spirv::MemoryModel::Vulkan);
   return success();
 }
 
 ConvertToSubgroupBallot::ConvertToSubgroupBallot(MLIRContext *context)
-    : RewritePattern("test.convert_to_subgroup_ballot_op",
-                     {"spv.SubgroupBallotKHR"}, 1, context) {}
+    : RewritePattern("test.convert_to_subgroup_ballot_op", 1, context,
+                     {"spv.SubgroupBallotKHR"}) {}
 
 LogicalResult
 ConvertToSubgroupBallot::matchAndRewrite(Operation *op,
                                          PatternRewriter &rewriter) const {
   Value predicate = op->getOperand(0);
 
   rewriter.replaceOpWithNewOp<spirv::SubgroupBallotKHROp>(
       op, op->getResult(0).getType(), predicate);

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

 };
 
 //===----------------------------------------------------------------------===//
 // Type-Conversion Rewrite Testing
 
 /// This patterns erases a region operation that has had a type conversion.
 struct TestDropOpSignatureConversion : public ConversionPattern {
   TestDropOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
-      : ConversionPattern("test.drop_region_op", 1, converter, ctx) {}
+      : ConversionPattern(converter, "test.drop_region_op", 1, ctx) {}
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Region &region = op->getRegion(0);
     Block *entry = &region.front();
 
     // Convert the original entry arguments.
     TypeConverter &converter = *getTypeConverter();

 
 //===----------------------------------------------------------------------===//
 // 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, MatchAnyOpTypeTag()) {}
+  RemoveTestDialectOps(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
 
   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::RewritePatternSet patterns(&getContext());
-    patterns.add<RemoveTestDialectOps>();
+    patterns.add<RemoveTestDialectOps>(&getContext());
 
     mlir::ConversionTarget target(getContext());
     target.addIllegalDialect<TestDialect>();
     if (failed(
             applyPartialConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
   }
 };

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

static void applyPatterns(FuncOp funcOp) {
   // Linalg distribution patterns.
   //===--------------------------------------------------------------------===//
   LinalgLoopDistributionOptions distributionOptions;
 
   //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.add<LinalgVectorizationPattern>(
-      LinalgTransformationFilter(Identifier::get("VECTORIZE", ctx))
+      ctx, LinalgTransformationFilter(Identifier::get("VECTORIZE", ctx))
                .addOpFilter<MatmulOp, FillOp, CopyOp, GenericOp>());
 
   //===--------------------------------------------------------------------===//
   // Linalg generic permutation patterns.
   //===--------------------------------------------------------------------===//
   patterns.add<LinalgInterchangePattern<GenericOp>>(
       ctx,
       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
       LinalgTransformationFilter(ArrayRef<Identifier>{},
patternsVector.emplace_back(
           LinalgTransformationFilter(Identifier::get("L1", ctx),
                                      Identifier::get("VEC", ctx))));
 
   patternsVector.emplace_back(
       ctx, std::make_unique<LinalgVectorizationPattern>(
                MatmulOp::getOperationName(), ctx, LinalgVectorizationOptions(),
                LinalgTransformationFilter(Identifier::get("VEC", ctx))));
   patternsVector.back().add<LinalgVectorizationPattern>(
-      LinalgTransformationFilter().addFilter(
+      ctx, LinalgTransformationFilter().addFilter(
                [](Operation *op) { return success(isa<FillOp, CopyOp>(op)); }));
 }
 
 //===----------------------------------------------------------------------===//
 // Test promotion callbacks
 //===----------------------------------------------------------------------===//
 
 // Allocation call back
 static Optional<Value> allocCallBackFn(OpBuilder &b, memref::SubViewOp subView,
static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   forwardPattern.add<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
   forwardPattern.add<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(forwardPattern));
 }
 
 static void applyLinalgToVectorPatterns(FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   patterns.add<LinalgVectorizationPattern>(
+      funcOp.getContext(),
       LinalgTransformationFilter()
           .addOpFilter<ContractionOpInterface, FillOp, CopyOp, GenericOp>());
   patterns.add<PadTensorOpVectorizationPattern>(funcOp.getContext());
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   RewritePatternSet foldPattern(funcOp.getContext());

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

 }
 
 void TCParser::printCanonicalizersAndFolders(llvm::raw_ostream &os,
                                              StringRef cppOpName) {
   const char *canonicalizersAndFoldersFmt = R"FMT(
     void {0}::getCanonicalizationPatterns(
         RewritePatternSet &results,
         MLIRContext *context) {{
-      results.add<EraseDeadLinalgOp>();
-      results.add<FoldTensorCastOp>();
+      results.add<EraseDeadLinalgOp>(context);
+      results.add<FoldTensorCastOp>(context);
     }
     LogicalResult {0}::fold(ArrayRef<Attribute>,
                             SmallVectorImpl<OpFoldResult> &) {{
       return foldMemRefCast(*this);
     }
     void {0}::getEffects(SmallVectorImpl<
         SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
       getGenericEffectsImpl(effects,

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

 
 // Implementations of getCanonicalizationPatterns, fold and getEffects.
 // Parameters:
 // {0}: Class name
 const char structuredOpCanonicalizersAndFoldersFormat[] = R"FMT(
 void {0}::getCanonicalizationPatterns(
     RewritePatternSet &results,
     MLIRContext *context) {{
-  results.add<EraseDeadLinalgOp>();
-  results.add<FoldTensorCastOp>();
+  results.add<EraseDeadLinalgOp>(context);
+  results.add<FoldTensorCastOp>(context);
 }
 LogicalResult {0}::fold(ArrayRef<Attribute>,
                         SmallVectorImpl<OpFoldResult> &) {{
   return foldMemRefCast(*this);
 }
 void {0}::getEffects(SmallVectorImpl<
     SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
   getGenericEffectsImpl(effects,

mlir/tools/mlir-tblgen/RewriterGen.cpp

void PatternEmitter::emit(StringRef rewriteName) {
   LLVM_DEBUG(llvm::dbgs() << "done collecting ops used in result patterns\n");
 
   // Emit RewritePattern for Pattern.
   auto locs = pattern.getLocation();
   os << formatv("/* Generated from:\n    {0:$[ instantiating\n    ]}\n*/\n",
                 make_range(locs.rbegin(), locs.rend()));
   os << formatv(R"(struct {0} : public ::mlir::RewritePattern {
   {0}(::mlir::MLIRContext *context)
-      : ::mlir::RewritePattern("{1}", {{)",
-                rewriteName, rootName);
+      : ::mlir::RewritePattern("{1}", {2}, context, {{)",
+                rewriteName, rootName, pattern.getBenefit());
   // Sort result operators by name.
   llvm::SmallVector<const Operator *, 4> sortedResultOps(resultOps.begin(),
                                                          resultOps.end());
   llvm::sort(sortedResultOps, [&](const Operator *lhs, const Operator *rhs) {
     return lhs->getOperationName() < rhs->getOperationName();
   });
   llvm::interleaveComma(sortedResultOps, os, [&](const Operator *op) {
     os << '"' << op->getOperationName() << '"';
   });
-  os << formatv(R"(}, {0}, context) {{})", pattern.getBenefit()) << "\n";
+  os << "}) {}\n";
 
   // Emit matchAndRewrite() function.
   {
     auto classScope = os.scope();
     os.reindent(R"(
     ::mlir::LogicalResult matchAndRewrite(::mlir::Operation *op0,
         ::mlir::PatternRewriter &rewriter) const override {)")
         << '\n';

mlir/unittests/Rewrite/PatternBenefit.cpp

matchAndRewrite(ModuleOp /*op*/,
       return failure();
     }
 
   private:
     bool *called;
   };
 
   struct Pattern2 : public RewritePattern {
-    Pattern2(bool *called)
-        : RewritePattern(/*benefit*/ 2, MatchAnyOpTypeTag{}), called(called) {}
+    Pattern2(MLIRContext *context, bool *called)
+        : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/2, context),
+          called(called) {}
 
     mlir::LogicalResult
     matchAndRewrite(Operation * /*op*/,
                     mlir::PatternRewriter & /*rewriter*/) const override {
       *called = true;
       return failure();
     }
 
   private:
     bool *called;
   };
 
   RewritePatternSet patterns(&context);
 
   bool called1 = false;
   bool called2 = false;
 
   patterns.add<Pattern1>(&context, &called1);
-  patterns.add<Pattern2>(&called2);
+  patterns.add<Pattern2>(&context, &called2);
 
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
   PatternApplicator pa(frozenPatterns);
   pa.applyDefaultCostModel();
 
   class MyPatternRewriter : public PatternRewriter {
   public:
     MyPatternRewriter(MLIRContext *ctx) : PatternRewriter(ctx) {}