[mlir] move PDL-related transform ops into an extension

[ftynse]

Summary

The initial bring-up of the Transform dialect relied on PDL to provide
the default handle type (!pdl.operation) and the matching capability.
Both are now provided natively by the Transform dialect removing the
reason to have a hard dependency on the PDL dialect and its interpreter.
Move PDL-related transform operations into a separate extension.

This requires us to introduce a dialect state extension mechanism into
the Transform dialect so it no longer needs to know about PDL constraint
functions that may be injected by extensions similarly to operations and
types. This mechanism will be reused to connect pattern application
drivers and the Transform dialect.

This completes the restructuring of the Transform dialect to remove
overrilance on PDL.

Note to downstreams: flow that are using !pdl.operation with Transform
dialect operations will now require transform::PDLExtension to be
applied to the transform dialect in order to provide the transform
handle type interface for !pdl.operation.

[springerm]

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td

19

Is this op now deprecated?

[ftynse]

Add dedicated test for the constraint injection

[ftynse]

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td

19

Not really, there may be cases where using PDL matching is interesting, but not necessarily in the Linalg/Vector case.

[ftynse]

Add some more Bazel

[ftynse]

More bazel.

[ftynse]

Bazel is annoying.

mlir/include/mlir/Dialect/Transform/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(PDLExtension)
 add_subdirectory(Transforms)

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h

 //===- TransformDialect.h - Transform Dialect Definition --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include <optional>
 
 namespace mlir {
 namespace transform {
+
+namespace detail {
+/// Concrete base class for CRTP TransformDialectDataBase. Must not be used
+/// directly.
+class TransformDialectDataBase {
+public:
+  virtual ~TransformDialectDataBase() = default;
+
+  /// Returns the dynamic type ID of the subclass.
+  TypeID getTypeID() const { return typeID; }
+
+protected:
+  /// Must be called by the subclass with the appropriate type ID.
+  explicit TransformDialectDataBase(TypeID typeID) : typeID(typeID) {}
+
+private:
+  /// The type ID of the subclass.
+  const TypeID typeID;
+};
+} // namespace detail
+
+/// Base class for additional data owned by the Transform dialect. Extensions
+/// may communicate with each other using this data. The data object is
+/// identified by the TypeID of the specific data subclass, querying the data of
+/// the same subclass returns a reference to the same object. When a Transform
+/// dialect extension is initialized, it can populate the data in the specific
+/// subclass. When a Transform op is applied, it can read (but not mutate) the
+/// data in the specific subclass, including the data provided by other
+/// extensions.
+///
+/// This follows CRTP: derived classes must list themselves as template
+/// argument.
+template <typename DerivedTy>
+class TransformDialectData : public detail::TransformDialectDataBase {
+protected:
+  /// Forward the TypeID of the derived class to the base.
+  TransformDialectData() : TransformDialectDataBase(TypeID::get<DerivedTy>()) {}
+};
+
 #ifndef NDEBUG
 namespace detail {
 /// Asserts that the operations provided as template arguments implement the
 /// TransformOpInterface and MemoryEffectsOpInterface. This must be a dynamic
 /// assertion since interface implementations may be registered at runtime.
 void checkImplementsTransformOpInterface(StringRef name, MLIRContext *context);
 
 /// Asserts that the type provided as template argument implements the
for (const DialectLoader &loader : dialectLoaders)
       loader(context);
 
     // Only load generated dialects if the user intends to apply
     // transformations specified by the extension.
     if (!buildOnly)
       for (const DialectLoader &loader : generatedDialectLoaders)
         loader(context);
 
-    for (const Initializer &init : opInitializers)
+    for (const Initializer &init : initializers)
       init(transformDialect);
-    transformDialect->mergeInPDLMatchHooks(std::move(pdlMatchConstraintFns));
   }
 
 protected:
   using Base = TransformDialectExtension<DerivedTy, ExtraDialects...>;
 
   /// Extension constructor. The argument indicates whether to skip generated
   /// dialects when applying the extension.
   explicit TransformDialectExtension(bool buildOnly = false)
       : buildOnly(buildOnly) {
     static_cast<DerivedTy *>(this)->init();
   }
 
+  /// Registers a custom initialization step to be performed when the extension
+  /// is applied to the dialect while loading. This is discouraged in favor of
+  /// more specific calls `declareGeneratedDialect`, `addDialectDataInitializer`
+  /// etc. `Func` must be convertible to the `void (MLIRContext *)` form. It
+  /// will be called during the extension initialization and given the current
+  /// MLIR context. This may be used to attach additional interfaces that cannot
+  /// be attached elsewhere.
+  template <typename Func>
+  void addCustomInitializationStep(Func &&func) {
+    std::function<void(MLIRContext *)> initializer = func;
+    dialectLoaders.push_back(
+        [init = std::move(initializer)](MLIRContext *ctx) { init(ctx); });
+  }
+
+  /// Registers the given function as one of the initializers for the
+  /// dialect-owned data of the kind specified as template argument. The
+  /// function must be convertible to the `void (DataTy &)` form. It will be
+  /// called during the extension initialization and will be given a mutable
+  /// reference to `DataTy`. The callback is expected to append data to the
+  /// given storage, and is not allowed to remove or destructively mutate the
+  /// existing data. The order in which callbacks from different extensions are
+  /// executed is unspecified so the callbacks may not rely on data being
+  /// already present. `DataTy` must be a class deriving `TransformDialectData`.
+  template <typename DataTy, typename Func>
+  void addDialectDataInitializer(Func &&func) {
+    static_assert(std::is_base_of_v<detail::TransformDialectDataBase, DataTy>,
+                  "only classes deriving TransformDialectData are accepted");
+
+    std::function<void(DataTy &)> initializer = func;
+    initializers.push_back(
+        [init = std::move(initializer)](TransformDialect *transformDialect) {
+          init(transformDialect->getOrCreateExtraData<DataTy>());
+        });
+  }
+
   /// Hook for derived classes to inject constructor behavior.
   void init() {}
 
   /// Injects the operations into the Transform dialect. The operations must
   /// implement the TransformOpInterface and MemoryEffectsOpInterface, and the
   /// implementations must be already available when the operation is injected.
   template <typename... OpTys>
   void registerTransformOps() {
-    opInitializers.push_back([](TransformDialect *transformDialect) {
+    initializers.push_back([](TransformDialect *transformDialect) {
       transformDialect->addOperationsChecked<OpTys...>();
     });
   }
 
   /// Injects the types into the Transform dialect. The types must implement
   /// the TransformHandleTypeInterface and the implementation must be already
   /// available when the type is injected. Furthermore, the types must provide
   /// a `getMnemonic` static method returning an object convertible to
   /// `StringRef` that is unique across all injected types.
   template <typename... TypeTys>
   void registerTypes() {
-    opInitializers.push_back([](TransformDialect *transformDialect) {
+    initializers.push_back([](TransformDialect *transformDialect) {
       transformDialect->addTypesChecked<TypeTys...>();
     });
   }
 
   /// Declares that this Transform dialect extension depends on the dialect
   /// provided as template parameter. When the Transform dialect is loaded,
   /// dependent dialects will be loaded as well. This is intended for dialects
   /// that contain attributes and types used in creation and canonicalization of
protected:
   /// dependent dialects of a pass. These dialects will be loaded along with the
   /// transform dialect unless the extension is in the build-only mode.
   template <typename DialectTy>
   void declareGeneratedDialect() {
     generatedDialectLoaders.push_back(
         [](MLIRContext *context) { context->loadDialect<DialectTy>(); });
   }
 
-  /// Injects the named constraint to make it available for use with the
-  /// PDLMatchOp in the transform dialect.
-  void registerPDLMatchConstraintFn(StringRef name,
-                                    PDLConstraintFunction &&fn) {
-    pdlMatchConstraintFns.try_emplace(name,
-                                      std::forward<PDLConstraintFunction>(fn));
-  }
-  template <typename ConstraintFnTy>
-  void registerPDLMatchConstraintFn(StringRef name, ConstraintFnTy &&fn) {
-    pdlMatchConstraintFns.try_emplace(
-        name, ::mlir::detail::pdl_function_builder::buildConstraintFn(
-                  std::forward<ConstraintFnTy>(fn)));
-  }
-
 private:
-  SmallVector<Initializer> opInitializers;
+  /// Callbacks performing extension initialization, e.g., registering ops,
+  /// types and defining the additional data.
+  SmallVector<Initializer> initializers;
 
   /// Callbacks loading the dependent dialects, i.e. the dialect needed for the
   /// extension ops.
   SmallVector<DialectLoader> dialectLoaders;
 
   /// Callbacks loading the generated dialects, i.e. the dialects produced when
   /// applying the transformations.
   SmallVector<DialectLoader> generatedDialectLoaders;
 
-  /// A list of constraints that should be made available to PDL patterns
-  /// processed by PDLMatchOp in the Transform dialect.
-  ///
-  /// Declared as mutable so its contents can be moved in the `apply` const
-  /// method, which is only called once.
-  mutable llvm::StringMap<PDLConstraintFunction> pdlMatchConstraintFns;
-
   /// Indicates that the extension is in build-only mode.
   bool buildOnly;
 };
 
 template <typename OpTy>
 void TransformDialect::addOperationIfNotRegistered() {
   StringRef name = OpTy::getOperationName();
   std::optional<RegisteredOperationName> opName =
void TransformDialect::addTypeIfNotRegistered() {
   addTypes<Type>();
 
 #ifndef NDEBUG
   detail::checkImplementsTransformHandleTypeInterface(TypeID::get<Type>(),
                                                       getContext());
 #endif // NDEBUG
 }
 
+template <typename DataTy>
+DataTy &TransformDialect::getOrCreateExtraData() {
+  TypeID typeID = TypeID::get<DataTy>();
+  auto it = extraData.find(typeID);
+  if (it != extraData.end())
+    return static_cast<DataTy &>(*it->getSecond());
+
+  auto emplaced = extraData.try_emplace(typeID, std::make_unique<DataTy>());
+  return static_cast<DataTy &>(*emplaced.first->getSecond());
+}
+
 /// A wrapper for transform dialect extensions that forces them to be
 /// constructed in the build-only mode.
 template <typename DerivedTy>
 class BuildOnly : public DerivedTy {
 public:
   BuildOnly() : DerivedTy(/*buildOnly=*/true) {}
 };
 
 } // namespace transform
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

 
 def Transform_Dialect : Dialect {
   let summary = "Fine-grain transformation control dialect";
   // For description, see docs/Dialects/Transform.md.
 
   let name = "transform";
   let cppNamespace = "::mlir::transform";
 
-  let dependentDialects = [
-    "::mlir::pdl::PDLDialect",
-    "::mlir::pdl_interp::PDLInterpDialect",
-  ];
-
   let hasOperationAttrVerify = 1;
   let usePropertiesForAttributes = 1;
 
   let extraClassDeclaration = [{
       /// Name of the attribute attachable to the symbol table operation
       /// containing named sequences. This is used to trigger verification.
-      constexpr const static llvm::StringLiteral
+      constexpr const static ::llvm::StringLiteral
           kWithNamedSequenceAttrName = "transform.with_named_sequence";
 
       /// Names of the attribute attachable to an operation so it can be
       /// identified as root by the default interpreter pass.
-      constexpr const static llvm::StringLiteral
+      constexpr const static ::llvm::StringLiteral
           kTargetTagAttrName = "transform.target_tag";
 
       /// Names of the attributes indicating whether an argument of an external
       /// transform dialect symbol is consumed or only read.
-      constexpr const static llvm::StringLiteral
+      constexpr const static ::llvm::StringLiteral
           kArgConsumedAttrName = "transform.consumed";
-      constexpr const static llvm::StringLiteral
+      constexpr const static ::llvm::StringLiteral
           kArgReadOnlyAttrName = "transform.readonly";
 
-      /// Returns the named PDL constraint functions available in the dialect
-      /// as a map from their name to the function.
-      const ::llvm::StringMap<::mlir::PDLConstraintFunction> &
-      getPDLConstraintHooks() const;
+      template <typename DataTy>
+      const DataTy &getExtraData() const {
+        return *static_cast<const DataTy *>(extraData.at(::mlir::TypeID::get<DataTy>()).get());
+      }
 
       /// Parses a type registered by this dialect or one of its extensions.
       ::mlir::Type parseType(::mlir::DialectAsmParser &parser) const override;
 
       /// Prints a type registered by this dialect or one of its extensions.
       void printType(::mlir::Type type,
                      ::mlir::DialectAsmPrinter &printer) const override;
 
private:
       }
       template <typename Type>
       void addTypeIfNotRegistered();
 
       /// Reports a repeated registration error of a type with the given
       /// mnemonic.
       [[noreturn]] void reportDuplicateTypeRegistration(StringRef mnemonic);
 
+      /// Registers dialect types with the context.
       void initializeTypes();
 
+      // Give extensions access to injection functions.
       template <typename, typename...>
       friend class TransformDialectExtension;
 
-      /// Takes ownership of the named PDL constraint function from the given
-      /// map and makes them available for use by the operations in the dialect.
-      void mergeInPDLMatchHooks(
-          ::llvm::StringMap<::mlir::PDLConstraintFunction> &&constraintFns);
+      /// Gets a mutable reference to extra data of the kind specified as
+      /// template argument. Allocates the data on the first call.
+      template <typename DataTy>
+      DataTy &getOrCreateExtraData();
 
       //===----------------------------------------------------------------===//
       // Data fields
       //===----------------------------------------------------------------===//
 
-      /// A container for PDL constraint function that can be used by
-      /// operations in this dialect.
-      ::mlir::PDLPatternModule pdlMatchHooks;
+      /// Additional data associated with and owned by the dialect. Accessible
+      /// to extensions.
+      ::llvm::DenseMap<::mlir::TypeID, std::unique_ptr<
+            ::mlir::transform::detail::TransformDialectDataBase>>
+          extraData;
 
       /// A map from type mnemonic to its parsing function for the remainder of
       /// the syntax. The parser has access to the mnemonic, so it is used for
       /// further dispatch.
       ::llvm::StringMap<ExtensionTypeParsingHook> typeParsingHooks;
 
       /// A map from type TypeID to its printing function. No need to do string
       /// lookups when the type is fully constructed.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

 /// the payload IR, depending on what is available in the context.
 LogicalResult
 mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
                                              Operation *op, Region &region);
 
 /// Verification hook for PossibleTopLevelTransformOpTrait.
 LogicalResult verifyPossibleTopLevelTransformOpTrait(Operation *op);
 
+/// Populates `effects` with side effects implied by
+/// PossibleTopLevelTransformOpTrait for the given operation. The operation may
+/// have an optional `root` operand, indicating it is not in fact top-level. It
+/// is also expected to have a single-block body.
+void getPotentialTopLevelEffects(
+    Operation *operation, Value root, Block &body,
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+
 /// Verification hook for TransformOpInterface.
 LogicalResult verifyTransformOpInterface(Operation *op);
 
 /// Populates `mappings` with mapped values associated with the given transform
 /// IR values in the given `state`.
 void prepareValueMappings(
     SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings,
     ValueRange values, const transform::TransformState &state);

 TransformState::make_isolated_region_scope(Region &region) {
   return RegionScope(*this, region, RegionScope::Isolated());
 }
 
 /// This trait is supposed to be attached to Transform dialect operations that
 /// can be standalone top-level transforms. Such operations typically contain
 /// other Transform dialect operations that can be executed following some
 /// control flow logic specific to the current operation. The operations with
-/// this trait are expected to have at least one single-block region with one
-/// argument of PDL Operation type. The operations are also expected to be valid
-/// without operands, in which case they are considered top-level, and with one
-/// or more arguments, in which case they are considered nested. Top-level
-/// operations have the block argument of the entry block in the Transform IR
-/// correspond to the root operation of Payload IR. Nested operations have the
-/// block argument of the entry block in the Transform IR correspond to a list
-/// of Payload IR operations mapped to the first operand of the Transform IR
-/// operation. The operation must implement TransformOpInterface.
+/// this trait are expected to have at least one single-block region with at
+/// least one argument of type implementing TransformHandleTypeInterface. The
+/// operations are also expected to be valid without operands, in which case
+/// they are considered top-level, and with one or more arguments, in which case
+/// they are considered nested. Top-level operations have the block argument of
+/// the entry block in the Transform IR correspond to the root operation of
+/// Payload IR. Nested operations have the block argument of the entry block in
+/// the Transform IR correspond to a list of Payload IR operations mapped to the
+/// first operand of the Transform IR operation. The operation must implement
+/// TransformOpInterface.
 template <typename OpTy>
 class PossibleTopLevelTransformOpTrait
     : public OpTrait::TraitBase<OpTy, PossibleTopLevelTransformOpTrait> {
 public:
   /// Verifies that `op` satisfies the invariants of this trait. Not expected to
   /// be called directly.
   static LogicalResult verifyTrait(Operation *op) {
     return detail::verifyPossibleTopLevelTransformOpTrait(op);
   }
 
   /// Returns the single block of the given region.
   Block *getBodyBlock(unsigned region = 0) {
     return &this->getOperation()->getRegion(region).front();
   }
 
+  /// Populates `effects` with side effects implied by this trait.
+  void getPotentialTopLevelEffects(
+      SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+    detail::getPotentialTopLevelEffects(
+        this->getOperation(), cast<OpTy>(this->getOperation()).getRoot(),
+        *getBodyBlock(), effects);
+  }
+
   /// Sets up the mapping between the entry block of the given region of this op
   /// and the relevant list of Payload IR operations in the given state. The
   /// state is expected to be already scoped at the region of this operation.
   LogicalResult mapBlockArguments(TransformState &state, Region &region) {
     assert(region.getParentOp() == this->getOperation() &&
            "op comes from the wrong region");
     return detail::mapPossibleTopLevelTransformOpBlockArguments(
         state, this->getOperation(), region);

mlir/include/mlir/Dialect/Transform/IR/TransformOps.h

 //===- TransformDialect.h - Transform dialect operations --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
 
-#include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformAttrs.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let description = [{
     but the general association check may fail if the parameter type does not
     accept the given kind of attribute as valid.
   }];
   let arguments = (ins AnyAttr:$value);
   let results = (outs TransformParamTypeInterface:$param);
   let assemblyFormat = "$value attr-dict `->` type($param)";
 }
 
-def PDLMatchOp : TransformDialectOp<"pdl_match",
-    [DeclareOpInterfaceMethods<TransformOpInterface>,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let summary = "Finds ops that match the named PDL pattern";
-  let description = [{
-    Find Payload IR ops nested within the Payload IR op associated with the
-    operand that match the PDL pattern identified by its name. The pattern is
-    expected to be defined in the closest surrounding `WithPDLPatternsOp`.
-
-    Produces a Transform IR value associated with the list of Payload IR ops
-    that matched the pattern. The order of results in the list is that of the
-    Operation::walk, clients are advised not to rely on a specific order though.
-    If the operand is associated with multiple Payload IR ops, finds matching
-    ops nested within each of those and produces a single list containing all
-    of the matched ops.
-
-    The transformation is considered successful regardless of whether some
-    Payload IR ops actually matched the pattern and only fails if the pattern
-    could not be looked up or compiled.
-  }];
-
-  let arguments = (ins
-    Arg<TransformHandleTypeInterface, "Payload IR scope to match within">:$root,
-    SymbolRefAttr:$pattern_name);
-  let results = (outs
-    Res<TransformHandleTypeInterface, "Handle to the matched Payload IR ops">:$matched);
-
-  let assemblyFormat = "$pattern_name `in` $root attr-dict `:` "
-                       "functional-type(operands, results)";
-}
-
 def PrintOp : TransformDialectOp<"print",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Dump each payload op";
   let description = [{
     This op dumps each payload op that is associated with the `target` operand
     to stderr. It also prints the `name` string attribute. If no target is
     specified, the top-level op is dumped.
def SequenceOp : TransformDialectOp<"sequence",
   let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "transform"; }
   }];
 
   let hasVerifier = 1;
 }
 
-def WithPDLPatternsOp : TransformDialectOp<"with_pdl_patterns",
-    [DeclareOpInterfaceMethods<TransformOpInterface>, NoTerminator,
-     OpAsmOpInterface, PossibleTopLevelTransformOpTrait,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     SymbolTable]> {
-  let summary = "Contains PDL patterns available for use in transforms";
-  let description = [{
-    This op contains a set of named PDL patterns that are available for the
-    Transform dialect operations to be used for pattern matching. For example,
-    PDLMatchOp can be used to produce a Transform IR value associated with all
-    Payload IR operations that match the pattern as follows:
-
-    ```mlir
-    transform.with_pdl_patterns {
-    ^bb0(%arg0: !transform.any_op):
-      pdl.pattern @my_pattern : benefit(1) {
-        %0 = pdl.operation //...
-        // Regular PDL goes here.
-        pdl.rewrite %0 with "transform.dialect"
-      }
-
-      sequence %arg0 failures(propagate) {
-      ^bb0(%arg1: !transform.any_op):
-        %1 = pdl_match @my_pattern in %arg1
-        // Use %1 as handle
-      }
-    }
-    ```
-
-    Note that the pattern is expected to finish with a `pdl.rewrite` terminator
-    that points to the custom rewriter named "transform.dialect". The rewriter
-    actually does nothing, but the transform application will keep track of the
-    operations that matched the pattern.
-
-    This op is expected to contain `pdl.pattern` operations and exactly one
-    another Transform dialect operation that gets executed with all patterns
-    available. This op is a possible top-level Transform IR op, the argument of
-    its entry block corresponds to either the root op of the payload IR or the
-    ops associated with its operand when provided.
-  }];
-
-  let arguments = (ins
-    Arg<Optional<TransformHandleTypeInterface>, "Root operation of the Payload IR"
-        >:$root);
-  let regions = (region SizedRegion<1>:$body);
-  let assemblyFormat = "($root^ `:` type($root))? attr-dict-with-keyword regions";
-
-  let hasVerifier = 1;
-
-  let extraClassDeclaration = [{
-    /// Allow the dialect prefix to be omitted.
-    static StringRef getDefaultDialect() { return "transform"; }
-  }];
-}
-
 def YieldOp : TransformDialectOp<"yield",
     [Terminator, DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Yields operation handles from a transform IR region";
   let description = [{
     This terminator operation yields operation handles from regions of the
     transform IR ops back to the containing op. It is not itself associated with
     any transformation on the payload IR and is used for flow purposes only.
   }];

mlir/include/mlir/Dialect/Transform/PDLExtension/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS PDLExtensionOps.td)
+mlir_tablegen(PDLExtensionOps.h.inc -gen-op-decls)
+mlir_tablegen(PDLExtensionOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRTransformDialectPDLExtensionOpsIncGen)
+
+add_mlir_doc(PDLExtensionOps PDLExtensionOps Dialects/ -gen-op-doc)

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtension.h

+//===- PDLExtension.h - PDL extension for Transform dialect -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+class DialectRegistry;
+
+namespace transform {
+/// Registers the PDL extension of the Transform dialect in the given registry.
+void registerPDLExtension(DialectRegistry &dialectRegistry);
+} // namespace transform
+} // namespace mlir

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h

+//===- PDLExtensionOps.h - PDL extension for Transform dialect --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS_H
+#define MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h.inc"
+
+namespace mlir {
+namespace transform {
+/// PDL constraint callbacks that can be used by the PDL extension of the
+/// Transform dialect. These are owned by the Transform dialect and can be
+/// populated by extensions.
+class PDLMatchHooks : public TransformDialectData<PDLMatchHooks> {
+public:
+  /// Takes ownership of the named PDL constraint function from the given
+  /// map and makes them available for use by the operations in the dialect.
+  void
+  mergeInPDLMatchHooks(llvm::StringMap<PDLConstraintFunction> &&constraintFns);
+
+  /// Returns the named PDL constraint functions available in the dialect
+  /// as a map from their name to the function.
+  const llvm::StringMap<::mlir::PDLConstraintFunction> &
+  getPDLConstraintHooks() const;
+
+private:
+  /// A container for PDL constraint function that can be used by
+  /// operations in this dialect.
+  PDLPatternModule pdlMatchHooks;
+};
+} // namespace transform
+} // namespace mlir
+
+MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::transform::PDLMatchHooks)
+
+#endif // MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS_H

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td

+//===- TransformOps.td - Transform dialect operations ------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS
+#define MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/SymbolInterfaces.td"
+
+def PDLMatchOp : TransformDialectOp<"pdl_match",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Finds ops that match the named PDL pattern";
+  let description = [{
+    Find Payload IR ops nested within the Payload IR op associated with the
+    operand that match the PDL pattern identified by its name. The pattern is
+    expected to be defined in the closest surrounding `WithPDLPatternsOp`.
+
+    Produces a Transform IR value associated with the list of Payload IR ops
+    that matched the pattern. The order of results in the list is that of the
+    Operation::walk, clients are advised not to rely on a specific order though.
+    If the operand is associated with multiple Payload IR ops, finds matching
+    ops nested within each of those and produces a single list containing all
+    of the matched ops.
+
+    The transformation is considered successful regardless of whether some
+    Payload IR ops actually matched the pattern and only fails if the pattern
+    could not be looked up or compiled.
+  }];
+
+  let arguments = (ins
+    Arg<TransformHandleTypeInterface, "Payload IR scope to match within">:$root,
+    SymbolRefAttr:$pattern_name);
+  let results = (outs
+    Res<TransformHandleTypeInterface, "Handle to the matched Payload IR ops">:$matched);
+
+  let assemblyFormat = "$pattern_name `in` $root attr-dict `:` "
+                       "functional-type(operands, results)";
+}
+
+def WithPDLPatternsOp : TransformDialectOp<"with_pdl_patterns",
+    [DeclareOpInterfaceMethods<TransformOpInterface>, NoTerminator,
+     OpAsmOpInterface, PossibleTopLevelTransformOpTrait,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     SymbolTable]> {
+  let summary = "Contains PDL patterns available for use in transforms";
+  let description = [{
+    This op contains a set of named PDL patterns that are available for the
+    Transform dialect operations to be used for pattern matching. For example,
+    PDLMatchOp can be used to produce a Transform IR value associated with all
+    Payload IR operations that match the pattern as follows:
+
+    ```mlir
+    transform.with_pdl_patterns {
+    ^bb0(%arg0: !transform.any_op):
+      pdl.pattern @my_pattern : benefit(1) {
+        %0 = pdl.operation //...
+        // Regular PDL goes here.
+        pdl.rewrite %0 with "transform.dialect"
+      }
+
+      sequence %arg0 failures(propagate) {
+      ^bb0(%arg1: !transform.any_op):
+        %1 = pdl_match @my_pattern in %arg1
+        // Use %1 as handle
+      }
+    }
+    ```
+
+    Note that the pattern is expected to finish with a `pdl.rewrite` terminator
+    that points to the custom rewriter named "transform.dialect". The rewriter
+    actually does nothing, but the transform application will keep track of the
+    operations that matched the pattern.
+
+    This op is expected to contain `pdl.pattern` operations and exactly one
+    another Transform dialect operation that gets executed with all patterns
+    available. This op is a possible top-level Transform IR op, the argument of
+    its entry block corresponds to either the root op of the payload IR or the
+    ops associated with its operand when provided.
+  }];
+
+  let arguments = (ins
+    Arg<Optional<TransformHandleTypeInterface>, "Root operation of the Payload IR"
+        >:$root);
+  let regions = (region SizedRegion<1>:$body);
+  let assemblyFormat = "($root^ `:` type($root))? attr-dict-with-keyword regions";
+
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    /// Allow the dialect prefix to be omitted.
+    static StringRef getDefaultDialect() { return "transform"; }
+  }];
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_PDLEXTENSION_PDLEXTENSIONOPS

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
inline void registerAllDialects(DialectRegistry &registry) {
   // Register all dialect extensions.
   affine::registerTransformDialectExtension(registry);
   bufferization::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
   tensor::registerTransformDialectExtension(registry);
+  transform::registerPDLExtension(registry);
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);

