[mlir] add types to the transform dialect

[ftynse]

Summary

Introduce a type system for the transform dialect. A transform IR type
captures the expectations of the transform IR on the payload IR
operations that are being transformed, such as being of a certain kind
or implementing an interface that enables the transformation. This
provides stricter checking and better readability of the transform IR
than using the catch-all "handle" type.

This change implements the basic support for a type system amendable to
dialect extensions and adds a drop-in replacement for the unrestricted
"handle" type. The actual switch of transform dialect ops to that type
will happen in a separate commit.

See https://discourse.llvm.org/t/rfc-type-system-for-the-transform-dialect/65702

[nicolasvasilache]

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

380

Does this need to be a public method?
Also, add an Impl suffix to make it clearer it is an impl detail?

398

nit: I like to add a redundant public: to separate the methods from contained objects

407

Can you comment on the asymmetry here ?
Why StringRef for one set of hooks and typeId for the second?

[ftynse]

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

380

It's not public, there is a private: in line 324.

398

This one is not public either, I wouldn't want to make it such.

[ftynse]

Address review.

mlir/docs/Dialects/Transform.md

 # Transform Dialect
 
 [TOC]
 
+[include "Dialects/TransformTypes.md"]
+
 [include "Dialects/TransformOps.md"]
 
 ## Bufferization Transform Operations
 
 [include "Dialects/BufferizationTransformOps.md"]
 
 ## Loop (SCF) Transform Operations
 
 [include "Dialects/SCFLoopTransformOps.md"]
 
 ## Structured (Linalg) Transform Operations
 
 [include "Dialects/LinalgStructuredTransformOps.md"]
 
+[include "Dialects/TransformTypeInterfaces.md"]
+
 [include "Dialects/TransformOpInterfaces.md"]

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

 # Generate the dialect files from the dialect .td.
 #
 # TODO: Make it possible to use XDialect instead of XOpsDialect in
 # add_mlir_dialect.
 set(LLVM_TARGET_DEFINITIONS TransformDialect.td)
 mlir_tablegen(TransformDialect.h.inc -gen-dialect-decls -dialect=transform)
 mlir_tablegen(TransformDialect.cpp.inc -gen-dialect-defs -dialect=transform)
 add_public_tablegen_target(MLIRTransformDialectIncGen)
 add_dependencies(mlir-headers MLIRTransformDialectIncGen)
 
+set(LLVM_TARGET_DEFINITIONS TransformTypes.td)
+mlir_tablegen(TransformTypes.h.inc -gen-typedef-decls)
+mlir_tablegen(TransformTypes.cpp.inc -gen-typedef-defs)
+add_public_tablegen_target(MLIRTransformTypesIncGen)
+add_dependencies(mlir-headers MLIRTransformTypesIncGen)
+add_mlir_doc(TransformTypes TransformTypes Dialects/ -gen-typedef-docs)
+
 set(LLVM_TARGET_DEFINITIONS TransformAttrs.td)
 mlir_tablegen(TransformDialectEnums.h.inc -gen-enum-decls)
 mlir_tablegen(TransformDialectEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRTransformDialectEnumIncGen)
 add_dependencies(mlir-headers MLIRTransformDialectEnumIncGen)
 
 add_mlir_dialect(TransformOps transform)
 add_mlir_doc(TransformOps TransformOps Dialects/ -gen-dialect-doc -dialect=transform)
 
+# Contrary to what the name claims, this only produces the _op_ interface.
 add_mlir_interface(TransformInterfaces)
 add_mlir_doc(TransformInterfaces TransformOpInterfaces Dialects/ -gen-op-interface-docs)
