[mlir][transform] Add ApplyPatternsOp and PatternRegistry

[springerm]

Summary

Add a new transform op that applies patterns to a targeted payload op. Patterns can be registered by transform dialect extensions in a pattern registry.

Depends On: D151888

[ftynse]

Nice!

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

165

Nit: I'd take PopulatePatternsFn && here to avoid an extra constructor of std::function with dynamic allocation.

178

On a second thought, how about making this a DenseMap<StringAttr, Fn>? When verifying/applying the op, we already have a StringAttr so we will only need uniquing at extension injection time, which is likely cheaper that doing string-based search when verifying/applying (which happens more often).

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

142

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

191

I suppose this may also need to reset errorCounter.

207

Nit: it's now possible to do auto &&[index, value] here. Also prefer && to const & for enumerate as that is slated for change.

227–228

Ultra-nit: this can check it != patterns.end() below.

454

Nit: braces.

462

It's preferable to explicitly silence the diagnostic here by calling .silence(). Optionally, we can also print the contents to the debug stream.

[springerm]

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

462

You mean silencing the error in case getFailOnPayloadReplacementNotFound() == false, right?

[springerm]

address comments

[springerm]

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

165

I kept StringRef here so there's no boilerplate at the call site. But that's why I needed the MLIRContext in TransformDialectData.

[ftynse]

LGTM

[springerm]

update

[ftynse]

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

462

Yep.

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

 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) {}
+  explicit TransformDialectDataBase(TypeID typeID, MLIRContext *ctx)
+      : typeID(typeID), ctx(ctx) {}
+
+  /// Return the MLIR context.
+  MLIRContext *getContext() const { return ctx; }
 
 private:
   /// The type ID of the subclass.
   const TypeID typeID;
+
+  /// The MLIR context.
+  MLIRContext *ctx;
 };
 } // 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>()) {}
+  TransformDialectData(MLIRContext *ctx)
+      : TransformDialectDataBase(TypeID::get<DerivedTy>(), ctx) {}
 };
 
 #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);

 
 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>());
+  auto emplaced =
+      extraData.try_emplace(typeID, std::make_unique<DataTy>(getContext()));
   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/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/Bytecode/BytecodeOpInterface.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/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 namespace mlir {
 namespace transform {
+class ApplyPatternsOp;
+
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 
 /// A builder function that populates the body of a SequenceOp.
 using SequenceBodyBuilderFn = ::llvm::function_ref<void(
     ::mlir::OpBuilder &, ::mlir::Location, ::mlir::BlockArgument)>;
 using SequenceBodyBuilderArgsFn =
     ::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location,
private:
   void notifyOperationRemoved(Operation *op) override;
 
   void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
 
   /// The transform op in which this TrackingListener is used.
   TransformOpInterface transformOp;
 };
 
+/// A specialized listener that keeps track of cases in which no replacement
+/// payload could be found. The error state of this listener must be checked
+/// before the end of its lifetime.
+class ErrorCheckingTrackingListener : public TrackingListener {
+public:
+  using transform::TrackingListener::TrackingListener;
+
+  ~ErrorCheckingTrackingListener() override;
+
+  /// Check and return the current error state of this listener. Afterwards,
+  /// resets the error state to "success".
+  DiagnosedSilenceableFailure checkAndResetError();
+
+  /// Return "true" if this tracking listener had a failure.
+  bool failed() const;
+
+protected:
+  void notifyPayloadReplacementNotFound(Operation *op,
+                                        ValueRange values) override;
+
+private:
+  /// The error state of this listener. "Success" indicates that no error
+  /// happened so far.
+  DiagnosedSilenceableFailure status = DiagnosedSilenceableFailure::success();
+
+  /// The number of errors that have been encountered.
+  int64_t errorCounter = 0;
+};
+
+/// The PatternRegistry stores callbacks to functions that populate a
+/// `RewritePatternSet`. Registered patterns can be applied with the
+/// "transform.apply_patterns" op.
+class PatternRegistry : public TransformDialectData<PatternRegistry> {
+public:
+  PatternRegistry(MLIRContext *ctx) : TransformDialectData(ctx), builder(ctx) {}
+
+  /// A function that populates a `RewritePatternSet`.
+  using PopulatePatternsFn = std::function<void(RewritePatternSet &)>;
+
+  /// Registers patterns with the specified identifier. The identifier should
+  /// be prefixed with the dialect to which the patterns belong.
+  void registerPatterns(StringRef identifier, PopulatePatternsFn &&fn);
+
+protected:
+  friend class ApplyPatternsOp;
+
+  /// Returns "true" if patterns are registered with the specified identifier.
+  bool hasPatterns(StringAttr identifier) const;
+
+  /// Populates the given pattern set with the specified patterns.
+  void populatePatterns(StringAttr identifier,
+                        RewritePatternSet &patternSet) const;
+
+private:
+  /// A builder for creating StringAttrs.
+  Builder builder;
+
+  DenseMap<StringAttr, PopulatePatternsFn> patterns;
+};
+
 } // namespace transform
 } // namespace mlir
 
+MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::transform::PatternRegistry)
+
 #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

let arguments = (ins TransformHandleTypeInterface:$target,
                        Optional<TransformParamTypeInterface>:$param);
   let results = (outs);
 
   let assemblyFormat =
     "$target $name attr-dict (`=` $param^)?"
     "`:` type($target) (`,` type($param)^)?";
 }
 
+def ApplyPatternsOp : TransformDialectOp<"apply_patterns",
+    [TransformOpInterface, TransformEachOpTrait,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Greedily applies patterns to the body of the targeted op";
+  let description = [{
+    This transform greedily applies the specified patterns to the body of the
+    targeted op until a fixpoint was reached. Patterns are not applied to the
+    targeted op itself.
+
+    Only patterns that were registered in the transform dialect's
+    `PatternRegistry` are available. Additional patterns can be registered as
+    part of transform dialect extensions.
+
+    This transform only reads the target handle and modifies the payload. If a
+    pattern erases or replaces a tracked op, the mapping is updated accordingly.
+
+    Only replacements via `RewriterBase::replaceOp` or `replaceOpWithNewOp` are
+    considered "payload op replacements". Furthermore, only if the replacement
+    values are defined by the same op and that op has the same type as the
+    original op, the mapping is updated. Otherwise, this transform fails
+    silently unless `fail_on_payload_replacement_not_found` is set to "false".
+    More details can be found at the documentation site of `TrackingListener`.
+
+    This transform also fails silently if the pattern application did not
+    converge within the default number of iterations/rewrites of the greedy
+    pattern rewrite driver.
+  }];
+
+  let arguments = (ins
+    TransformHandleTypeInterface:$target, ArrayAttr:$patterns,
+    DefaultValuedAttr<BoolAttr, "true">:$fail_on_payload_replacement_not_found);
+  let results = (outs);
+  let assemblyFormat = "$patterns `to` $target attr-dict `:` type($target)";
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::Operation *target,
+      ::mlir::transform::ApplyToEachResultList &results,
+      ::mlir::transform::TransformState &state);
+  }];
+}
+
 def CastOp : TransformDialectOp<"cast",
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<CastOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$input);
   let results = (outs TransformHandleTypeInterface:$output);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 

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

 
 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:
+  PDLMatchHooks(MLIRContext *ctx) : TransformDialectData(ctx) {}
+
   /// 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> &

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

 #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/Transforms/GreedyPatternRewriteDriver.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 "] ")
 
 #define DEBUG_TYPE_MATCHER "transform-matcher"
 #define DBGS_MATCHER() (llvm::dbgs() << "[" DEBUG_TYPE_MATCHER "] ")
 #define DEBUG_MATCHER(x) DEBUG_WITH_TYPE(DEBUG_TYPE_MATCHER, x)
 
 using namespace mlir;
 
+MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::transform::PatternRegistry)
+
 static ParseResult parseSequenceOpOperands(
     OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &root,
     Type &rootType,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &extraBindings,
     SmallVectorImpl<Type> &extraBindingTypes);
 static void printSequenceOpOperands(OpAsmPrinter &printer, Operation *op,
                                     Value root, Type rootType,
                                     ValueRange extraBindings,
void transform::TrackingListener::notifyOperationReplaced(
   Operation *replacement = findReplacementOp(op, newValues);
   // If the op is tracked but no replacement op was found, send a
   // notification.
   if (!replacement)
     notifyPayloadReplacementNotFound(op, newValues);
   (void)replacePayloadOp(op, replacement);
 }
 
+transform::ErrorCheckingTrackingListener::~ErrorCheckingTrackingListener() {
+  // The state of the ErrorCheckingTrackingListener must be checked and reset
+  // if there was an error. This is to prevent errors from accidentally being
+  // missed.
+  assert(status.succeeded() && "listener state was not checked");
+}
+
+DiagnosedSilenceableFailure
+transform::ErrorCheckingTrackingListener::checkAndResetError() {
+  DiagnosedSilenceableFailure s = std::move(status);
+  status = DiagnosedSilenceableFailure::success();
+  errorCounter = 0;
+  return s;
+}
+
+bool transform::ErrorCheckingTrackingListener::failed() const {
+  return !status.succeeded();
+}
+
+void transform::ErrorCheckingTrackingListener::notifyPayloadReplacementNotFound(
+    Operation *op, ValueRange values) {
+  if (status.succeeded()) {
+    status = emitSilenceableFailure(
+        getTransformOp(), "tracking listener failed to find replacement op");
+  }
+
+  status.attachNote(op->getLoc()) << "[" << errorCounter << "] replaced op";
+  for (auto &&[index, value] : llvm::enumerate(values))
+    status.attachNote(value.getLoc())
+        << "[" << errorCounter << "] replacement value " << index;
+
+  ++errorCounter;
+}
+
+//===----------------------------------------------------------------------===//
+// PatternRegistry
+//===----------------------------------------------------------------------===//
+
+void transform::PatternRegistry::registerPatterns(StringRef identifier,
+                                                  PopulatePatternsFn &&fn) {
+  StringAttr attr = builder.getStringAttr(identifier);
+  assert(!patterns.contains(attr) && "patterns identifier is already in use");
+  patterns.try_emplace(attr, std::move(fn));
+}
+
+void transform::PatternRegistry::populatePatterns(
+    StringAttr identifier, RewritePatternSet &patternSet) const {
+  auto it = patterns.find(identifier);
+  assert(it != patterns.end() && "patterns not registered in registry");
+  it->second(patternSet);
+}
+
+bool transform::PatternRegistry::hasPatterns(StringAttr identifier) const {
+  return patterns.contains(identifier);
+}
+
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
 OperandRange transform::AlternativesOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {
   if (index && getOperation()->getNumOperands() == 1)
     return getOperation()->getOperands();

 void transform::AnnotateOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getTarget(), effects);
   onlyReadsHandle(getParam(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// ApplyPatternsOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::ApplyPatternsOp::applyToOne(Operation *target,
+                                       ApplyToEachResultList &results,
+                                       transform::TransformState &state) {
+  // Gather all specified patterns.
+  MLIRContext *ctx = target->getContext();
+  RewritePatternSet patterns(ctx);
+  const auto &registry = getContext()
+                             ->getLoadedDialect<transform::TransformDialect>()
+                             ->getExtraData<transform::PatternRegistry>();
+  for (Attribute attr : getPatterns())
+    registry.populatePatterns(attr.cast<StringAttr>(), patterns);
+
+  // Configure the GreedyPatternRewriteDriver.
+  ErrorCheckingTrackingListener listener(state, *this);
+  GreedyRewriteConfig config;
+  config.listener = &listener;
+
+  // Manually gather list of ops because the other GreedyPatternRewriteDriver
+  // overloads only accepts ops that are isolated from above. This way, patterns
+  // can be applied to ops that are not isolated from above.
+  SmallVector<Operation *> ops;
+  target->walk([&](Operation *nestedOp) {
+    if (target != nestedOp)
+      ops.push_back(nestedOp);
+  });
+  LogicalResult result =
+      applyOpPatternsAndFold(ops, std::move(patterns), config);
+  // A failure typically indicates that the pattern application did not
+  // converge.
+  if (failed(result)) {
+    return emitSilenceableFailure(target)
+           << "greedy pattern application failed";
+  }
+
+  // Check listener state for tracking errors.
+  if (listener.failed()) {
+    DiagnosedSilenceableFailure status = listener.checkAndResetError();
+    if (getFailOnPayloadReplacementNotFound())
+      return status;
+    (void)status.silence();
+  }
+
+  return DiagnosedSilenceableFailure::success();
+}
+
+LogicalResult transform::ApplyPatternsOp::verify() {
+  const auto &registry = getContext()
+                             ->getLoadedDialect<transform::TransformDialect>()
+                             ->getExtraData<transform::PatternRegistry>();
+  for (Attribute attr : getPatterns()) {
+    auto strAttr = attr.dyn_cast<StringAttr>();
+    if (!strAttr)
+      return emitOpError() << "expected " << getPatternsAttrName()
+                           << " to be an array of strings";
+    if (!registry.hasPatterns(strAttr))
+      return emitOpError() << "patterns not registered: " << strAttr.strref();
+  }
+  return success();
+}
+
+void transform::ApplyPatternsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();

mlir/test/Dialect/Transform/ops-invalid.mlir

module attributes { transform.with_named_sequence } {
 
   transform.sequence failures(propagate) {
   ^bb0(%root: !transform.any_op):
     // expected-error @below {{'transform.foreach_match' op does not expect matcher symbol to consume its operand}}
     transform.foreach_match in %root
       @match -> @action : (!transform.any_op) -> !transform.any_op
   }
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{patterns not registered: transform.invalid_pattern_identifier}}
+  transform.apply_patterns ["transform.invalid_pattern_identifier"] to %arg0 : !transform.any_op
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{expected "patterns" to be an array of strings}}
+  transform.apply_patterns [3, 9] to %arg0 : !transform.any_op
+}