mlir/lib/Dialect/Transform/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(PDLExtension)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRTransformDialect
   MatchInterfaces.cpp
   TransformAttrs.cpp
   TransformDialect.cpp
   TransformInterfaces.cpp
   TransformOps.cpp
   TransformTypes.cpp
 
   DEPENDS
   MLIRMatchInterfacesIncGen
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRParser
-  MLIRPDLDialect
-  MLIRPDLInterpDialect
   MLIRRewrite
   MLIRSideEffectInterfaces
   MLIRTransforms
   MLIRTransformDialectUtils
   )

mlir/lib/Dialect/Transform/IR/TransformDialect.cpp

 //===- TransformDialect.cpp - Transform Dialect Definition ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Analysis/CallGraph.h"
-#include "mlir/Dialect/PDL/IR/PDL.h"
-#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/SCCIterator.h"
 
 using namespace mlir;
 
assert((abstractType.hasInterface(
               TransformParamTypeInterface::getInterfaceID()) ||
           abstractType.hasInterface(
               TransformValueHandleTypeInterface::getInterfaceID())) &&
          "expected Transform dialect type to implement one of the three "
          "interfaces");
 }
 #endif // NDEBUG
 
-namespace {
-struct PDLOperationTypeTransformHandleTypeInterfaceImpl
-    : public transform::TransformHandleTypeInterface::ExternalModel<
-          PDLOperationTypeTransformHandleTypeInterfaceImpl,
-          pdl::OperationType> {
-  DiagnosedSilenceableFailure
-  checkPayload(Type type, Location loc, ArrayRef<Operation *> payload) const {
-    return DiagnosedSilenceableFailure::success();
-  }
-};
-} // namespace
-
 void transform::TransformDialect::initialize() {
   // Using the checked versions to enable the same assertions as for the ops
   // from extensions.
   addOperationsChecked<
 #define GET_OP_LIST
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
       >();
   initializeTypes();
-
-  pdl::OperationType::attachInterface<
-      PDLOperationTypeTransformHandleTypeInterfaceImpl>(*getContext());
-}
-
-void transform::TransformDialect::mergeInPDLMatchHooks(
-    llvm::StringMap<PDLConstraintFunction> &&constraintFns) {
-  // Steal the constraint functions from the given map.
-  for (auto &it : constraintFns)
-    pdlMatchHooks.registerConstraintFunction(it.getKey(), std::move(it.second));
-}
-
-const llvm::StringMap<PDLConstraintFunction> &
-transform::TransformDialect::getPDLConstraintHooks() const {
-  return pdlMatchHooks.getConstraintFunctions();
 }
 
 Type transform::TransformDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   SMLoc loc = parser.getCurrentLocation();
   if (failed(parser.parseKeyword(&keyword)))
     return nullptr;
 

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