+
+set(LLVM_TARGET_DEFINITIONS TransformInterfaces.td)
+mlir_tablegen(TransformTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(TransformTypeInterfaces.cpp.inc -gen-type-interface-defs)
+add_public_tablegen_target(MLIRTransformDialectTypeInterfacesIncGen)
+add_dependencies(mlir-headers MLIRTransformDialectTypeInterfacesIncGen)
+add_mlir_doc(TransformInterfaces TransformTypeInterfaces Dialects/ -gen-type-interface-docs)

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/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 namespace transform {
 #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.
 template <typename OpTy>
 static inline void checkImplementsTransformInterface(MLIRContext *context) {
   // Since the operation is being inserted into the Transform dialect and the
   // dialect does not implement the interface fallback, only check for the op
   // itself having the interface implementation.
   RegisteredOperationName opName =
       *RegisteredOperationName::lookup(OpTy::getOperationName(), context);
   assert((opName.hasInterface<TransformOpInterface>() ||
           opName.hasTrait<OpTrait::IsTerminator>()) &&
          "non-terminator ops injected into the transform dialect must "
          "implement TransformOpInterface");
   assert(opName.hasInterface<MemoryEffectOpInterface>() &&
          "ops injected into the transform dialect must implement "
          "MemoryEffectsOpInterface");
 }
+
+/// Asserts that the type provided as template argument implements the
+/// TransformTypeInterface. This must be a dynamic assertion since interface
+/// implementations may be registered at runtime.
+void checkImplementsTransformTypeInterface(TypeID typeID, MLIRContext *context);
 } // namespace detail
 #endif // NDEBUG
 } // namespace transform
 } // namespace mlir
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h.inc"
 
 namespace mlir {
protected:
   /// implementations must be already available when the operation is injected.
   template <typename... OpTys>
   void registerTransformOps() {
     opInitializers.push_back([](TransformDialect *transformDialect) {
       transformDialect->addOperationsChecked<OpTys...>();
     });
   }
 
+  /// Injects the types into the Transform dialect. The types must implement
+  /// the TransformTypeInterface 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) {
+      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
   /// the injected operations, similarly to how the dialect definition may list
   /// dependent dialects. This is *not* intended for dialects entities from
   /// which may be produced when applying the transformations specified by ops
   /// registered by this extension.
private:
   /// 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 Type>
+void TransformDialect::addTypeIfNotRegistered() {
+  // Use the address of the parse method as a proxy for identifying whether we
+  // are registering the same type class for the same mnemonic.
+  StringRef mnemonic = Type::getMnemonic();
+  auto [it, inserted] = typeParsingHooks.try_emplace(mnemonic, Type::parse);
+  if (!inserted) {
+    const ExtensionTypeParsingHook &parsingHook = it->getValue();
+    if (*parsingHook.target<mlir::Type (*)(AsmParser &)>() != &Type::parse)
+      reportDuplicateTypeRegistration(mnemonic);
+  }
+  typePrintingHooks.try_emplace(
+      TypeID::get<Type>(), +[](mlir::Type type, AsmPrinter &printer) {
+        printer << Type::getMnemonic();
+        cast<Type>(type).print(printer);
+      });
+  addTypes<Type>();
+}
+
 /// 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
 
 #include "mlir/Dialect/Transform/IR/TransformDialectEnums.h.inc"
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT_H

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

def Transform_Dialect : Dialect {
   ];
 
   let extraClassDeclaration = [{
       /// 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;
 
+      /// 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;
+
+      /// Parser callback for an individual type registered by this dialect or
+      /// its extensions.
+      using ExtensionTypeParsingHook =
+          std::function<::mlir::Type (::mlir::AsmParser &)>;
+
+      /// Printer callback for an individual type registered by this dialect or
+      /// its extensions.
+      using ExtensionTypePrintingHook =
+          std::function<void (::mlir::Type, ::mlir::AsmPrinter &)>;
+
     private:
       template <typename OpTy>
       void addOperationIfNotRegistered() {
         Optional<RegisteredOperationName> opName =
             RegisteredOperationName::lookup(OpTy::getOperationName(),
                                             getContext());
         if (!opName)
           return addOperations<OpTy>();
private:
       void addOperationsChecked() {
         (addOperationIfNotRegistered<OpTys>(),...);
 
         #ifndef NDEBUG
         (detail::checkImplementsTransformInterface<OpTys>(getContext()),...);
         #endif // NDEBUG
       }
 
+      /// Registers the types specified as template parameters with the
+      /// Transform dialect. Checks that they meet the requirements for
+      /// Transform IR types.
+      template <typename... TypeTys>
+      void addTypesChecked() {
+        (addTypeIfNotRegistered<TypeTys>(), ...);
+
+        #ifndef NDEBUG
+        (detail::checkImplementsTransformTypeInterface(
+            TypeID::get<TypeTys>(), getContext()), ...);
+        #endif // NDEBUG
+      }
+
+      /// Implementation of the type registration for a single type, should
+      /// not be called directly, use addTypesChecked instead.
+      template <typename Type>
+      void addTypeIfNotRegistered();
+
+      /// Reports a repeated registration error of a type with the given
+      /// mnemonic.
+      [[noreturn]] void reportDuplicateTypeRegistration(StringRef mnemonic);
+
       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);
 
+      //===----------------------------------------------------------------===//
+      // Data fields
+      //===----------------------------------------------------------------===//
+
       /// A container for PDL constraint function that can be used by
       /// operations in this dialect.
-      PDLPatternModule pdlMatchHooks;
+      ::mlir::PDLPatternModule pdlMatchHooks;
+
+      /// 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.
+      ::llvm::DenseMap<::mlir::TypeID, ExtensionTypePrintingHook>
+      typePrintingHooks;
   }];
 }
 
 // Base class for ops that belong to the tranfsorm dialect. Ops defined in
 // extensions of this dialect may also use this.
 class TransformDialectOp<string mnemonic, list<Trait> traits = []>
     : Op<Transform_Dialect, mnemonic, traits>;
 

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

 //===- TransformInterfaces.h - Transform Dialect Interfaces -----*- 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_TRANSFORMINTERFACES_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ScopeExit.h"
 
 namespace mlir {
 
 /// The result of a transform IR operation application. This can have one of the
 /// three states:
 ///   - success;
 ///   - silenceable (recoverable) failure with yet-unreported diagnostic;
public:
   /// Applies the transformation specified by the given transform op and updates
   /// the state accordingly.
   DiagnosedSilenceableFailure applyTransform(TransformOpInterface transform);
 
   /// Records the mapping between a block argument in the transform IR and a
   /// list of operations in the payload IR. The arguments must be defined in
   /// blocks of the currently processed transform IR region, typically after a
   /// region scope is defined.
-  void mapBlockArguments(BlockArgument argument,
+  ///
+  /// Returns failure if the payload does not satisfy the conditions associated
+  /// with the type of the handle value.
+  LogicalResult mapBlockArguments(BlockArgument argument,
                                   ArrayRef<Operation *> operations) {
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
     assert(argument.getParentRegion() == regionStack.back() &&
            "mapping block arguments from a region other than the active one");
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
-    setPayloadOps(argument, operations);
+    return setPayloadOps(argument, operations);
   }
 
   // Forward declarations to support limited visibility.
   class RegionScope;
 
   /// Creates a new region scope for the given region. The region is expected to
   /// be nested in the currently processed region.
   // Implementation note: this method is inline but implemented outside of the
private:
   /// rename function" transform on both handles:
   ///
   ///   transform.rewrite_and_rename %0 { new_name = "func" }
   ///   transform.rewrite_and_rename %1 { new_name = "func" }
   ///
   /// is invalid given the transformation "consumes" the handle as expressed
   /// by side effects. Practically, a transformation consuming a handle means
   /// that the associated payload operation may no longer exist.
-  void setPayloadOps(Value value, ArrayRef<Operation *> targets);
+  ///
+  /// Returns failure if the payload does not satisfy the conditions associated
+  /// with the type of the handle value.
+  LogicalResult setPayloadOps(Value value, ArrayRef<Operation *> targets);
 
   /// Forgets the payload IR ops associated with the given transform IR value.
   void removePayloadOps(Value value);
 
   /// Updates the payload IR ops associated with the given transform IR value.
   /// The callback function is called once per associated operation and is
   /// expected to return the modified operation or nullptr. In the latter case,
   /// the corresponding operation is no longer associated with the transform IR
   /// value.
-  void updatePayloadOps(Value value,
-                        function_ref<Operation *(Operation *)> callback);
+  ///
+  /// Returns failure if the payload does not satisfy the conditions associated
+  /// with the type of the handle value.
+  LogicalResult
+  updatePayloadOps(Value value,
 
   /// If the operand is a handle consumed by the operation, i.e. has the "free"
   /// memory effect associated with it, identifies other handles that are
   /// pointing to payload IR operations nested in the operations pointed to by
   /// the consumed handle. Marks all such handles as invalidated to trigger
   /// errors if they are used.
   void recordHandleInvalidation(OpOperand &handle);
   void recordHandleInvalidationOne(OpOperand &handle, Operation *payloadOp,

 TransformState::RegionScope TransformState::make_region_scope(Region &region) {
   return RegionScope(*this, region);
 }
 
 namespace detail {
 /// Maps the only block argument of the op with PossibleTopLevelTransformOpTrait
 /// to either the list of operations associated with its operand or the root of
 /// the payload IR, depending on what is available in the context.
-void mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
-                                                  Operation *op,
-                                                  Region &region);
+LogicalResult
+mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
+                                             Operation *op, Region &region);
 
 /// Verification hook for PossibleTopLevelTransformOpTrait.
 LogicalResult verifyPossibleTopLevelTransformOpTrait(Operation *op);
 } // namespace detail
 
 /// 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
public:
   /// Returns the single block of the given region.
   Block *getBodyBlock(unsigned region = 0) {
     return &this->getOperation()->getRegion(region).front();
   }
 
   /// 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.
-  void mapBlockArguments(TransformState &state, Region &region) {
+  LogicalResult mapBlockArguments(TransformState &state, Region &region) {
     assert(region.getParentOp() == this->getOperation() &&
            "op comes from the wrong region");
-    detail::mapPossibleTopLevelTransformOpBlockArguments(
+    return detail::mapPossibleTopLevelTransformOpBlockArguments(
         state, this->getOperation(), region);
   }
-  void mapBlockArguments(TransformState &state) {
+  LogicalResult mapBlockArguments(TransformState &state) {
     assert(
         this->getOperation()->getNumRegions() == 1 &&
         "must indicate the region to map if the operation has more than one");
-    mapBlockArguments(state, this->getOperation()->getRegion(0));
+    return mapBlockArguments(state, this->getOperation()->getRegion(0));
   }
 };
 
 /// Trait implementing the TransformOpInterface for operations applying a
 /// transformation to a single operation handle and producing zero, one or
 /// multiple operation handles.
 /// The op must implement a method with the following signature:
 ///   - DiagnosedSilenceableFailure applyToOne(OpTy,

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

DiagnosedSilenceableFailure emitDefaultSilenceableFailure(
       Diagnostic diag($_op->getLoc(), DiagnosticSeverity::Error);
       diag << $_op->getName() << " failed to apply";
       diag.attachNote(target->getLoc()) << "when applied to this op";
       return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
     }
   }];
 }
 
+def TransformTypeInterface : TypeInterface<"TransformTypeInterface"> {
+  let description = [{
+    Types that can be used for Transform dialect handle values. Such types
+    define the properties of Payload IR operations associated with the handle.
+    A user of such a handle can assume that these properties have been verified
+    for any Payload IR operation associated with it.
+  }];
+
+  let cppNamespace = "::mlir::transform";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Checks if the given list of associated Payload IR operations satisfy
+        the conditions defined by this type. If not, produces a silenceable
+        error at the specified location.
+      }],
+      /*returnType=*/"::mlir::DiagnosedSilenceableFailure",
+      /*name=*/"checkPayload",
+      /*arguments=*/(ins "::mlir::Location":$loc,
+                         "::mlir::ArrayRef<::mlir::Operation *>":$payload)
+    >
+  ];
+}
+
 def FunctionalStyleTransformOpTrait
     : NativeOpTrait<"FunctionalStyleTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def TransformEachOpTrait : NativeOpTrait<"TransformEachOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def NavigationTransformOpTrait : NativeOpTrait<"NavigationTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_INTERFACES_TD

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/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 namespace mlir {
 namespace transform {
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 } // namespace transform
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.h.inc"
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H

mlir/include/mlir/Dialect/Transform/IR/TransformOps.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_IR_TRANSFORMOPS
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
 
+include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Dialect/Transform/IR/TransformAttrs.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
def AlternativesOp : TransformDialectOp<"alternatives",
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region VariadicRegion<SizedRegion<1>>:$alternatives);
 
   let assemblyFormat =
     "($scope^)? (`->` type($results)^)? attr-dict-with-keyword regions";
   let hasVerifier = 1;
 }
 