mlir/test/Dialect/Transform/test-pattern-application.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: func @update_tracked_op_mapping()
+//       CHECK:   "test.container"() ({
+//       CHECK:     %0 = "test.foo"() {annotated} : () -> i32
+//       CHECK:   }) : () -> ()
+func.func @update_tracked_op_mapping() {
+  "test.container"() ({
+    %0 = "test.foo"() {replace_with_new_op = "test.foo"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  // Add an attribute to %1, which is now mapped to a new op.
+  transform.annotate %1 "annotated" : !transform.any_op
+}
+
+// -----
+
+func.func @replacement_op_not_found() {
+  "test.container"() ({
+    // expected-note @below {{[0] replaced op}}
+    // expected-note @below {{[0] replacement value 0}}
+    %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-error @below {{tracking listener failed to find replacement op}}
+  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  // %1 must be used in some way. If no replacement payload op could be found,
+  // an error is thrown only if the handle is not dead.
+  transform.annotate %1 "annotated" : !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: func @replacement_op_for_dead_handle_not_found()
+//       CHECK:   "test.container"() ({
+//       CHECK:     %0 = "test.bar"() : () -> i32
+//       CHECK:   }) : () -> ()
+func.func @replacement_op_for_dead_handle_not_found() {
+  "test.container"() ({
+    %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // No error because %1 is dead.
+  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: func @replacement_op_not_found_silenced()
+//       CHECK:   "test.container"() ({
+//       CHECK:     %0 = "test.bar"() : () -> i32
+//       CHECK:   }) : () -> ()
+func.func @replacement_op_not_found_silenced() {
+  "test.container"() ({
+    %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {fail_on_payload_replacement_not_found = false}: !transform.any_op
+  transform.annotate %1 "annotated" : !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: func @patterns_apply_only_to_target_body()
+//       CHECK:   %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> i32
+func.func @patterns_apply_only_to_target_body() {
+  %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: func @erase_tracked_op()
+//       CHECK:   "test.container"() ({
+//  CHECK-NEXT:   ^bb0:
+//  CHECK-NEXT:   }) : () -> ()
+func.func @erase_tracked_op() {
+  "test.container"() ({
+    // expected-remark @below {{matched op}}
+    %0 = "test.erase_op"() {replace_with_new_op = "test.foo"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.erase_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.test_print_remark_at_operand %1, "matched op" : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.test_print_remark_at_operand %1, "op was deleted" : !transform.any_op
+}

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

for (Operation *op : state.getPayloadOps(getIn())) {
     }
   }
 
   emitRemark() << numIterations << " iterations";
   return DiagnosedSilenceableFailure::success();
 }
 
 namespace {
+// Test pattern to replace an operation with a new op.
+class ReplaceWithNewOp : public RewritePattern {
+public:
+  ReplaceWithNewOp(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    auto newName = op->getAttrOfType<StringAttr>("replace_with_new_op");
+    if (!newName)
+      return failure();
+    Operation *newOp = rewriter.create(
+        op->getLoc(), OperationName(newName, op->getContext()).getIdentifier(),
+        op->getOperands(), op->getResultTypes());
+    rewriter.replaceOp(op, newOp->getResults());
+    return success();
+  }
+};
+
+// Test pattern to erase an operation.
+class EraseOp : public RewritePattern {
+public:
+  EraseOp(MLIRContext *context)
+      : RewritePattern("test.erase_op", /*benefit=*/1, context) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+void populateTestPatterns(RewritePatternSet &patterns) {
+  patterns.insert<ReplaceWithNewOp, EraseOp>(patterns.getContext());
+}
+
 /// 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;
#include "TestTransformDialectExtensionTypes.cpp.inc"
     };
 
     addDialectDataInitializer<transform::PDLMatchHooks>(
         [&](transform::PDLMatchHooks &hooks) {
           llvm::StringMap<PDLConstraintFunction> constraints;
           constraints.try_emplace("verbose_constraint", verboseConstraint);
           hooks.mergeInPDLMatchHooks(std::move(constraints));
         });
+
+    addDialectDataInitializer<transform::PatternRegistry>(
+        [&](transform::PatternRegistry &registry) {
+          registry.registerPatterns("transform.test", populateTestPatterns);
+        });
   }
 };
 } // 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);