for (auto &&[terminatorOperand, result] :
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 
+/// Appends to `effects` the memory effect instances on `target` with the same
+/// resource and effect as the ones the operation `iface` having on `source`.
+static void
+remapEffects(MemoryEffectOpInterface iface, BlockArgument source, Value target,
+             SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  SmallVector<MemoryEffects::EffectInstance> nestedEffects;
+  iface.getEffectsOnValue(source, nestedEffects);
+  for (const auto &effect : nestedEffects)
+    effects.emplace_back(effect.getEffect(), target, effect.getResource());
+}
+
+/// Appends to `effects` the same effects as the operations of `block` have on
+/// block arguments but associated with `operands.`
+static void
+remapArgumentEffects(Block &block, ValueRange operands,
+                     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  for (Operation &op : block) {
+    auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
+    if (!iface)
+      continue;
+
+    for (auto &&[source, target] : llvm::zip(block.getArguments(), operands)) {
+      remapEffects(iface, source, target, effects);
+    }
+
+    SmallVector<MemoryEffects::EffectInstance> nestedEffects;
+    iface.getEffectsOnResource(transform::PayloadIRResource::get(),
+                               nestedEffects);
+    llvm::append_range(effects, nestedEffects);
+  }
+}
+
+void transform::detail::getPotentialTopLevelEffects(
+    Operation *operation, Value root, Block &body,
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(operation->getOperands(), effects);
+  transform::producesHandle(operation->getResults(), effects);
+
+  if (!root) {
+    for (Operation &op : body) {
+      auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
+      if (!iface)
+        continue;
+
+      SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
+      iface.getEffects(effects);
+    }
+    return;
+  }
+
+  // Carry over all effects on arguments of the entry block as those on the
+  // operands, this is the same value just remapped.
+  remapArgumentEffects(body, operation->getOperands(), effects);
+}
+
 LogicalResult transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
     TransformState &state, Operation *op, Region &region) {
   SmallVector<Operation *> targets;
   SmallVector<SmallVector<MappedValue>> extraMappings;
   if (op->getNumOperands() != 0) {
     llvm::append_range(targets, state.getPayloadOps(op->getOperand(0)));
     prepareValueMappings(extraMappings, op->getOperands().drop_front(), state);
   } else {

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 //===- TransformDialect.cpp - Transform dialect operations ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
-#include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformAttrs.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
-#include "mlir/Rewrite/FrozenRewritePatternSet.h"
-#include "mlir/Rewrite/PatternApplicator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "transform-dialect"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")

 static ParseResult parseForeachMatchSymbols(OpAsmParser &parser,
                                             ArrayAttr &matchers,
                                             ArrayAttr &actions);
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
-// PatternApplicatorExtension
-//===----------------------------------------------------------------------===//
-
-namespace {
-/// A TransformState extension that keeps track of compiled PDL pattern sets.
-/// This is intended to be used along the WithPDLPatterns op. The extension
-/// can be constructed given an operation that has a SymbolTable trait and
-/// contains pdl::PatternOp instances. The patterns are compiled lazily and one
-/// by one when requested; this behavior is subject to change.
-class PatternApplicatorExtension : public transform::TransformState::Extension {
-public:
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PatternApplicatorExtension)
-
-  /// Creates the extension for patterns contained in `patternContainer`.
-  explicit PatternApplicatorExtension(transform::TransformState &state,
-                                      Operation *patternContainer)
-      : Extension(state), patterns(patternContainer) {}
-
-  /// Appends to `results` the operations contained in `root` that matched the
-  /// PDL pattern with the given name. Note that `root` may or may not be the
-  /// operation that contains PDL patterns. Reports an error if the pattern
-  /// cannot be found. Note that when no operations are matched, this still
-  /// succeeds as long as the pattern exists.
-  LogicalResult findAllMatches(StringRef patternName, Operation *root,
-                               SmallVectorImpl<Operation *> &results);
-
-private:
-  /// Map from the pattern name to a singleton set of rewrite patterns that only
-  /// contains the pattern with this name. Populated when the pattern is first
-  /// requested.
-  // TODO: reconsider the efficiency of this storage when more usage data is
-  // available. Storing individual patterns in a set and triggering compilation
-  // for each of them has overhead. So does compiling a large set of patterns
-  // only to apply a handlful of them.
-  llvm::StringMap<FrozenRewritePatternSet> compiledPatterns;
-
-  /// A symbol table operation containing the relevant PDL patterns.
-  SymbolTable patterns;
-};
-
-LogicalResult PatternApplicatorExtension::findAllMatches(
-    StringRef patternName, Operation *root,
-    SmallVectorImpl<Operation *> &results) {
-  auto it = compiledPatterns.find(patternName);
-  if (it == compiledPatterns.end()) {
-    auto patternOp = patterns.lookup<pdl::PatternOp>(patternName);
-    if (!patternOp)
-      return failure();
-
-    // Copy the pattern operation into a new module that is compiled and
-    // consumed by the PDL interpreter.
-    OwningOpRef<ModuleOp> pdlModuleOp = ModuleOp::create(patternOp.getLoc());
-    auto builder = OpBuilder::atBlockEnd(pdlModuleOp->getBody());
-    builder.clone(*patternOp);
-    PDLPatternModule patternModule(std::move(pdlModuleOp));
-
-    // Merge in the hooks owned by the dialect. Make a copy as they may be
-    // also used by the following operations.
-    auto *dialect =
-        root->getContext()->getLoadedDialect<transform::TransformDialect>();
-    for (const auto &[name, constraintFn] : dialect->getPDLConstraintHooks())
-      patternModule.registerConstraintFunction(name, constraintFn);
-
-    // Register a noop rewriter because PDL requires patterns to end with some
-    // rewrite call.
-    patternModule.registerRewriteFunction(
-        "transform.dialect", [](PatternRewriter &, Operation *) {});
-
-    it = compiledPatterns
-             .try_emplace(patternOp.getName(), std::move(patternModule))
-             .first;
-  }
-
-  PatternApplicator applicator(it->second);
-  // We want to discourage direct use of PatternRewriter in APIs but In this
-  // very specific case, an IRRewriter is not enough.
-  struct TrivialPatternRewriter : public PatternRewriter {
-  public:
-    explicit TrivialPatternRewriter(MLIRContext *context)
-        : PatternRewriter(context) {}
-  };
-  TrivialPatternRewriter rewriter(root->getContext());
-  applicator.applyDefaultCostModel();
-  root->walk([&](Operation *op) {
-    if (succeeded(applicator.matchAndRewrite(op, rewriter)))
-      results.push_back(op);
-  });
-
-  return success();
-}
-} // namespace
-
-//===----------------------------------------------------------------------===//
 // TrackingListener
 //===----------------------------------------------------------------------===//
 
 Operation *transform::TrackingListener::getCommonDefiningOp(ValueRange values) {
   Operation *defOp = nullptr;
   for (Value v : values) {
     // Skip empty values.
     if (!v)
void transform::CastOp::getEffects(
   producesHandle(getOutput(), effects);
 }
 
 bool transform::CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   assert(inputs.size() == 1 && "expected one input");
   assert(outputs.size() == 1 && "expected one output");
   return llvm::all_of(
       std::initializer_list<Type>{inputs.front(), outputs.front()},
-      [](Type ty) {
-        return llvm::isa<pdl::OperationType,
-                         transform::TransformHandleTypeInterface>(ty);
-      });
+      [](Type ty) { return isa<transform::TransformHandleTypeInterface>(ty); });
 }
 
 //===----------------------------------------------------------------------===//
 // ForeachMatchOp
 //===----------------------------------------------------------------------===//
 
 /// Applies matcher operations from the given `block` assigning `op` as the
 /// payload of the block's first argument. Updates `state` accordingly. If any
transform::IncludeOp::apply(transform::TransformResults &results,
   mappings.clear();
   detail::prepareValueMappings(
       mappings, callee.getBody().front().getTerminator()->getOperands(), state);
   for (auto &&[result, mapping] : llvm::zip_equal(getResults(), mappings))
     results.setMappedValues(result, mapping);
   return result;
 }
 
-/// Appends to `effects` the memory effect instances on `target` with the same
-/// resource and effect as the ones the operation `iface` having on `source`.
-static void
-remapEffects(MemoryEffectOpInterface iface, BlockArgument source, Value target,
-             SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  SmallVector<MemoryEffects::EffectInstance> nestedEffects;
-  iface.getEffectsOnValue(source, nestedEffects);
-  for (const auto &effect : nestedEffects)
-    effects.emplace_back(effect.getEffect(), target, effect.getResource());
-}
-
-/// Appends to `effects` the same effects as the operations of `block` have on
-/// block arguments but associated with `operands.`
-static void
-remapArgumentEffects(Block &block, ValueRange operands,
-                     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  for (Operation &op : block) {
-    auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
-    if (!iface)
-      continue;
-
-    for (auto &&[source, target] : llvm::zip(block.getArguments(), operands)) {
-      remapEffects(iface, source, target, effects);
-    }
-
-    SmallVector<MemoryEffects::EffectInstance> nestedEffects;
-    iface.getEffectsOnResource(transform::PayloadIRResource::get(),
-                               nestedEffects);
-    llvm::append_range(effects, nestedEffects);
-  }
-}
-
 static DiagnosedSilenceableFailure
 verifyNamedSequenceOp(transform::NamedSequenceOp op);
 
 void transform::IncludeOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   // Always mark as modifying the payload.
   // TODO: a mechanism to annotate effects on payload. Even when all handles are
   // only read, the payload may still be modified, so we currently stay on the

 
 //===----------------------------------------------------------------------===//
 // SplitHandleOp
 //===----------------------------------------------------------------------===//
 
 void transform::SplitHandleOp::build(OpBuilder &builder, OperationState &result,
                                      Value target, int64_t numResultHandles) {
   result.addOperands(target);
-  auto pdlOpType = pdl::OperationType::get(builder.getContext());
-  result.addTypes(SmallVector<pdl::OperationType>(numResultHandles, pdlOpType));
+  result.addTypes(SmallVector<Type>(numResultHandles, target.getType()));
 }
 
 DiagnosedSilenceableFailure
 transform::SplitHandleOp::apply(transform::TransformResults &results,
                                 transform::TransformState &state) {
   int64_t numPayloadOps = llvm::range_size(state.getPayloadOps(getHandle()));
   auto produceNumOpsError = [&]() {
     return emitSilenceableError()

 LogicalResult transform::SplitHandleOp::verify() {
   if (getOverflowResult().has_value() &&
       !(*getOverflowResult() >= 0 && *getOverflowResult() < getNumResults()))
     return emitOpError("overflow_result is not a valid result index");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// PDLMatchOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::PDLMatchOp::apply(transform::TransformResults &results,
-                             transform::TransformState &state) {
-  auto *extension = state.getExtension<PatternApplicatorExtension>();
-  assert(extension &&
-         "expected PatternApplicatorExtension to be attached by the parent op");
-  SmallVector<Operation *> targets;
-  for (Operation *root : state.getPayloadOps(getRoot())) {
-    if (failed(extension->findAllMatches(
-            getPatternName().getLeafReference().getValue(), root, targets))) {
-      emitDefiniteFailure()
-          << "could not find pattern '" << getPatternName() << "'";
-    }
-  }
-  results.set(llvm::cast<OpResult>(getResult()), targets);
-  return DiagnosedSilenceableFailure::success();
-}
-
-void transform::PDLMatchOp::getEffects(
-    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  onlyReadsHandle(getRoot(), effects);
-  producesHandle(getMatched(), effects);
-  onlyReadsPayload(effects);
-}
-
-//===----------------------------------------------------------------------===//
 // ReplicateOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ReplicateOp::apply(transform::TransformResults &results,
                               transform::TransformState &state) {
   unsigned numRepetitions = llvm::range_size(state.getPayloadOps(getPattern()));
   for (const auto &en : llvm::enumerate(getHandles())) {
InFlightDiagnostic diag = emitOpError()
                               << "expects the types of the terminator operands "
                                  "to match the types of the result";
     diag.attachNote(getBodyBlock()->getTerminator()->getLoc()) << "terminator";
     return diag;
   }
   return success();
 }
 
-/// Populate `effects` with transform dialect memory effects for the potential
-/// top-level operation. Such operations have recursive effects from nested
-/// operations. When they have an operand, we can additionally remap effects on
-/// the block argument to be effects on the operand.
-template <typename OpTy>
-static void getPotentialTopLevelEffects(
-    OpTy operation, SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  transform::onlyReadsHandle(operation->getOperands(), effects);
-  transform::producesHandle(operation->getResults(), effects);
-
-  if (!operation.getRoot()) {
-    for (Operation &op : *operation.getBodyBlock()) {
-      auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
-      if (!iface)
-        continue;
-
-      SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
-      iface.getEffects(effects);
-    }
-    return;
-  }
-
-  // Carry over all effects on arguments of the entry block as those on the
-  // operands, this is the same value just remapped.
-  remapArgumentEffects(*operation.getBodyBlock(), operation->getOperands(),
-                       effects);
-}
-
 void transform::SequenceOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  getPotentialTopLevelEffects(*this, effects);
+  getPotentialTopLevelEffects(effects);
 }
 
 OperandRange transform::SequenceOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {
   assert(index && *index == 0 && "unexpected region index");
   if (getOperation()->getNumOperands() > 0)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),
void transform::SequenceOp::build(OpBuilder &builder, OperationState &state,
                                   Type bbArgType, TypeRange extraBindingTypes,
                                   SequenceBodyBuilderArgsFn bodyBuilder) {
   build(builder, state, resultTypes, failurePropagationMode, /*root=*/Value(),
         /*extra_bindings=*/ValueRange());
   buildSequenceBody(builder, state, bbArgType, extraBindingTypes, bodyBuilder);
 }
 
 //===----------------------------------------------------------------------===//
-// WithPDLPatternsOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
-                                    transform::TransformState &state) {
-  TransformOpInterface transformOp = nullptr;
-  for (Operation &nested : getBody().front()) {
-    if (!isa<pdl::PatternOp>(nested)) {
-      transformOp = cast<TransformOpInterface>(nested);
-      break;
-    }
-  }
-
-  state.addExtension<PatternApplicatorExtension>(getOperation());
-  auto guard = llvm::make_scope_exit(
-      [&]() { state.removeExtension<PatternApplicatorExtension>(); });
-
-  auto scope = state.make_region_scope(getBody());
-  if (failed(mapBlockArguments(state)))
-    return DiagnosedSilenceableFailure::definiteFailure();
-  return state.applyTransform(transformOp);
-}
-
-void transform::WithPDLPatternsOp::getEffects(
-    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  getPotentialTopLevelEffects(*this, effects);
-}
-
-LogicalResult transform::WithPDLPatternsOp::verify() {
-  Block *body = getBodyBlock();
-  Operation *topLevelOp = nullptr;
-  for (Operation &op : body->getOperations()) {
-    if (isa<pdl::PatternOp>(op))
-      continue;
-
-    if (op.hasTrait<::mlir::transform::PossibleTopLevelTransformOpTrait>()) {
-      if (topLevelOp) {
-        InFlightDiagnostic diag =
-            emitOpError() << "expects only one non-pattern op in its body";
-        diag.attachNote(topLevelOp->getLoc()) << "first non-pattern op";
-        diag.attachNote(op.getLoc()) << "second non-pattern op";
-        return diag;
-      }
-      topLevelOp = &op;
-      continue;
-    }
-
-    InFlightDiagnostic diag =
-        emitOpError()
-        << "expects only pattern and top-level transform ops in its body";
-    diag.attachNote(op.getLoc()) << "offending op";
-    return diag;
-  }
-
-  if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
-    InFlightDiagnostic diag = emitOpError() << "cannot be nested";
-    diag.attachNote(parent.getLoc()) << "parent operation";
-    return diag;
-  }
-
-  if (!topLevelOp) {
-    InFlightDiagnostic diag = emitOpError()
-                              << "expects at least one non-pattern op";
-    return diag;
-  }
-
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // PrintOp
 //===----------------------------------------------------------------------===//
 
 void transform::PrintOp::build(OpBuilder &builder, OperationState &result,
                                StringRef name) {
   if (!name.empty()) {
     result.addAttribute(PrintOp::getNameAttrName(result.name),
                         builder.getStrArrayAttr(name));

mlir/lib/Dialect/Transform/PDLExtension/CMakeLists.txt

+add_mlir_dialect_library(MLIRTransformPDLExtension
+  PDLExtension.cpp
+  PDLExtensionOps.cpp
+
+  DEPENDS
+  MLIRTransformDialectPDLExtensionOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRTransformDialect
+  MLIRPDLDialect
+  MLIRPDLInterpDialect
+)

mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp

+//===- PDLExtension.cpp - PDL extension for the Transform dialect ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h"
+#include "mlir/IR/DialectRegistry.h"
+
+using namespace mlir;
+
+namespace {
+/// Implementation of the TransformHandleTypeInterface for the PDL
+/// OperationType. Accepts any payload operation.
+struct PDLOperationTypeTransformHandleTypeInterfaceImpl
+    : public transform::TransformHandleTypeInterface::ExternalModel<
+          PDLOperationTypeTransformHandleTypeInterfaceImpl,
+          pdl::OperationType> {
+
+  /// Accept any operation.
+  DiagnosedSilenceableFailure
+  checkPayload(Type type, Location loc, ArrayRef<Operation *> payload) const {
+    return DiagnosedSilenceableFailure::success();
+  }
+};
+} // namespace
+
+namespace {
+/// PDL extension of the Transform dialect. This provides transform operations
+/// that connect to PDL matching as well as interfaces for PDL types to be used
+/// with Transform dialect operations.
+class PDLExtension : public transform::TransformDialectExtension<PDLExtension> {
+public:
+  void init() {
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.cpp.inc"
+        >();
+
+    addDialectDataInitializer<transform::PDLMatchHooks>(
+        [](transform::PDLMatchHooks &) {});
+
+    // Declare PDL as dependent so we can attach an interface to its type in the
+    // later step.
+    declareDependentDialect<pdl::PDLDialect>();
+
+    // PDLInterp is only relevant if we actually apply the transform IR so
+    // declare it as generated.
+    declareGeneratedDialect<pdl_interp::PDLInterpDialect>();
+
+    // Make PDL OperationType usable as a transform dialect type.
+    addCustomInitializationStep([](MLIRContext *context) {
+      pdl::OperationType::attachInterface<
+          PDLOperationTypeTransformHandleTypeInterfaceImpl>(*context);
+    });
+  }
+};
+} // namespace
+
+void mlir::transform::registerPDLExtension(DialectRegistry &dialectRegistry) {
+  dialectRegistry.addExtensions<PDLExtension>();
+}

mlir/lib/Dialect/Transform/PDLExtension/PDLExtensionOps.cpp

+//===- PDLExtensionOps.cpp - PDL extension for the Transform dialect ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h"
+#include "mlir/Dialect/PDL/IR/PDLOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"
+#include "mlir/Rewrite/PatternApplicator.h"
+#include "llvm/ADT/ScopeExit.h"
+
+using namespace mlir;
+
+MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::transform::PDLMatchHooks)
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// PatternApplicatorExtension
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// A TransformState extension that keeps track of compiled PDL pattern sets.
+/// This is intended to be used along the WithPDLPatterns op. The extension
+/// can be constructed given an operation that has a SymbolTable trait and
+/// contains pdl::PatternOp instances. The patterns are compiled lazily and one
+/// by one when requested; this behavior is subject to change.
+class PatternApplicatorExtension : public transform::TransformState::Extension {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PatternApplicatorExtension)
+
+  /// Creates the extension for patterns contained in `patternContainer`.
+  explicit PatternApplicatorExtension(transform::TransformState &state,
+                                      Operation *patternContainer)
+      : Extension(state), patterns(patternContainer) {}
+
+  /// Appends to `results` the operations contained in `root` that matched the
+  /// PDL pattern with the given name. Note that `root` may or may not be the
+  /// operation that contains PDL patterns. Reports an error if the pattern
+  /// cannot be found. Note that when no operations are matched, this still
+  /// succeeds as long as the pattern exists.
+  LogicalResult findAllMatches(StringRef patternName, Operation *root,
+                               SmallVectorImpl<Operation *> &results);
+
+private:
+  /// Map from the pattern name to a singleton set of rewrite patterns that only
+  /// contains the pattern with this name. Populated when the pattern is first
+  /// requested.
+  // TODO: reconsider the efficiency of this storage when more usage data is
+  // available. Storing individual patterns in a set and triggering compilation
+  // for each of them has overhead. So does compiling a large set of patterns
+  // only to apply a handful of them.
+  llvm::StringMap<FrozenRewritePatternSet> compiledPatterns;
+
+  /// A symbol table operation containing the relevant PDL patterns.
+  SymbolTable patterns;
+};
+
+LogicalResult PatternApplicatorExtension::findAllMatches(
+    StringRef patternName, Operation *root,
+    SmallVectorImpl<Operation *> &results) {
+  auto it = compiledPatterns.find(patternName);
+  if (it == compiledPatterns.end()) {
+    auto patternOp = patterns.lookup<pdl::PatternOp>(patternName);
+    if (!patternOp)
+      return failure();
+
+    // Copy the pattern operation into a new module that is compiled and
+    // consumed by the PDL interpreter.
+    OwningOpRef<ModuleOp> pdlModuleOp = ModuleOp::create(patternOp.getLoc());
+    auto builder = OpBuilder::atBlockEnd(pdlModuleOp->getBody());
+    builder.clone(*patternOp);
+    PDLPatternModule patternModule(std::move(pdlModuleOp));
+
+    // Merge in the hooks owned by the dialect. Make a copy as they may be
+    // also used by the following operations.
+    auto *dialect =
+        root->getContext()->getLoadedDialect<transform::TransformDialect>();
+    for (const auto &[name, constraintFn] :
+         dialect->getExtraData<transform::PDLMatchHooks>()
+             .getPDLConstraintHooks()) {
+      patternModule.registerConstraintFunction(name, constraintFn);
+    }
+
+    // Register a noop rewriter because PDL requires patterns to end with some
+    // rewrite call.
+    patternModule.registerRewriteFunction(
+        "transform.dialect", [](PatternRewriter &, Operation *) {});
+
+    it = compiledPatterns
+             .try_emplace(patternOp.getName(), std::move(patternModule))
+             .first;
+  }
+
+  PatternApplicator applicator(it->second);
+  // We want to discourage direct use of PatternRewriter in APIs but In this
+  // very specific case, an IRRewriter is not enough.
+  struct TrivialPatternRewriter : public PatternRewriter {
+  public:
+    explicit TrivialPatternRewriter(MLIRContext *context)
+        : PatternRewriter(context) {}
+  };
+  TrivialPatternRewriter rewriter(root->getContext());
+  applicator.applyDefaultCostModel();
+  root->walk([&](Operation *op) {
+    if (succeeded(applicator.matchAndRewrite(op, rewriter)))
+      results.push_back(op);
+  });
+
+  return success();
+}
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// PDLMatchHooks
+//===----------------------------------------------------------------------===//
+
+void transform::PDLMatchHooks::mergeInPDLMatchHooks(
+    llvm::StringMap<PDLConstraintFunction> &&constraintFns) {
+  // Steal the constraint functions from the given map.
+  for (auto &it : constraintFns)
+    pdlMatchHooks.registerConstraintFunction(it.getKey(), std::move(it.second));
+}
+
+const llvm::StringMap<PDLConstraintFunction> &
+transform::PDLMatchHooks::getPDLConstraintHooks() const {
+  return pdlMatchHooks.getConstraintFunctions();
+}
+
+//===----------------------------------------------------------------------===//
+// PDLMatchOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::PDLMatchOp::apply(transform::TransformResults &results,
+                             transform::TransformState &state) {
+  auto *extension = state.getExtension<PatternApplicatorExtension>();
+  assert(extension &&
+         "expected PatternApplicatorExtension to be attached by the parent op");
+  SmallVector<Operation *> targets;
+  for (Operation *root : state.getPayloadOps(getRoot())) {
+    if (failed(extension->findAllMatches(
+            getPatternName().getLeafReference().getValue(), root, targets))) {
+      emitDefiniteFailure()
+          << "could not find pattern '" << getPatternName() << "'";
+    }
+  }
+  results.set(llvm::cast<OpResult>(getResult()), targets);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::PDLMatchOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getRoot(), effects);
+  producesHandle(getMatched(), effects);
+  onlyReadsPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
+// WithPDLPatternsOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
+                                    transform::TransformState &state) {
+  TransformOpInterface transformOp = nullptr;
+  for (Operation &nested : getBody().front()) {
+    if (!isa<pdl::PatternOp>(nested)) {
+      transformOp = cast<TransformOpInterface>(nested);
+      break;
+    }
+  }
+
+  state.addExtension<PatternApplicatorExtension>(getOperation());
+  auto guard = llvm::make_scope_exit(
+      [&]() { state.removeExtension<PatternApplicatorExtension>(); });
+
+  auto scope = state.make_region_scope(getBody());
+  if (failed(mapBlockArguments(state)))
+    return DiagnosedSilenceableFailure::definiteFailure();
+  return state.applyTransform(transformOp);
+}
+
+void transform::WithPDLPatternsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  getPotentialTopLevelEffects(effects);
+}
+
+LogicalResult transform::WithPDLPatternsOp::verify() {
+  Block *body = getBodyBlock();
+  Operation *topLevelOp = nullptr;
+  for (Operation &op : body->getOperations()) {
+    if (isa<pdl::PatternOp>(op))
+      continue;
+
+    if (op.hasTrait<::mlir::transform::PossibleTopLevelTransformOpTrait>()) {
+      if (topLevelOp) {
+        InFlightDiagnostic diag =
+            emitOpError() << "expects only one non-pattern op in its body";
+        diag.attachNote(topLevelOp->getLoc()) << "first non-pattern op";
+        diag.attachNote(op.getLoc()) << "second non-pattern op";
+        return diag;
+      }
+      topLevelOp = &op;
+      continue;
+    }
+
+    InFlightDiagnostic diag =
+        emitOpError()
+        << "expects only pattern and top-level transform ops in its body";
+    diag.attachNote(op.getLoc()) << "offending op";
+    return diag;
+  }
+
+  if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {
+    InFlightDiagnostic diag = emitOpError() << "cannot be nested";
+    diag.attachNote(parent.getLoc()) << "parent operation";
+    return diag;
+  }
+
+  if (!topLevelOp) {
+    InFlightDiagnostic diag = emitOpError()
+                              << "expects at least one non-pattern op";
+    return diag;
+  }
+
+  return success();
+}