+def CastOp : TransformDialectOp<"cast",
+    [TransformOpInterface, TransformEachOpTrait,
+     DeclareOpInterfaceMethods<CastOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  // TODO: temporarily fallback support for casting from PDL_Operation type.
+  let arguments = (ins AnyType:$input);
+  let results = (outs AnyType:$output);
+  let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::Operation *target,
+      ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+      ::mlir::transform::TransformState &state);
+  }];
+}
+
 def ForeachOp : TransformDialectOp<"foreach",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface, [
          "getSuccessorRegions", "getSuccessorEntryOperands"]>,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">
     ]> {
   let summary = "Executes the body for each payload op";

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

+//===- TransformTypes.h - Transform dialect types ---------------*- 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_TRANSFORMTYPES_H
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMTYPES_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class DiagnosedSilenceableFailure;
+class Operation;
+class Type;
+} // namespace mlir
+
+#include "mlir/Dialect/Transform/IR/TransformTypeInterfaces.h.inc"
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformTypes.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMTYPES_H

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

+//===- TransformTypes.td - Transform dialect types ---------*- 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_IR_TRANSFORMTYPES
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMTYPES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+
+def Transform_AnyOpType : TypeDef<Transform_Dialect, "AnyOp",
+    [DeclareTypeInterfaceMethods<TransformTypeInterface>]> {
+  let description = [{
+    Transform IR handle that can be associated with a list of arbitrary
+    Payload IR operations.
+  }];
+  let mnemonic = "any_op";
+  let assemblyFormat = "";
+}
+
+#endif  // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMTYPES

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

 add_mlir_dialect_library(MLIRTransformDialect
   TransformDialect.cpp
   TransformInterfaces.cpp
   TransformOps.cpp
+  TransformTypes.cpp
 
   DEPENDS
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRParser
   MLIRPDLDialect
   MLIRPDLInterpDialect
   MLIRRewrite
   MLIRSideEffectInterfaces
   )

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/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/IR/DialectImplementation.h"
 
 using namespace mlir;
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
 
+#ifndef NDEBUG
+void transform::detail::checkImplementsTransformTypeInterface(
+    TypeID typeID, MLIRContext *context) {
+  const auto &abstractType = AbstractType::lookup(typeID, context);
+  assert(abstractType.hasInterface(TransformTypeInterface::getInterfaceID()));
+}
+#endif // NDEBUG
+
 void transform::TransformDialect::initialize() {
-  // Using the checked version to enable the same assertions as for the ops from
-  // extensions.
+  // 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"
       >();
+  addTypesChecked<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/Transform/IR/TransformTypes.cpp.inc"
+      >();
 }
 
 void transform::TransformDialect::mergeInPDLMatchHooks(
     llvm::StringMap<PDLConstraintFunction> &&constraintFns) {
   // Steal the constraint functions form 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;
+
+  auto it = typeParsingHooks.find(keyword);
+  if (it == typeParsingHooks.end()) {
+    parser.emitError(loc) << "unknown type mnemonic: " << keyword;
+    return nullptr;
+  }
+
+  return it->getValue()(parser);
+}
+
+void transform::TransformDialect::printType(Type type,
+                                            DialectAsmPrinter &printer) const {
+  auto it = typePrintingHooks.find(type.getTypeID());
+  assert(it != typePrintingHooks.end() && "printing unknown type");
+  it->getSecond()(type, printer);
+}
+
+void transform::TransformDialect::reportDuplicateTypeRegistration(
+    StringRef mnemonic) {
+  std::string buffer;
+  llvm::raw_string_ostream msg(buffer);
+  msg << "error: extensible dialect type '" << mnemonic
+      << "' is already registered with a different implementation";
+  msg.flush();
+  llvm::report_fatal_error(StringRef(buffer));
+}
+
 #include "mlir/Dialect/Transform/IR/TransformDialectEnums.cpp.inc"

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

 //===- TransformInterfaces.cpp - Transform Dialect Interfaces -------------===//
 //
 // 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/TransformInterfaces.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DEBUG_PRINT_AFTER_ALL "transform-dialect-print-top-level-after-all"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