mlir/python/CMakeLists.txt

declare_mlir_dialect_python_bindings(
   TD_FILE dialects/LinalgOps.td
   SOURCES
     dialects/_linalg_ops_ext.py
   SOURCES_GLOB
     dialects/linalg/*.py
   DIALECT_NAME linalg
   DEPENDS LinalgOdsGen)
 
+declare_mlir_dialect_extension_python_bindings(
+ADD_TO_PARENT MLIRPythonSources.Dialects
+ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/TransformPDLExtensionOps.td
+  SOURCES
+    dialects/_transform_pdl_extension_ops_ext.py
+    dialects/transform/pdl.py
+  DIALECT_NAME transform
+  EXTENSION_NAME transform_pdl_extension)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformOps.td
   SOURCES
     dialects/_transform_ops_ext.py
     dialects/transform/__init__.py
     _mlir_libs/_mlir/dialects/transform/__init__.pyi

mlir/python/mlir/dialects/TransformPDLExtensionOps.td

+//===-- TransformPDLExtensionOps.td - Binding entry point --*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Entry point of the generated Python bindings for the PDL extension of the
+// Transform dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_TRANSFORM_PDL_EXTENSION_OPS
+#define PYTHON_BINDINGS_TRANSFORM_PDL_EXTENSION_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td"
+
+#endif // PYTHON_BINDINGS_TRANSFORM_PDL_EXTENSION_OPS

mlir/python/mlir/dialects/_transform_ops_ext.py

):
     super().__init__(
         [_get_op_result_or_value(h) for h in handles],
         deduplicate=deduplicate,
         loc=loc,
         ip=ip,
     )
 
 
-class PDLMatchOp:
-
-  def __init__(
-      self,
-      result_type: Type,
-      target: Union[Operation, Value],
-      pattern_name: Union[Attribute, str],
-      *,
-      loc=None,
-      ip=None,
-  ):
-    super().__init__(
-        result_type,
-        _get_op_result_or_value(target),
-        pattern_name,
-        loc=loc,
-        ip=ip,
-    )
-
-
 class ReplicateOp:
 
   def __init__(
       self,
       pattern: Union[Operation, Value],
       handles: Sequence[Union[Operation, Value]],
       *,
       loc=None,
class SequenceOp:
   def bodyTarget(self) -> Value:
     return self.body.arguments[0]
 
   @property
   def bodyExtraArgs(self) -> BlockArgumentList:
     return self.body.arguments[1:]
 
 
-class WithPDLPatternsOp:
-
-  def __init__(self,
-               target: Union[Operation, Value, Type],
-               *,
-               loc=None,
-               ip=None):
-    root = _get_op_result_or_value(target) if not isinstance(target,
-                                                             Type) else None
-    root_type = target if isinstance(target, Type) else root.type
-    super().__init__(root=root, loc=loc, ip=ip)
-    self.regions[0].blocks.append(root_type)
-
-  @property
-  def body(self) -> Block:
-    return self.regions[0].blocks[0]
-
-  @property
-  def bodyTarget(self) -> Value:
-    return self.body.arguments[0]
-
-
 class YieldOp:
 
   def __init__(
       self,
       operands: Optional[Union[Operation, Sequence[Value]]] = None,
       *,
       loc=None,
       ip=None,
   ):
     if operands is None:
       operands = []
     super().__init__(_get_op_results_or_values(operands), loc=loc, ip=ip)

mlir/python/mlir/dialects/_transform_pdl_extension_ops_ext.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+try:
+  from ..ir import *
+  from ._ods_common import (
+      get_op_result_or_value as _get_op_result_or_value,
+      get_op_results_or_values as _get_op_results_or_values,
+  )
+except ImportError as e:
+  raise RuntimeError("Error loading imports from extension module") from e
+
+from typing import Union
+
+class PDLMatchOp:
+
+  def __init__(
+      self,
+      result_type: Type,
+      target: Union[Operation, Value],
+      pattern_name: Union[Attribute, str],
+      *,
+      loc=None,
+      ip=None,
+  ):
+    super().__init__(
+        result_type,
+        _get_op_result_or_value(target),
+        pattern_name,
+        loc=loc,
+        ip=ip,
+    )
+
+
+class WithPDLPatternsOp:
+
+  def __init__(self,
+               target: Union[Operation, Value, Type],
+               *,
+               loc=None,
+               ip=None):
+    root = _get_op_result_or_value(target) if not isinstance(target,
+                                                             Type) else None
+    root_type = target if isinstance(target, Type) else root.type
+    super().__init__(root=root, loc=loc, ip=ip)
+    self.regions[0].blocks.append(root_type)
+
+  @property
+  def body(self) -> Block:
+    return self.regions[0].blocks[0]
+
+  @property
+  def bodyTarget(self) -> Value:
+    return self.body.arguments[0]

mlir/python/mlir/dialects/transform/pdl.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._transform_pdl_extension_ops_gen import *

mlir/test/Dialect/Transform/test-interpreter.mlir

^bb0(%arg1: !transform.any_op):
     yield %1 : !transform.any_op
   }
   // expected-remark @below {{succeeded}}
   test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand" : !transform.any_op
 }
 
 // -----
 
-transform.with_pdl_patterns {
-^bb0(%arg0: !transform.any_op):
-  sequence %arg0 : !transform.any_op failures(propagate) {
-  ^bb0(%arg1: !transform.any_op):
-    %0 = pdl_match @some in %arg1 : (!transform.any_op) -> !transform.any_op
-    test_print_remark_at_operand %0, "matched" : !transform.any_op
-  }
-
-  pdl.pattern @some : benefit(1) {
-    %0 = pdl.operation "test.some_op"
-    pdl.rewrite %0 with "transform.dialect"
-  }
-
-  pdl.pattern @other : benefit(1) {
-    %0 = pdl.operation "test.other_op"
-    pdl.rewrite %0 with "transform.dialect"
-  }
-}
-
-// expected-remark @below {{matched}}
-"test.some_op"() : () -> ()
-"test.other_op"() : () -> ()
-// expected-remark @below {{matched}}
-"test.some_op"() : () -> ()
-
-// -----
-
 // expected-remark @below {{parent function}}
 func.func @foo() {
   %0 = arith.constant 0 : i32
   return
 }
 
 // expected-remark @below {{parent function}}
 func.func @bar() {

mlir/test/Dialect/Transform/test-pdl-extension.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !transform.any_op):
+  sequence %arg0 : !transform.any_op failures(propagate) {
+  ^bb0(%arg1: !transform.any_op):
+    %0 = pdl_match @some in %arg1 : (!transform.any_op) -> !transform.any_op
+    test_print_remark_at_operand %0, "matched" : !transform.any_op
+  }
+
+  pdl.pattern @some : benefit(1) {
+    %0 = pdl.operation "test.some_op"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  pdl.pattern @other : benefit(1) {
+    %0 = pdl.operation "test.other_op"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+}
+
+// expected-remark @below {{matched}}
+"test.some_op"() : () -> ()
+"test.other_op"() : () -> ()
+// expected-remark @below {{matched}}
+"test.some_op"() : () -> ()
+
+
+// -----
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !transform.any_op):
+  sequence %arg0 : !transform.any_op failures(propagate) {
+  ^bb1(%arg1: !transform.any_op):
+    %0 = pdl_match @some in %arg1 : (!transform.any_op) -> !transform.any_op
+  }
+
+  pdl.pattern @some : benefit(1) {
+    %0 = pdl.operation "test.some_op"
+    pdl.apply_native_constraint "verbose_constraint"(%0 : !pdl.operation)
+    pdl.rewrite %0 with "transform.dialect"
+  }
+}
+
+// expected-warning @below {{from PDL constraint}}
+"test.some_op"() : () -> ()
+"test.other_op"() : () -> ()

mlir/test/lib/Dialect/Transform/CMakeLists.txt

add_mlir_library(MLIRTestTransformDialect
   MLIRTestTransformDialectExtensionIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRPDLDialect
   MLIRTransformDialect
   MLIRTransformDialectTransforms
+  MLIRTransformPDLExtension
 )

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 //===----------------------------------------------------------------------===//
 
 #include "TestTransformDialectExtension.h"
 #include "TestTransformStateExtension.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 namespace {
registerTransformOps<TestTransformOp,
                          TestTransformUnrestrictedOpNoInterface,
 #define GET_OP_LIST
 #include "TestTransformDialectExtension.cpp.inc"
                          >();
     registerTypes<
 #define GET_TYPEDEF_LIST
 #include "TestTransformDialectExtensionTypes.cpp.inc"
         >();
+
+    auto verboseConstraint = [](PatternRewriter &rewriter,
+                                ArrayRef<PDLValue> pdlValues) {
+      for (const PDLValue &pdlValue : pdlValues) {
+        if (Operation *op = pdlValue.dyn_cast<Operation *>()) {
+          op->emitWarning() << "from PDL constraint";
+        }
+      }
+      return success();
+    };
+
+    addDialectDataInitializer<transform::PDLMatchHooks>(
+        [&](transform::PDLMatchHooks &hooks) {
+          llvm::StringMap<PDLConstraintFunction> constraints;
+          constraints.try_emplace("verbose_constraint", verboseConstraint);
+          hooks.mergeInPDLMatchHooks(std::move(constraints));
+        });
   }
 };
 } // namespace
 
 // These are automatically generated by ODS but are not used as the Transform
 // dialect uses a different dispatch mechanism to support dialect extensions.
 LLVM_ATTRIBUTE_UNUSED static OptionalParseResult
 generatedTypeParser(AsmParser &parser, StringRef *mnemonic, Type &value);

mlir/test/python/dialects/transform.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import transform
-from mlir.dialects import pdl
+from mlir.dialects.transform import pdl as transform_pdl
 
 
 def run(f):
   with Context(), Location.unknown():
     module = Module.create()
     with InsertionPoint(module.body):
       print("\nTEST:", f.__name__)
       f()
def testNestedSequenceOpWithExtras():
   # CHECK-LABEL: TEST: testNestedSequenceOpWithExtras
   # CHECK: transform.sequence failures(propagate)
   # CHECK: ^{{.*}}(%[[ARG0:.*]]: !transform.any_op, %[[ARG1:.*]]: !transform.any_op, %[[ARG2:.*]]: !transform.op<"foo.bar">):
   # CHECK:   sequence %[[ARG0]], %[[ARG1]], %[[ARG2]] : (!transform.any_op, !transform.any_op, !transform.op<"foo.bar">)
 
 
 @run
 def testTransformPDLOps():
-  withPdl = transform.WithPDLPatternsOp(transform.AnyOpType.get())
+  withPdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())
   with InsertionPoint(withPdl.body):
     sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
                                     [transform.AnyOpType.get()],
                                     withPdl.bodyTarget)
     with InsertionPoint(sequence.body):
-      match = transform.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "pdl_matcher")
+      match = transform_pdl.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "pdl_matcher")
       transform.YieldOp(match)
   # CHECK-LABEL: TEST: testTransformPDLOps
   # CHECK: transform.with_pdl_patterns {
   # CHECK: ^{{.*}}(%[[ARG0:.+]]: !transform.any_op):
   # CHECK:   = sequence %[[ARG0]] : !transform.any_op -> !transform.any_op failures(propagate) {
   # CHECK:   ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):
   # CHECK:     %[[RES:.+]] = pdl_match @pdl_matcher in %[[ARG1]]
   # CHECK:     yield %[[RES]] : !transform.any_op
def testMergeHandlesOp():
   # CHECK-LABEL: TEST: testMergeHandlesOp
   # CHECK: transform.sequence
   # CHECK: ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):
   # CHECK:   = merge_handles %[[ARG1]]
 
 
 @run
 def testReplicateOp():
-  with_pdl = transform.WithPDLPatternsOp(transform.AnyOpType.get())
+  with_pdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())
   with InsertionPoint(with_pdl.body):
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.PROPAGATE, [], with_pdl.bodyTarget)
     with InsertionPoint(sequence.body):
-      m1 = transform.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "first")
-      m2 = transform.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "second")
+      m1 = transform_pdl.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "first")
+      m2 = transform_pdl.PDLMatchOp(transform.AnyOpType.get(), sequence.bodyTarget, "second")
       transform.ReplicateOp(m1, [m2])
       transform.YieldOp()
   # CHECK-LABEL: TEST: testReplicateOp
   # CHECK: %[[FIRST:.+]] = pdl_match
   # CHECK: %[[SECOND:.+]] = pdl_match
   # CHECK: %{{.*}} = replicate num(%[[FIRST]]) %[[SECOND]]

mlir/test/python/dialects/transform_structured_ext.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import transform
 from mlir.dialects import pdl
 from mlir.dialects.transform import structured
+from mlir.dialects.transform import pdl as transform_pdl
 
 
 def run(f):
   with Context(), Location.unknown():
     module = Module.create()
     with InsertionPoint(module.body):
       print("\nTEST:", f.__name__)
       f()
with InsertionPoint(sequence.body):
     transform.YieldOp()
   # CHECK-LABEL: TEST: testTileZero
   # CHECK: transform.sequence
   # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 0, 2, 0]
   # CHECK: interchange = [0, 1, 2, 3]
 
 @run
 def testTileDynamic():
-  with_pdl = transform.WithPDLPatternsOp(pdl.OperationType.get())
+  with_pdl = transform_pdl.WithPDLPatternsOp(pdl.OperationType.get())
   with InsertionPoint(with_pdl.body):
     sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [],
                                     with_pdl.bodyTarget)
     with InsertionPoint(sequence.body):
-      m1 = transform.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "first")
-      m2 = transform.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "second")
+      m1 = transform_pdl.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "first")
+      m2 = transform_pdl.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "second")
       structured.TileOp(sequence.bodyTarget,
                         sizes=[m1, 3, m2, 0])
       transform.YieldOp()
   # CHECK-LABEL: TEST: testTileDynamic
   # CHECK: %[[FIRST:.+]] = pdl_match
   # CHECK: %[[SECOND:.+]] = pdl_match
   # CHECK: %{{.+}}, %{{.+}}:3 = transform.structured.tile %{{.*}}[%[[FIRST]], 3, %[[SECOND]], 0]
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":TensorInferTypeOpInterfaceImpl",
         ":TensorTilingInterfaceImpl",
         ":TensorTransformOps",
         ":TensorTransforms",
         ":TosaDialect",
         ":TosaToLinalg",
         ":TransformDialect",
         ":TransformDialectTransforms",
+        ":TransformPDLExtension",
         ":Transforms",
         ":TransformsPassIncGen",
         ":VectorDialect",
         ":VectorToLLVM",
         ":VectorToSCF",
         ":VectorToSPIRV",
         ":VectorTransformOps",
         ":VectorTransforms",

 td_library(
     name = "TransformDialectTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/IR/*.td"]),
     deps = [
         ":CastInterfacesTdFiles",
         ":ControlFlowInterfacesTdFiles",
         ":InferTypeOpInterfaceTdFiles",
         ":OpBaseTdFiles",
-        ":PDLDialectTdFiles",
         ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "TransformDialectEnumsIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
cc_library(
     name = "TransformDialect",
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":Analysis",
         ":CallOpInterfaces",
         ":ControlFlowInterfaces",
         ":IR",
-        ":PDLDialect",
-        ":PDLInterpDialect",
         ":Rewrite",
         ":SideEffectInterfaces",
         ":Support",
         ":TransformDialectEnumsIncGen",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
         ":TransformDialectMatchInterfacesIncGen",
         ":TransformDialectUtils",
         ":TransformOpsIncGen",
         ":TransformTypesIncGen",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 td_library(
+    name = "TransformPDLExtensionTdFiles",
+    srcs = glob(["include/mlir/Dialect/Transform/PDLExtension/*.td"]),
+    deps = [
+        ":PDLDialectTdFiles",
+        ":TransformDialectTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TransformPDLExtensionOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            [
+                "-gen-op-decls",
+            ],
+            "include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h.inc",
+        ),
+        (
+            [
+                "-gen-op-defs",
+            ],
+            "include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td",
+    deps = [":TransformPDLExtensionTdFiles"],
+)
+
+cc_library(
+    name = "TransformPDLExtension",
+    srcs = glob(["lib/Dialect/Transform/PDLExtension/*.cpp"]),
+    hdrs = glob(["include/mlir/Dialect/Transform/PDLExtension/*.h"]),
+    deps = [
+        ":IR",
+        ":PDLDialect",
+        ":PDLInterpDialect",
+        ":SideEffectInterfaces",
+        ":Support",
+        ":TransformDialect",
+        ":TransformDialectTransformsIncGen",
+        "//llvm:Support",
+    ],
+)
+
+td_library(
     name = "TransformDialectTransformsTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/Transforms/*.td"]),
     deps = [
         ":PassBaseTdFiles",
         ":TransformDialectTdFiles",
     ],
 )
 

utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

tbl_outs = [
             ],
             "mlir/dialects/_loop_transform_ops_gen.py",
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",
     td_file = "mlir/dialects/SCFLoopTransformOps.td",
     deps = [
         ":TransformOpsPyTdFiles",
+        "//mlir:TransformPDLExtensionTdFiles",
+    ],
+)
+
+gentbl_filegroup(
+    name = "PDLTransformOpsPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=transform",
+                "-dialect-extension=transform_pdl_extension",
+            ],
+            "mlir/dialects/_transform_pdl_extension_ops_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/TransformPDLExtensionOps.td",
+    deps = [
+        ":TransformOpsPyTdFiles",
         "//mlir:SCFTransformOpsTdFiles",
     ],
 )
 
 filegroup(
     name = "TransformOpsPyFiles",
     srcs = [
         "mlir/dialects/_loop_transform_ops_ext.py",
         "mlir/dialects/_structured_transform_ops_ext.py",
         "mlir/dialects/_transform_ops_ext.py",
         ":LoopTransformOpsPyGen",
+        ":PDLTransformOpsPyGen",
         ":StructuredTransformOpsPyGen",
         ":TransformOpsPyGen",
     ],
 )
 
 filegroup(
     name = "TransformOpsPackagePyFiles",
     srcs = glob(["mlir/dialects/transform/*.py"]),

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

gentbl_cc_library(
     ],
     tblgen = "//mlir:mlir-tblgen",
     td_file = "lib/Dialect/Transform/TestTransformDialectExtension.td",
     test = True,
     deps = [
         ":TransformDialectTdFiles",
         "//mlir:PDLDialectTdFiles",
         "//mlir:TransformDialectTdFiles",
+        "//mlir:TransformPDLExtension",
     ],
 )
 
 cc_library(
     name = "TestTransformDialect",
     srcs = glob(["lib/Dialect/Transform/*.cpp"]),
     hdrs = glob(["lib/Dialect/Transform/*.h"]),
     includes = ["lib/Dialect/Transform"],