if (iterator != mapping.reverse.end()) {
       llvm::append_range(handles, iterator->getSecond());
       found = true;
     }
   }
 
   return success(found);
 }
 
-void transform::TransformState::setPayloadOps(Value value,
-                                              ArrayRef<Operation *> targets) {
+LogicalResult
+transform::TransformState::setPayloadOps(Value value,
   assert(value != kTopLevelValue &&
          "attempting to reset the transformation root");
 
   // TODO: this may go now
   if (value.use_empty())
-    return;
+    return success();
+
+  if (auto iface = value.getType().dyn_cast<TransformTypeInterface>()) {
+    DiagnosedSilenceableFailure result =
+        iface.checkPayload(value.getLoc(), targets);
+    if (failed(result.checkAndReport()))
+      return failure();
+  }
 
   // Setting new payload for the value without cleaning it first is a misuse of
   // the API, assert here.
   SmallVector<Operation *> storedTargets(targets.begin(), targets.end());
   Mappings &mappings = getMapping(value);
   bool inserted =
       mappings.direct.insert({value, std::move(storedTargets)}).second;
   assert(inserted && "value is already associated with another list");
   (void)inserted;
 
   for (Operation *op : targets)
     mappings.reverse[op].push_back(value);
+
+  return success();
 }
 
 void transform::TransformState::dropReverseMapping(Mappings &mappings,
                                                    Operation *op, Value value) {
   auto it = mappings.reverse.find(op);
   if (it != mappings.reverse.end())
     return;
 
   llvm::erase_value(it->getSecond(), value);
   if (it->getSecond().empty())
     mappings.reverse.erase(it);
 }
 
 void transform::TransformState::removePayloadOps(Value value) {
   Mappings &mappings = getMapping(value);
   for (Operation *op : mappings.direct[value])
     dropReverseMapping(mappings, op, value);
   mappings.direct.erase(value);
 }
 
-void transform::TransformState::updatePayloadOps(
+LogicalResult transform::TransformState::updatePayloadOps(
     Value value, function_ref<Operation *(Operation *)> callback) {
   Mappings &mappings = getMapping(value);
   auto it = mappings.direct.find(value);
   assert(it != mappings.direct.end() && "unknown handle");
   SmallVector<Operation *> &association = it->getSecond();
   SmallVector<Operation *> updated;
   updated.reserve(association.size());
 
   for (Operation *op : association) {
     dropReverseMapping(mappings, op, value);
     if (Operation *updatedOp = callback(op)) {
       updated.push_back(updatedOp);
       mappings.reverse[updatedOp].push_back(value);
     }
   }
 
+  if (auto iface = value.getType().dyn_cast<TransformTypeInterface>()) {
+    DiagnosedSilenceableFailure result =
+        iface.checkPayload(value.getLoc(), updated);
+    if (failed(result.checkAndReport()))
+      return failure();
+  }
+
   std::swap(association, updated);
+  return success();
 }
 
 void transform::TransformState::recordHandleInvalidationOne(
     OpOperand &handle, Operation *payloadOp, Value otherHandle) {
   ArrayRef<Operation *> potentialAncestors = getPayloadOps(handle.get());
   // If the op is associated with invalidated handle, skip the check as it
   // may be reading invalid IR.
   if (invalidatedHandles.count(otherHandle))
if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
       removePayloadOps(target.get());
     }
   }
 
   for (OpResult result : transform->getResults()) {
     assert(result.getDefiningOp() == transform.getOperation() &&
            "payload IR association for a value other than the result of the "
            "current transform op");
-    setPayloadOps(result, results.get(result.getResultNumber()));
+    if (failed(setPayloadOps(result, results.get(result.getResultNumber()))))
+      return DiagnosedSilenceableFailure::definiteFailure();
   }
 
   printOnFailureRAII.release();
   DEBUG_WITH_TYPE(DEBUG_PRINT_AFTER_ALL, {
     DBGS() << "Top-level payload:\n";
     getTopLevel()->print(llvm::dbgs());
   });
   return result;
 }
 
 //===----------------------------------------------------------------------===//
 // TransformState::Extension
 //===----------------------------------------------------------------------===//
 
 transform::TransformState::Extension::~Extension() = default;
 
 LogicalResult
 transform::TransformState::Extension::replacePayloadOp(Operation *op,
                                                        Operation *replacement) {
   SmallVector<Value> handles;
   if (failed(state.getHandlesForPayloadOp(op, handles)))
     return failure();
 
   for (Value handle : handles) {
+    LogicalResult result =
         state.updatePayloadOps(handle, [&](Operation *current) {
           return current == op ? replacement : current;
         });
+    if (failed(result))
+      return failure();
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TransformResults
 //===----------------------------------------------------------------------===//
 
transform::TransformResults::get(unsigned resultNumber) const {
   assert(segments[resultNumber].data() != nullptr && "querying unset results");
   return segments[resultNumber];
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 
-void transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
+LogicalResult transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
     TransformState &state, Operation *op, Region &region) {
   SmallVector<Operation *> targets;
   if (op->getNumOperands() != 0)
     llvm::append_range(targets, state.getPayloadOps(op->getOperand(0)));
   else
     targets.push_back(state.getTopLevel());
 
-  state.mapBlockArguments(region.front().getArgument(0), targets);
+  return state.mapBlockArguments(region.front().getArgument(0), targets);
 }
 
 LogicalResult
 transform::detail::verifyPossibleTopLevelTransformOpTrait(Operation *op) {
   // Attaching this trait without the interface is a misuse of the API, but it
   // cannot be caught via a static_assert because interface registration is
   // dynamic.
   assert(isa<TransformOpInterface>(op) &&

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

 #include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/OpImplementation.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/Support/Debug.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
 
 using namespace mlir;
 
for (Region &reg : getAlternatives()) {
     // the transformation from accessing any IR outside the scope.
     auto scope = state.make_region_scope(reg);
     auto clones = llvm::to_vector(
         llvm::map_range(originals, [](Operation *op) { return op->clone(); }));
     auto deleteClones = llvm::make_scope_exit([&] {
       for (Operation *clone : clones)
         clone->erase();
     });
-    state.mapBlockArguments(reg.front().getArgument(0), clones);
+    if (failed(state.mapBlockArguments(reg.front().getArgument(0), clones)))
+      return DiagnosedSilenceableFailure::definiteFailure();
 
     bool failed = false;
     for (Operation &transform : reg.front().without_terminator()) {
       DiagnosedSilenceableFailure result =
           state.applyTransform(cast<TransformOpInterface>(transform));
       if (result.isSilenceableFailure()) {
         LLVM_DEBUG(DBGS() << "alternative failed: " << result.getMessage()
                           << "\n");
LogicalResult transform::AlternativesOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ForeachOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
+transform::CastOp::applyToOne(Operation *target,
+                              SmallVectorImpl<Operation *> &results,
+                              transform::TransformState &state) {
+  results.push_back(target);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::CastOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsPayload(effects);
+  consumesHandle(getInput(), effects);
+  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 ty.isa<pdl::OperationType, transform::TransformTypeInterface>();
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// ForeachOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
 transform::ForeachOp::apply(transform::TransformResults &results,
                             transform::TransformState &state) {
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
   SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});
 
   for (Operation *op : payloadOps) {
     auto scope = state.make_region_scope(getBody());
-    state.mapBlockArguments(getIterationVariable(), {op});
+    if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
+      return DiagnosedSilenceableFailure::definiteFailure();
 
     // Execute loop body.
     for (Operation &transform : getBody().front().without_terminator()) {
       DiagnosedSilenceableFailure result = state.applyTransform(
           cast<transform::TransformOpInterface>(transform));
       if (!result.succeeded())
         return result;
     }

 // SequenceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::SequenceOp::apply(transform::TransformResults &results,
                              transform::TransformState &state) {
   // Map the entry block argument to the list of operations.
   auto scope = state.make_region_scope(*getBodyBlock()->getParent());
-  mapBlockArguments(state);
+  if (failed(mapBlockArguments(state)))
+    return DiagnosedSilenceableFailure::definiteFailure();
 
   // Apply the sequenced ops one by one.
   for (Operation &transform : getBodyBlock()->without_terminator()) {
     DiagnosedSilenceableFailure result =
         state.applyTransform(cast<TransformOpInterface>(transform));
     if (result.isDefiniteFailure())
       return result;
 
for (Operation &nested : getBody().front()) {
     }
   }
 
   state.addExtension<PatternApplicatorExtension>(getOperation());
   auto guard = llvm::make_scope_exit(
       [&]() { state.removeExtension<PatternApplicatorExtension>(); });
 
   auto scope = state.make_region_scope(getBody());
-  mapBlockArguments(state);
+  if (failed(mapBlockArguments(state)))
+    return DiagnosedSilenceableFailure::definiteFailure();
   return state.applyTransform(transformOp);
 }
 
 LogicalResult transform::WithPDLPatternsOp::verify() {
   Block *body = getBodyBlock();
   Operation *topLevelOp = nullptr;
   for (Operation &op : body->getOperations()) {
     if (isa<pdl::PatternOp>(op))

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

+//===- TransformTypes.cpp - Transform Dialect Type Definitions ------------===//
+//
+// 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/TransformTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/Types.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Compiler.h"
+
+using namespace mlir;
+
+#include "mlir/Dialect/Transform/IR/TransformTypeInterfaces.cpp.inc"
+
+// 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);
+LLVM_ATTRIBUTE_UNUSED static LogicalResult
+generatedTypePrinter(Type def, AsmPrinter &printer);
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformTypes.cpp.inc"
+
+DiagnosedSilenceableFailure
+transform::AnyOpType::checkPayload(Location loc,
+                                   ArrayRef<Operation *> payload) const {
+  return DiagnosedSilenceableFailure::success();
+}

mlir/test/Dialect/Transform/ops.mlir

 // CHECK: transform.sequence
 // CHECK: foreach
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   transform.foreach %arg0 {
   ^bb1(%arg1: !pdl.operation):
   }
 }
+
+// CHECK: transform.sequence
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  // CHECK: cast %{{.*}} : !pdl.operation to !transform.any_op
+  %0 = cast %arg0: !pdl.operation to !transform.any_op
+}

mlir/test/Dialect/Transform/test-dialect-injection.mlir

 // RUN: mlir-opt %s | FileCheck %s
 
-// These ops are defined by a test extension but should be okay to roundtrip.
+// These types and ops are defined by a test extension but should be okay to
+// roundtrip.
 
 // CHECK: transform.test_transform_op
 transform.test_transform_op
 
 // CHECK: = transform.test_produce_param_or_forward_operand 42 {foo = "bar"}
 %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
 
 // CHECK: transform.test_consume_operand_if_matches_param_or_fail %{{.*}}[42]
 transform.test_consume_operand_if_matches_param_or_fail %0[42]
+
+// Ensure that the extension type is roundtripped correctly.
+// CHECK: transform.cast %{{.*}} : !pdl.operation to !transform.test_dialect_op
+%1 = transform.cast %0: !pdl.operation to !transform.test_dialect_op

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

^bb1(%fun: !pdl.operation):
   %muli = transform.structured.match ops{["arith.muli"]} in %fun
   %h:2 = split_handles %muli in [2]
   // expected-remark @below {{1}}
   transform.test_print_number_of_associated_payload_ir_ops %h#0
   %muli_2 = transform.structured.match ops{["arith.muli"]} in %fun
   // Silenceable failure and all handles are now empty.
   %h_2:3 = split_handles %muli_2 in [3]
   // expected-remark @below {{0}}
   transform.test_print_number_of_associated_payload_ir_ops %h_2#0
 }
+
+// -----
+
+"test.some_op"() : () -> ()
+"other_dialect.other_op"() : () -> ()
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @some : benefit(1) {
+    %0 = pdl.operation "test.some_op"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @some in %arg1
+    %2 = transform.cast %0 : !pdl.operation to !transform.test_dialect_op
+    transform.cast %2 : !transform.test_dialect_op to !pdl.operation
+  }
+}
+
+// -----
+
+"test.some_op"() : () -> ()
+"other_dialect.other_op"() : () -> ()
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @other : benefit(1) {
+    %0 = pdl.operation "other_dialect.other_op"
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @other in %arg1
+    // expected-error @below {{expected the payload operation to belong to the 'test' dialect}}
+    %2 = transform.cast %0 : !pdl.operation to !transform.test_dialect_op
+    transform.cast %2 : !transform.test_dialect_op to !pdl.operation
+  }
+}

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

 set(LLVM_TARGET_DEFINITIONS TestTransformDialectExtension.td)
 mlir_tablegen(TestTransformDialectExtension.h.inc -gen-op-decls)
 mlir_tablegen(TestTransformDialectExtension.cpp.inc -gen-op-defs)
+mlir_tablegen(TestTransformDialectExtensionTypes.h.inc -gen-typedef-decls -typedefs-dialect=transform)
+mlir_tablegen(TestTransformDialectExtensionTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=transform)
 add_public_tablegen_target(MLIRTestTransformDialectExtensionIncGen)
 
 add_mlir_library(MLIRTestTransformDialect
   TestTransformDialectExtension.cpp
   TestTransformDialectInterpreter.cpp
 
   EXCLUDE_FROM_LIBMLIR
 

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTTRANSFORMDIALECTEXTENSION_H
 #define MLIR_TESTTRANSFORMDIALECTEXTENSION_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 class DialectRegistry;
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "TestTransformDialectExtension.h.inc"
 
+#define GET_TYPEDEF_CLASSES
+#include "TestTransformDialectExtensionTypes.h.inc"
+
 namespace test {
 /// Registers the test extension to the Transform dialect.
 void registerTestTransformDialectExtension(::mlir::DialectRegistry &registry);
 } // namespace test
 
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_H

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/IR/OpImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Compiler.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple transform op defined outside of the dialect. Just emits a remark when
 /// applied. This op is defined in C++ to test that C++ definitions also work
 /// for op injection into the Transform dialect.
 class TestTransformOp

 
 DiagnosedSilenceableFailure
 mlir::test::TestCopyPayloadOp::apply(transform::TransformResults &results,
                                      transform::TransformState &state) {
   results.set(getCopy().cast<OpResult>(), state.getPayloadOps(getHandle()));
   return DiagnosedSilenceableFailure::success();
 }
 
+DiagnosedSilenceableFailure mlir::transform::TestDialectOpType::checkPayload(
+    Location loc, ArrayRef<Operation *> payload) const {
+  if (payload.empty())
+    return DiagnosedSilenceableFailure::success();
+
+  for (Operation *op : payload) {
+    if (op->getName().getDialectNamespace() != "test") {
+      Diagnostic diag(loc, DiagnosticSeverity::Error);
+      diag << "expected the payload operation to belong to the 'test' dialect";
+      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    }
+  }
+
+  return DiagnosedSilenceableFailure::success();
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:
   using Base::Base;
 
   void init() {
     declareDependentDialect<pdl::PDLDialect>();
     registerTransformOps<TestTransformOp,
                          TestTransformUnrestrictedOpNoInterface,
 #define GET_OP_LIST
 #include "TestTransformDialectExtension.cpp.inc"
                          >();
+    registerTypes<
+#define GET_TYPEDEF_LIST
+#include "TestTransformDialectExtensionTypes.cpp.inc"
+        >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "TestTransformDialectExtension.cpp.inc"
 
+// 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);
+LLVM_ATTRIBUTE_UNUSED static LogicalResult
+generatedTypePrinter(Type def, AsmPrinter &printer);
+
+#define GET_TYPEDEF_CLASSES
+#include "TestTransformDialectExtensionTypes.cpp.inc"
+
 void ::test::registerTestTransformDialectExtension(DialectRegistry &registry) {
   registry.addExtensions<TestTransformDialectExtension>();
 }

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

 //===- TestTransformDialectExtension.td --------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the operations that are injected into the Transform
 // dialect through the extension mechanism, as a test.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTTRANSFORMDIALECTEXTENSION_TD
 #define MLIR_TESTTRANSFORMDIALECTEXTENSION_TD
 
+include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 
+def TestTransformTestDialectHandleType
+  : TypeDef<Transform_Dialect, "TestDialectOp",
+      [DeclareTypeInterfaceMethods<TransformTypeInterface>]> {
+  let description = [{Handle pointing to an op from the Test dialect.}];
+  let mnemonic = "test_dialect_op";
+  let assemblyFormat = "";
+}
+
 def TestProduceParamOrForwardOperandOp
   : Op<Transform_Dialect, "test_produce_param_or_forward_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins 
      Arg<Optional<PDL_Operation>, "", [TransformMappingRead]>:$operand,
      OptionalAttr<I64Attr>:$parameter);
   let results = (outs 
      Res<PDL_Operation, "",

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

deps = [
         ":Transforms",
     ],
 )
 
 td_library(
     name = "TransformDialectTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/IR/*.td"]),
     deps = [
+        ":CastInterfacesTdFiles",
         ":ControlFlowInterfacesTdFiles",
         ":OpBaseTdFiles",
         ":PDLDialectTdFiles",
         ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
tbl_outs = [
             "include/mlir/Dialect/Transform/IR/TransformInterfaces.h.inc",
         ),
         (
             [
                 "-gen-op-interface-defs",
             ],
             "include/mlir/Dialect/Transform/IR/TransformInterfaces.cpp.inc",
         ),
+        (
+            [
+                "-gen-type-interface-decls",
+            ],
+            "include/mlir/Dialect/Transform/IR/TransformTypeInterfaces.h.inc",
+        ),
+        (
+            [
+                "-gen-type-interface-defs",
+            ],
+            "include/mlir/Dialect/Transform/IR/TransformTypeInterfaces.cpp.inc",
+        ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Transform/IR/TransformInterfaces.td",
     deps = [":TransformDialectTdFiles"],
 )
 
 gentbl_cc_library(
     name = "TransformDialectIncGen",
tbl_outs = [
             "include/mlir/Dialect/Transform/IR/TransformOps.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Transform/IR/TransformOps.td",
     deps = [":TransformDialectTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "TransformTypesIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-typedef-decls"],
+            "include/mlir/Dialect/Transform/IR/TransformTypes.h.inc",
+        ),
+        (
+            ["-gen-typedef-defs"],
+            "include/mlir/Dialect/Transform/IR/TransformTypes.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Transform/IR/TransformTypes.td",
+    deps = [":TransformDialectTdFiles"],
+)
+
 cc_library(
     name = "TransformDialect",
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":ControlFlowInterfaces",
         ":IR",
         ":PDLDialect",
         ":PDLInterpDialect",
         ":Rewrite",
         ":SideEffectInterfaces",
         ":Support",
         ":TransformDialectEnumsIncGen",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
         ":TransformOpsIncGen",
+        ":TransformTypesIncGen",
         "//llvm:Support",
     ],
 )
 
 td_library(
     name = "TransformDialectTransformsTdFiles",
     srcs = glob(["include/mlir/Dialect/Transform/Transforms/*.td"]),
     deps = [

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

tbl_outs = [
         (
             ["-gen-op-decls"],
             "lib/Dialect/Transform/TestTransformDialectExtension.h.inc",
         ),
         (
             ["-gen-op-defs"],
             "lib/Dialect/Transform/TestTransformDialectExtension.cpp.inc",
         ),
+        (
+            [
+                "-gen-typedef-decls",
+                "-typedefs-dialect=transform",
+            ],
+            "lib/Dialect/Transform/TestTransformDialectExtensionTypes.h.inc",
+        ),
+        (
+            [
+                "-gen-typedef-defs",
+                "-typedefs-dialect=transform",
+            ],
+            "lib/Dialect/Transform/TestTransformDialectExtensionTypes.cpp.inc",
+        ),
     ],
     tblgen = "//mlir:mlir-tblgen",
     td_file = "lib/Dialect/Transform/TestTransformDialectExtension.td",
     test = True,
     deps = [
         ":TransformDialectTdFiles",
         "//mlir:PDLDialectTdFiles",
         "//mlir:TransformDialectTdFiles",
     ],
 )
 
 cc_library(
     name = "TestTransformDialect",
     srcs = glob(["lib/Dialect/Transform/*.cpp"]),
     hdrs = glob(["lib/Dialect/Transform/*.h"]),
     includes = ["lib/Dialect/Transform"],
     deps = [
         ":TestTransformDialectExtensionIncGen",
+        "//llvm:Support",
         "//mlir:IR",
         "//mlir:PDLDialect",
         "//mlir:Pass",
         "//mlir:TransformDialect",
     ],
 )
 
 cc_library(