[mlir][transform] Add TransformRewriter

[springerm]

Summary

All apply functions now have a TransformRewriter & parameter. This rewriter should be used to modify the IR. It has a TrackingListener attached and updates the internal handle-payload mappings based on rewrites.

Implementations no longer need to create their own TrackingListener and IRRewriter. Error checking is integrated into applyTransform.

Additional API will be added to TransformRewriter in subsequent revisions. This revision just adds an "empty" TransformRewriter class and updates all apply implementations.

Depends On: D152446

[springerm]

rebase

[ftynse]

Documentation needs to be updated, in particular the tutorial and potentially the top-level Dialects/Transform.md description to account for this.

I am not entirely sure what contract do we want to have between the rewriter and the apply implementation. While it is tempting to have a blanket rule of the style "all IR modifications in apply must go through the rewriter, everything else is UB", most existing transformations don't follow that. Also, the replacement tracking is limited and we may create cases where replacement failed, but it wasn't actually necessary and the transform would have succeeded without it, i.e., without using the rewriter. What are your thoughts on this?

mlir/examples/transform/Ch2/lib/MyExtension.cpp

86–87

Please update the actual tutorial as well, it lives in mlir/docs/Tutorials/transform/. It should remain in sync with these sources.

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

53–54

Nit: namespace prefixing is not necessary here.

94–95

Ditto.

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

973

Please add doxygen documentation for this.

974

Why does this need to be an extension? So it can access TransformState::replacePayloadOp? If that's the only reason, let's just make it a friend of TransformState or even make replacePayloadOp public with a disclaimer that it's a low-level API and that reasonable clients should use the rewriter instead. Extensions were meant to carry extra state, but also were separating the ugliness of the early attempts at tracking. Since tracking becomes more first-class, I don't see the reason of jumping through hoops with this.

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

52

Update the documentation above to indicate that the implementation is expected to use the provided rewriter. It is also worth mentioning that its insertion point is not set and the implementation is in charge of setting it correctly given the payloads.

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

920

At this point, result may be a definite failure and we cannot attach notes to those (the diagnostic has been already reported). I suppose we can report a second error in this case.

1197

Nit: braces are required if there is a nested control flow construct.

1312

Nit: there is no op in the function below.

1316

Nit: I'd rather call this UB than invalid IR. Invalid IR suggests that it can be checked by the verifier but it cannot.

[springerm]

addres comments

[springerm]

In D152427#4407902, @ftynse wrote:

I am not entirely sure what contract do we want to have between the rewriter and the apply implementation. While it is tempting to have a blanket rule of the style "all IR modifications in apply must go through the rewriter, everything else is UB",

I was aiming for "everything has to go through the rewriter". Otherwise, it will be difficult to understand when a rewriter must be used and when it can be done without.

We have a similar situation with RewritePatterns. Creating a new with the rewriter, but then modifying it in-place (without updateRootInPlace) is actually fine (for the GreedyPatternRewriteDriver at least), but I always discouraged that when talking to people, because it makes the API more complicated and they may easily miss a required updateRootInPlace.

most existing transformations don't follow that.

Yes, I would update these gradually. (I think we have the same issue with some SCF RewritePatterns that call some util functions that just create a new OpBuilder...)

Also, the replacement tracking is limited and we may create cases where replacement failed, but it wasn't actually necessary and the transform would have succeeded without it, i.e., without using the rewriter. What are your thoughts on this?

I have a WIP change that allows users to install a custom replacement op finder:

rewriter.withPayloadReplacementHandler([&]() {
  //rewrite
}, /*handler=*/ [](ValueRange values) -> FailureOr<Operation *> {
 ...
});

There's also a default handler that silences all replacement errors:

rewriter.silencePayloadReplacementErrors([&]() { ... });

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

974

Ah yes this does not have to be an extension. The listener is an extension already.

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

1312

This refers to op of notifyOperationReplaced.

[ftynse]

Okay, my overall concern here is we may be breaking things that currently work in subtle ways. Both declaring the direct IR mutation as UB and by doing more tracking than before potentially aborting the transformation. If we can mitigate these two by (1) giving some slack for transformations/downstreams to adapt to using the rewriter + posting a PSA and (2) silencing replacement failures by default, without any action from the user, and separately discussing a switch in that behavior with another PSA, this should be fine.

mlir/docs/Tutorials/transform/Ch1.md

%biased = linalg.elemwise_binary { fun = #linalg.binary_fn<add> }
             ^
 test/Examples/transform/Ch1/invalidation-2.mlir:24:13: note: nested payload op
   %matmul = linalg.matmul ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)
 ```
 
 Note that the “add” elementwise operation is indicated as payload ancestor because it was used to produce the tile loop, and the loop therefore has its location.
 
 Finally, we would like to replace the call to the outlined function with a call to the microkernel. Unfortunately, the Transform dialect doesn’t have support for this transformation (and cannot have if the call is rewritten to a custom, out-of-tree operation). Therefore, we need to define new transform operations. The next chapters will describe how this can be done.
+
+## Tracking IR Modifications
+
+The transform dialect automatically tracks all IR changes that are made as part
+of transform ops. (Implementations must use the provided rewriter to modify IR.)
+If a payload op is erased, it is automatically removed from all handles that it
+is currently associated with. If a payload op is replaced, the transform dialect
+tries to find the replacement op and updates all handles accordingly. If a
+multi-result op is replaced with values that are defined by multiple ops, or if
+an op is replaced with an op of a different type, an error is produced. This is
+because it is unclear whether the direct replacements actually represent the
+computation of the original op. There are ways to customize this behavior. More
+details can be found at the documentation of `transform::TrackingListener`.

mlir/docs/Tutorials/transform/Ch2.md

def ChangeCallTargetOp : Op<Transform_Dialect, "my.change_call_target",
   // The results are empty as the transformation does not produce any new payload.
   let results = (outs);
 
   // Provide nice syntax.
   let assemblyFormat = "$call `,` $new_target attr-dict `:` type($call)";
 }
 ```
 
-To finalize the definition of the transform operation, we need to implement the interface methods. The `TransformOpInterface` currently requires only one method – `apply` – that performs the actual transformation. It is a good practice to limit the body of the method to manipulation of the Transform dialect constructs and have the actual transformation implemented as a standalone function so it can be used from other places in the code.
-
+To finalize the definition of the transform operation, we need to implement the
+interface methods. The `TransformOpInterface` currently requires only one method
+– `apply` – that performs the actual transformation. It is a good practice to
+limit the body of the method to manipulation of the Transform dialect constructs
+and have the actual transformation implemented as a standalone function so it
+can be used from other places in the code. Similar to rewrite patterns, all IR
+must be modified with the provided rewriter.
 
 ```c++
 // In MyExtension.cpp
 
 // Implementation of our transform dialect operation.
 // This operation returns a tri-state result that can be one of:
 // - success when the transformation succeeded;
-// - definite failure when the transformation failed in such a way that following 
-//   transformations are impossible or undesirable, typically it could have left payload 
-//   IR in an invalid state; it is expected that a diagnostic is emitted immediately 
-//   before returning the definite error;
-// - silenceable failure when the transformation failed but following transformations 
-//   are still applicable, typically this means a precondition for the transformation is 
-//   not satisfied and the payload IR has not been modified.
-// The silenceable failure additionally carries a Diagnostic that can be emitted to the 
-// user.
-::mlir::DiagnosedSilenceableFailure ChangeCallTargetOp::apply(
-    // The list of payload IR entities that will be associated with the transform IR 
-    // values defined by this transform operation. In this case, it can remain empty as 
-    // there are no results.
+// - definite failure when the transformation failed in such a way that
+//   following transformations are impossible or undesirable, typically it could
+//   have left payload IR in an invalid state; it is expected that a diagnostic
+//   is emitted immediately before returning the definite error;
+// - silenceable failure when the transformation failed but following
+//   transformations are still applicable, typically this means a precondition
+//   for the transformation is not satisfied and the payload IR has not been
+//   modified. The silenceable failure additionally carries a Diagnostic that
+//   can be emitted to the user.
+::mlir::DiagnosedSilenceableFailure mlir::transform::ChangeCallTargetOp::apply(
+    // The rewriter that should be used when modifying IR.
+    ::mlir::transform::TransformRewriter &rewriter,
+    // The list of payload IR entities that will be associated with the
+    // transform IR values defined by this transform operation. In this case, it
+    // can remain empty as there are no results.
     ::mlir::transform::TransformResults &results,
-    // The transform application state. This object can be used to query the current 
-    // associations between transform IR values and payload IR entities. It can also 
-    // carry additional user-defined state.
+    // The transform application state. This object can be used to query the
+    // current associations between transform IR values and payload IR entities.
+    // It can also carry additional user-defined state.
     ::mlir::transform::TransformState &state) {
 
-  // First, we need to obtain the list of payload operations that are associated with 
+  // First, we need to obtain the list of payload operations that are associated with
   // the operand handle.
   auto payload = state.getPayloadOps(getCall());
-  
-  // Then, we iterate over the list of operands and call the actual IR-mutating 
+  // Then, we iterate over the list of operands and call the actual IR-mutating
   // function. We also check the preconditions here.
   for (Operation *payloadOp : payload) {
     auto call = dyn_cast<::mlir::func::CallOp>(payloadOp);
     if (!call) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
           << "only applies to func.call payloads";
       diag.attachNote(payloadOp->getLoc()) << "offending payload";
       return diag;
     }
-    
     updateCallee(call, getNewTarget());
   }
 
   // If everything went well, return success.
   return DiagnosedSilenceableFailure::success();
 }
 ```
 

mlir/docs/Tutorials/transform/Ch3.md

 # Chapter 3: More than Simple Transform Operations
 
 ## Type Constraints and ApplyEach Trait
 
 A transform operation that applies to each payload operation individually and requires it to be of a specific kind is a repeated pattern. One can use Transform dialect types to specify the preconditions of the type. Specifically, we can change the expected operand type from the wide `TransformHandleTypeInterface` to the more narrow `Transform_ConcreteOp<"func.call">`. Furthermore, we use the `TransformEachOpTrait` trait to provide the skeleton implementation of the `apply` method that performs verification, iteration over payloads and result concatenation. The improved ODS op definition is as follows.
 
 ```tablegen
 // In MyExtension.td.
 
 // Define the new operation. By convention, prefix its name with the name of the dialect extension, "my.". The full operation name will be further prefixed with "transform.".
 def ChangeCallTargetOp : Op<Transform_Dialect, "my.change_call_target",
     // Indicate that the operation implements the required TransformOpInterface and
-    // MemoryEffectsOpInterface. Use the TransformEach trait to provide the 
+    // MemoryEffectsOpInterface. Use the TransformEach trait to provide the
     // implementation for TransformOpInterface.
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  // Provide a brief and a full description. It is recommended that the latter describes 
+  // Provide a brief and a full description. It is recommended that the latter describes
   // the effects on the operands and how the operation processes various failure modes.
   let summary = "Changes the callee of a call operation to the specified one";
   let description = [{
-    For each `func.call` payload operation associated with the handle, changes its 
+    For each `func.call` payload operation associated with the handle, changes its
     callee to be the symbol whose name is provided as an attribute to this operation.
 
-    Generates a silenceable failure if the operand is associated with payload operations 
+    Generates a silenceable failure if the operand is associated with payload operations
     that are not `func.call`.
     Only reads the operand.
   }];
 
-  // The arguments include the handle to the payload operations and the attribute that 
-  // specifies the new callee. The handle must implement TransformHandleTypeInterface.   
-  // We use a string attribute as the symbol may not exist in the transform IR so the 
-  // verification may fail. 
+  // The arguments include the handle to the payload operations and the attribute that
+  // specifies the new callee. The handle must implement TransformHandleTypeInterface.
+  // We use a string attribute as the symbol may not exist in the transform IR so the
+  // verification may fail.
   let arguments = (ins
     Transform_ConcreteOpType<"func.call">:$call,
     StrAttr:$new_target);
 
   // The results are empty as the transformation does not produce any new payload.
   let results = (outs);
 
   // Provide nice syntax.
   let assemblyFormat = "$call `,` $new_target attr-dict `:` type($call)";
 
   // Declare the function implementing the interface for a single payload operation.
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::func::CallOp call,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 ```
 
 Now, instead of defining the `apply` method with a loop, we can simply define a function that applies to an individual payload operation and the trait will take care of the rest.
 
 ```c++
 ::mlir::DiagnosedSilenceableFailure ChangeCallTargetOp::applyToOne(
-    ::mlir::func::CallOp call,,
+    ::mlir::transform::TransformRewriter &rewriter,
+    ::mlir::func::CallOp call,
     ::mlir::transform::ApplyToEachResultList &results,
     ::mlir::transform::TransformState &state) {
   // Call the actual transformation function.
   updateCallee(call, getNewTarget());
   // Indicate success.
   return DiagnosedSilenceableFailure::success();
 }
 ```
```mlir
   // Using our new operation.
   transform.my.change_call_target %casted, "microkernel" : !transform.my.call_op_interface
 ```
 
 ## Operand Consumption
 
 As an exercise, let us modify the rewriting operation to consume the operand. This would be necessary, for example, if the transformation were to rewrite the `func.call` operation into a custom operation `my.mm4`. Since the operand handle is now consumed, the operation can return a new handle to the newly produced payload operation. Otherwise, the ODS definition of the transform operation remains unchanged.
 
-
 ```tablegen
 // In MyExtension.td.
 
 // Define another transform operation.
 def CallToOp : Op<Transform_Dialect, "my.call_to_op",
      // Indicate that the operation implements the required TransformOpInterface and
-     // MemoryEffectsOpInterface. Use the TransformEach trait to provide the 
+     // MemoryEffectsOpInterface. Use the TransformEach trait to provide the
      // implementation for TransformOpInterface.
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   // Summary and description omitted for brevity.
 
   // The argument is the handle to the payload operations.
   let arguments = (ins CallOpInterfaceHandle:$call);
 
-  // The result is the handle to the payload operations produced during the 
+  // The result is the handle to the payload operations produced during the
   // transformation.
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   // Provide nice syntax.
   let assemblyFormat = "$call attr-dict `:` functional-type(inputs, outputs)";
 
   // Declare the function implementing the interface for a single payload operation.
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::CallOpInterface call,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 ```
 
 Now let’s look at the implementation of interface methods.
 
 ```c++
 // In MyExtension.cpp.
 
 ::mlir::DiagnosedSilenceableFailure CallToOp::applyToOne(
+    ::mlir::transform::TransformRewriter &rewriter,
     ::mlir::CallOpInterface call,
     ::mlir::transform::ApplyToEachResultList &results,
     ::mlir::transform::TransformState &state) {
   // Call the actual rewrite.
   Operation *rewritten = rewriteToOp(call);
 
-  // Report an error if the rewriter produced a null pointer. Note that it may have 
+  // Report an error if the rewriter produced a null pointer. Note that it may have
   // irreversibly modified the payload IR, so we produce a definite failure.
   if (!rewritten) {
     return emitDefiniteError() << "failed to rewrite call to operation";
   }
 
-  // On success, push the resulting operation into the result list. The list is expected 
-  // to contain exactly one entity per result and per application. The handles will be 
+  // On success, push the resulting operation into the result list. The list is expected
+  // to contain exactly one entity per result and per application. The handles will be
   // associated with lists of the respective values produced by each application.
   results.push_back(rewritten);
-  
   // If everything is fine, return success.
   return DiagnosedSilenceableFailure::success();
 }
 
 void CallToOp::getEffects(
     ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
   // Indicate using side effects that the operand handle is consumed, and the
   // result handle is produced.
   consumesHandle(getCall(), effects);
   producesHandle(getTransformed(), effects);
 
   // Indicate that the payload IR is modified.
   modifiesPayload(effects);
 }
 ```
 
 The overall flow of these implementations is similar to the previous one. The application also needs to specify the resulting entities that are going to be associated with the handles it produces. Operations are required to specify the entities to associate with _all_ results on success, even if the list is empty. An assertion will be triggered if it is not the case. In case of failure, the interpreter will automatically associate all results that are not yet defined with empty lists.
 
 Note that since `applyToOne` always expects one payload entity to be associated with each result handle in each application, it cannot be used to return handles associated with empty lists for non-empty operand handles. Instead, one would use `apply` directly.
 
 ```c++
 ::mlir::DiagnosedSilenceableFailure SomeOtherOp::apply(
+    ::mlir::transform::TransformRewriter &rewriter,
     ::mlir::transform::TransformResults &results,
     ::mlir::transform::TransformState &state) {
   // ...
 
   // Associate the result `transformed` with an empty list of payload operations.
   results.set(cast<OpResult>(getTransformed()), {});
   return DiagnosedSilenceableFailure::success();
 }
def CallToOp : Op<Transform_Dialect, "my.call_to_op",
      // Indicate that the operation implements the required MemoryEffectsOpInterface.
      // Use the FunctionalStyle trait to provide the implementation for this interface.
      MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]> {
   // Summary and description omitted for brevity.
 
   // The argument is the handle to the payload operations.
   let arguments = (ins CallOpInterfaceHandle:$call);
 
-  // The result is the handle to the payload operations produced during the 
+  // The result is the handle to the payload operations produced during the
   // transformation.
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   // Provide nice syntax.
   let assemblyFormat = "$call attr-dict `:` functional-type(operands, results)";
 
   // Declare the function implementing the interface for a single payload operation.
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::CallOpInterface call,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 ```
 
 ## Appendix: Autogenerated Documentation
 
 [include "MyExtensionCh3.md"]
 

mlir/examples/transform/Ch2/lib/MyExtension.cpp

 static void updateCallee(mlir::func::CallOp call, llvm::StringRef newTarget) {
   call.setCallee(newTarget);
 }
 
 // Implementation of our transform dialect operation.
 // This operation returns a tri-state result that can be one of:
 // - success when the transformation succeeded;
 // - definite failure when the transformation failed in such a way that
-// following
-//   transformations are impossible or undesirable, typically it could have left
-//   payload IR in an invalid state; it is expected that a diagnostic is emitted
-//   immediately before returning the definite error;
+//   following transformations are impossible or undesirable, typically it could
+//   have left payload IR in an invalid state; it is expected that a diagnostic
+//   is emitted immediately before returning the definite error;
 // - silenceable failure when the transformation failed but following
-// transformations
-//   are still applicable, typically this means a precondition for the
-//   transformation is not satisfied and the payload IR has not been modified.
-// The silenceable failure additionally carries a Diagnostic that can be emitted
-// to the user.
+//   transformations are still applicable, typically this means a precondition
+//   for the transformation is not satisfied and the payload IR has not been
+//   modified. The silenceable failure additionally carries a Diagnostic that
+//   can be emitted to the user.
 ::mlir::DiagnosedSilenceableFailure mlir::transform::ChangeCallTargetOp::apply(
+    // The rewriter that should be used when modifying IR.
+    ::mlir::transform::TransformRewriter &rewriter,
     // The list of payload IR entities that will be associated with the
     // transform IR values defined by this transform operation. In this case, it
     // can remain empty as there are no results.
     ::mlir::transform::TransformResults &results,
     // The transform application state. This object can be used to query the
     // current associations between transform IR values and payload IR entities.
     // It can also carry additional user-defined state.
     ::mlir::transform::TransformState &state) {

mlir/examples/transform/Ch3/include/MyExtension.td

def ChangeCallTargetOp : Op<Transform_Dialect, "my.change_call_target",
   let results = (outs);
 
   // Provide nice syntax.
   let assemblyFormat = "$call `,` $new_target attr-dict `:` qualified(type($call))";
 
   // Declare the function implementing the interface for a single payload operation.
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::func::CallOp call,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 // Define another transform operation.
 def CallToOp : Op<Transform_Dialect, "my.call_to_op",
def CallToOp : Op<Transform_Dialect, "my.call_to_op",
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   // Provide nice syntax.
   let assemblyFormat = "$call attr-dict `:` functional-type(operands, results)";
 
   // Declare the function implementing the interface for a single payload operation.
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::CallOpInterface call,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // MY_EXTENSION

mlir/examples/transform/Ch3/lib/MyExtension.cpp

 // - silenceable failure when the transformation failed but following
 // transformations
 //   are still applicable, typically this means a precondition for the
 //   transformation is not satisfied and the payload IR has not been modified.
 // The silenceable failure additionally carries a Diagnostic that can be emitted
 // to the user.
 ::mlir::DiagnosedSilenceableFailure
 mlir::transform::ChangeCallTargetOp::applyToOne(
+    // The rewriter that should be used when modifying IR.
+    ::mlir::transform::TransformRewriter &rewriter,
     // The single payload operation to which the transformation is applied.
     ::mlir::func::CallOp call,
     // The payload IR entities that will be appended to lists associated with
     // the results of this transform operation. This list contains one entry per
     // result.
     ::mlir::transform::ApplyToEachResultList &results,
     // The transform application state. This object can be used to query the
     // current associations between transform IR values and payload IR entities.
void mlir::transform::ChangeCallTargetOp::getEffects(
   // Indicate that the payload is modified by this operation.
   modifiesPayload(effects);
 }
 
 //===---------------------------------------------------------------------===//
 // CallToOp
 //===---------------------------------------------------------------------===//
 
-static mlir::Operation *replaceCallWithOp(mlir::CallOpInterface call) {
+static mlir::Operation *replaceCallWithOp(mlir::RewriterBase &rewriter,
+                                          mlir::CallOpInterface call) {
   // Construct an operation from an unregistered dialect. This is discouraged
   // and is only used here for brevity of the overall example.
   mlir::OperationState state(call.getLoc(), "my.mm4");
   state.types.assign(call->result_type_begin(), call->result_type_end());
   state.operands.assign(call->operand_begin(), call->operand_end());
 
-  mlir::OpBuilder builder(call);
-  mlir::Operation *replacement = builder.create(state);
-  call->replaceAllUsesWith(replacement->getResults());
-  call->erase();
+  mlir::Operation *replacement = rewriter.create(state);
+  rewriter.replaceOp(call, replacement->getResults());
   return replacement;
 }
 
 // See above for the signature description.
 mlir::DiagnosedSilenceableFailure mlir::transform::CallToOp::applyToOne(
-    mlir::CallOpInterface call, mlir::transform::ApplyToEachResultList &results,
+    mlir::transform::TransformRewriter &rewriter, mlir::CallOpInterface call,
+    mlir::transform::ApplyToEachResultList &results,
     mlir::transform::TransformState &state) {
 
   // Dispatch to the actual transformation.
-  Operation *replacement = replaceCallWithOp(call);
+  Operation *replacement = replaceCallWithOp(rewriter, call);
 
   // Associate the payload operation produced by the rewrite with the result
   // handle of this transform operation.
   results.push_back(replacement);
 
   // If everything went well, return success.
   return DiagnosedSilenceableFailure::success();
 }

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td

def EmptyTensorToAllocTensorOp
 
   let arguments = (ins Transform_EmptyOp:$target);
   let results = (outs Transform_AllocTensorOp:$transformed);
 
   let assemblyFormat = "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::tensor::EmptyOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // BUFFERIZATION_TRANSFORM_OPS

mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td

let assemblyFormat = [{
     `block_dims` `=` $block_dims
     (`warp_dims` `=` $warp_dims^)?
     (`sync_after_distribute` `=` $sync_after_distribute^)?
     attr-dict
     `:` functional-type($target, $result)
   }];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def MapForallToBlocks :
   Op<Transform_Dialect, "gpu.map_forall_to_blocks",
let assemblyFormat = [{
     $target
     (`generate_gpu_launch` $generate_gpu_launch^)?
     (`grid_dims` `=` $grid_dims^)?
     attr-dict
     `:` functional-type($target, $result)
   }];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // GPU_TRANSFORM_OPS

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

 
 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//
 
 def BufferizeToAllocationOp : Op<Transform_Dialect,
     "structured.bufferize_to_allocation",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     This transform materializes an allocation for the targeted tensor value. It
     replaces all original uses of the target with the newly allocated buffer,
     wrapped in a `bufferization.to_tensor` op. It returns a handle to the result
     of the `to_tensor` op.
 
     Example:
     ```

 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
     [FunctionalStyleTransformOpTrait, 
      MemoryEffectsOpInterface,
      TransformOpInterface, 
-     TransformEachOpTrait]> {
+     TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Decomposes named complex operations, such as higher-dimensional
     (depthwise) convolutions, into combinations of lower-dimensional equivalents
     when possible.
 
     #### Return modes
 
     This operation ignores non-Linalg ops and drops them in the return.
     If all the operations referred to by the `target` handle decompose
     properly, the transform succeeds. Otherwise the transform silently fails.
     The return handle points to only the subset of successfully produced
     computational operations, which can be empty.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$transformed);
   let assemblyFormat =
       "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // FuseOp
 //===----------------------------------------------------------------------===//
 
 def FuseOp : Op<Transform_Dialect, "structured.fuse",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+     DeclareOpInterfaceMethods<TransformOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Tiles the operations pointed to by the target handle and fuses their
     producers greedily using the options provided as attributes.
   }];
 
   let arguments =
     (ins TransformHandleTypeInterface:$target,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_interchange);
   let results = (outs TransformHandleTypeInterface:$transformed,
                       Variadic<TransformHandleTypeInterface>:$loops);
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FuseIntoContainingOp
 //===----------------------------------------------------------------------===//
 
 def FuseIntoContainingOp :
     Op<Transform_Dialect, "structured.fuse_into_containing_op",
       [DeclareOpInterfaceMethods<TransformOpInterface,
           ["allowsRepeatedHandleOperands"]>,
-       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       ReportTrackingListenerFailuresOpTrait]> {
   let summary = "Fuse a producer into a containing operation.";
 
   let description = [{
     Fuses the `producer_op` into the `containing_op`.
     Returns a handle to the fused ops and the `new_containing_op`.
 
     The producer is typically a slice of a tileable op (i.e., implements
     TilingInterface). In that case, this transform computes the accessed

 }
 
 //===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
 def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformOpInterface, TransformEachOpTrait]> {
+     TransformOpInterface, TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Transforms a named structured operation into the generic form with the
     explicit attached region.
 
     #### Return modes
 
     This operation ignores non-Linalg ops and drops them in the return.
     If all the operations referred to by the `target` handle generalize
     properly, the transform succeeds. Otherwise the transform silently fails.
     The return handle points to only the subset of successfully produced
     equivalent generic operations, which can be empty or contain the original
     ops if they were already in generic form.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$transformed);
   let assemblyFormat =
       "$target attr-dict `:` "
       "custom<SemiFunctionType>(type($target), type($transformed))";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // InterchangeOp
 //===----------------------------------------------------------------------===//
 
 def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-    TransformOpInterface, TransformEachOpTrait]> {
+    TransformOpInterface, TransformEachOpTrait,
+    ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Interchanges the iterators of the operations pointed to by the target handle
     using the iterator interchange attribute.
 
     #### Return modes
 
     This operation ignores non-linalg::Generic ops and drops them in the return.
     This operation fails if the interchange attribute is invalid.
let assemblyFormat = [{
     $target 
     (`iterator_interchange` `=` $iterator_interchange^)? attr-dict
     `:` custom<SemiFunctionType>(type($target), type($transformed))
   }];
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::GenericOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // LowerPackOp
 //===----------------------------------------------------------------------===//
 def LowerPackOp : Op<Transform_Dialect, "structured.lower_pack", [
                          FunctionalStyleTransformOpTrait,
                          MemoryEffectsOpInterface,
                          TransformEachOpTrait,
-                         TransformOpInterface]> {
+                         TransformOpInterface,
+                         ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Rewrite a tensor.pack into tensor.pad + tensor.expand_shape + linalg.transpose.
 
     #### Return modes
 
     This operation ignores non-pack ops and drops them in the return.
     This operation produces a silenceableFailure if the rewrite fails for any
     reason.
     If all the operations referred to by the `target` are rewritten, the
     transform succeeds.
     Return handles to the newly produced pad, expand_shape and transpose ops.
   }];
 
   let arguments = (ins Transform_ConcreteOpType<"tensor.pack">:$target);
   let results = (outs Transform_ConcreteOpType<"tensor.pad">:$pad_op,
                       Transform_ConcreteOpType<"tensor.expand_shape">:$expand_shape_op,
                       Transform_ConcreteOpType<"linalg.transpose">:$transpose_op);
   let assemblyFormat = [{
     $target attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::tensor::PackOp target,
         ::mlir::transform::ApplyToEachResultList &transformResults,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // LowerUnPackOp
 //===----------------------------------------------------------------------===//
 def LowerUnPackOp : Op<Transform_Dialect, "structured.lower_unpack", [
                          FunctionalStyleTransformOpTrait,
                          MemoryEffectsOpInterface,
                          TransformEachOpTrait,
-                         TransformOpInterface]> {
+                         TransformOpInterface,
+                         ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Lower a tensor.unpack into empty + linalg.transpose + tensor.collapse_shape + 
     tensor.extract_slice.
 
     #### Return modes
 
     This operation ignores non-unpack ops and drops them in the return.
     This operation produces a silenceableFailure if the rewrite fails for any
let results = (outs Transform_ConcreteOpType<"tensor.empty">:$empty_op,
                       Transform_ConcreteOpType<"tensor.collapse_shape">:$collapse_shape_op,
                       Transform_ConcreteOpType<"tensor.extract_slice">:$extract_slice_op);
   let assemblyFormat = [{ 
     $target attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::tensor::UnPackOp target,
         ::mlir::transform::ApplyToEachResultList &transformResults,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MatchOp

 }
 
 //===----------------------------------------------------------------------===//
 // MultiTileSizesOp
 //===----------------------------------------------------------------------===//
 
 def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     TransformOpInterface, TransformEachOpTrait]> {
+     TransformOpInterface, TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Emits the IR computing the tile sizes `s1` and `s2` such that:
 
       - there exists a combination of `n` tiles of size `s1` and `m` tiles of
         size `s2` that covers the entirety of the iteration space `dimension` of
         the target structured op;
       - `s1`, `s2` is less than or equal to `target_size`;
       - `s1` and `s2` are divisible by `divisor.
let results = (outs TransformParamTypeOrAnyHandle:$low_size,
                       TransformParamTypeOrAnyHandle:$split_point);
   let hasVerifier = 1;
   let assemblyFormat =
     "$target attr-dict `:` custom<MultitileSizesTypes>("
     "type($target), type($low_size), type($high_size), type($split_point))";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PackOp
 //===----------------------------------------------------------------------===//
 
 def PackOp : Op<Transform_Dialect, "structured.pack", [
                 DeclareOpInterfaceMethods<TransformOpInterface>,
-                DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+                DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Pack a LinalgOp by applying a data tiling transformation on the op and
     packing the operands according to the `packed_sizes` specification.
     
     Iterator dimensions are tiled in their canonical order in the op spec.
     Operands are packed according to the same canonical order of the op iterator
     dimensions.
 
def PackOp : Op<Transform_Dialect, "structured.pack", [
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PackGreedilyOp
 //===----------------------------------------------------------------------===//
 def PackGreedilyOp : Op<Transform_Dialect, "structured.pack_greedily", [
                         DeclareOpInterfaceMethods<TransformOpInterface>,
-                        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+                        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Target a Linalg op and rewrite it into packed LinalgOp form by trying to
     infer whether a known suboperation is embedded
 
     Different packing strategies are applied in order, when one applies 
     successfully, the transform returns:
       1. Matmul packing: Try to infer a matmul operation embedded in the target op.
          Specifically, this looks for 2 parallel dimensions that participate in

 }
 
 //===----------------------------------------------------------------------===//
 // PackTransposeOp
 //===----------------------------------------------------------------------===//
 def PackTransposeOp : Op<Transform_Dialect, "structured.pack_transpose", [
                          FunctionalStyleTransformOpTrait,
                          MemoryEffectsOpInterface,
-                         DeclareOpInterfaceMethods<TransformOpInterface>]> {
+                         DeclareOpInterfaceMethods<TransformOpInterface>,
+                         ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Apply a transposition to a single `tensor.pack` (resp. `tensor.unpack`) and 
     update the `linalg.generic` op that consumes (resp. produces) the operation.
 
     This transform allows composing a simple `structured.pack` with additional
     transpositions to e.g. match the data format required by a specific library
     call or ISA instruction.
 

 }
 
 //===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//
 
 def PadOp : Op<Transform_Dialect, "structured.pad",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformOpInterface, TransformEachOpTrait]> {
+     TransformOpInterface, TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Pads the operations pointed to by the target handle using the options
     provides as operation attributes.
 
     #### Return modes
 
     This operation ignores non-Linalg ops and drops them in the return.
     This operation may produce a definiteFailure if the padding fails for any
def PadOp : Op<Transform_Dialect, "structured.pad",
 
   let assemblyFormat =
     "$target attr-dict `:` "
     "custom<SemiFunctionType>(type($target), type($transformed))";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // HoistPadOp
 //===----------------------------------------------------------------------===//
 
 def HoistPadBuildPackingLoopNestOp :
     Op<Transform_Dialect,
        "structured.hoist_pad.build_packing_loop_nest",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+     DeclareOpInterfaceMethods<TransformOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Helper transform used to hoist a tensor.pad target operation. This operation
     creates the packing loop nest required by the hoist_pad operation and makes
     that functionality available independently.
 
     TODO: In the future, we should consider rewriting as a tensor.pack after
     hoisting since this abstraction is now available.
 
let assemblyFormat = [{
     (`,` `transpose` `by` $transpose^)? 
     attr-dict
     `:` functional-type(operands, results)
   }];
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::tensor::PadOp,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 
 def PromoteOp : Op<Transform_Dialect, "structured.promote",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-    TransformOpInterface, TransformEachOpTrait]> {
+    TransformOpInterface, TransformEachOpTrait,
+    ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Promotes the specified operands of the target into a separate memory buffer.
 
     At this point, this transform does not allow customizing alloc/dealloc
     functions nor the behavior on copy in/out operations.
 
     #### Return modes
 
def PromoteOp : Op<Transform_Dialect, "structured.promote",
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:`"
     "custom<SemiFunctionType>(type($target), type($transformed))";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
 def ReplaceOp : Op<Transform_Dialect, "structured.replace",
     [IsolatedFromAbove, DeclareOpInterfaceMethods<TransformOpInterface>,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>] # GraphRegionNoTerminator.traits> {
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     ReportTrackingListenerFailuresOpTrait] # GraphRegionNoTerminator.traits> {
   let description = [{
     Replace all `target` payload ops with the single op that is contained in
     this op's region. All targets must have zero arguments and must be isolated
     from above.
 
     This op is for debugging/experiments only.
 
     #### Return modes

 }
 
 //===----------------------------------------------------------------------===//
 // ScalarizeOp
 //===----------------------------------------------------------------------===//
 
 def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformOpInterface, TransformEachOpTrait]> {
+     TransformOpInterface, TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that ops of a specific kind in the given function should be
     scalarized (i.e. their dynamic dimensions tiled by 1).
 
     #### Return modes:
 
     This operation ignores non-Linalg ops and drops them in the return.
     This operation produces `definiteFailure` if the scalarization fails for any
def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
   let results = (outs TransformHandleTypeInterface:$result);
 
   let assemblyFormat = 
     "$target attr-dict `:`"
     "custom<SemiFunctionType>(type($target), type($result))";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // RewriteInDestinationPassingStyleOp.
 //===----------------------------------------------------------------------===//
 
 def RewriteInDestinationPassingStyleOp : Op<
     Transform_Dialect, "structured.rewrite_in_destination_passing_style",
     [FunctionalStyleTransformOpTrait, 
      MemoryEffectsOpInterface,
      TransformOpInterface, 
-     TransformEachOpTrait]> {
+     TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Rewrite a supported tensor operation that is not in destination-passing style
     into a form that is in destination-passing style.
     Currently supported operations are:
       - tensor.pad
       - tensor.generate
       - tensor.from_elements
     This dichotomy hints at a future interface, for now the implementation just 
def RewriteInDestinationPassingStyleOp : Op<
   let results = (outs TransformHandleTypeInterface:$transformed);
   let assemblyFormat = [{
     $target attr-dict
     `:` functional-type($target, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
 def SplitOp : Op<Transform_Dialect, "structured.split",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+     DeclareOpInterfaceMethods<TransformOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op should be split into two complementary
     parts, which combined cover the entire iteration domain of the original op.
     The split is performed along the iteration space dimension provided as
     attribute. In case of dimension overflow, the transformation fails. The
     split is performed at the dimension iterator value specified as either the
     static split point attribute when it is known at transform IR construction
     time or as the handle to an operation producing a single index-typed value

 }
 
 //===----------------------------------------------------------------------===//
 // SplitReductionOp
 //===----------------------------------------------------------------------===//
 
 def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-        TransformEachOpTrait, TransformOpInterface]> {
+        TransformEachOpTrait, TransformOpInterface,
+        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op should be transformed with the
     `splitReduction` transformation and split factor provided as attribute.
 
     The `splitReduction` transformation splits the first single linalg op
     reduction into a parallel and reduction dimension.
     A new `linalg.generic` op is created to perform the rest of the reduction.
 
OpBuilder<(ins "Value":$target,
                    "int64_t":$insertSplitDimension,
                    CArg<"bool", "false">:$innerParallel,
                    CArg<"bool", "false">:$useScalingAlgorithm,
                    CArg<"bool", "false">:$useAlloc)>
   ];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TileReductionUsingScfOp
 //===----------------------------------------------------------------------===//
 
 def TileReductionUsingScfOp : Op<Transform_Dialect, "structured.tile_reduction_using_scf",
        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-        TransformEachOpTrait, TransformOpInterface]> {
+        TransformEachOpTrait, TransformOpInterface,
+        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op should be transformed with the
     `tileReduction` transformation with the tile size provided as attribute.
 
     This transformation tiles the `target` along the reduction dimensions. It
     creates a tensor initialized with the identity value. Then it creates nested
     loops with a parallel version of `target` op inside. The parallel op
     dimensions are less or equal to the tile size passed by user.
let assemblyFormat = [{
     $target
     `by` `tile_sizes` `=` $tile_sizes
     attr-dict
     `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TileReductionUsingForallOp
 //===----------------------------------------------------------------------===//
 
 def TileReductionUsingForallOp :
   Op<Transform_Dialect, "structured.tile_reduction_using_forall",
        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-        TransformEachOpTrait, TransformOpInterface]> {
+        TransformEachOpTrait, TransformOpInterface,
+        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Tile a PartialReductionOpInterface op to a tiled `scf.forall` doing
     partial reduction.
 
     This transformation tiles the `target` along the reduction dimensions. It
     creates a tensor initialized with the identity value. Then it creates a
     `scf.forall` loops with the number threads given by `num_threads`.
     The op is tiled op with a size equal to `floordiv(size, num_threads)`.
let assemblyFormat = [{
     (`,` `tile_sizes` `=` $tile_sizes^)?
     (`,` `mapping` `=` $mapping^)?
     attr-dict
     `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 
 }
 
 //===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
 
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
-        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the given sizes.
     This transform generates a loop nest with a smaller ("tiled") target
     operation in its body. Currently limited to LinalgOps.
 
     Tile sizes may be known at transformation time, in which case they are
     expected to be provided in the `static_size` attribute, or not, in which
     case the tile value must be computed by the payload IR and the handle to the

 //===----------------------------------------------------------------------===//
 // TileToForallOp
 //===----------------------------------------------------------------------===//
 
 def TileToForallOp :
     Op<Transform_Dialect, "structured.tile_to_forall_op",
       [AttrSizedOperandSegments,
        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-       TransformOpInterface]> {
+       TransformOpInterface, ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Tile a TilingInterface op to a tiled `scf.forall`.
 
     Tiling is applied by either specifying `num_threads` or `tile_size`. If
     `num_threads` is specified, then the tile size for each dimension `i` is
     calculated dynamically via `ceilDiv(dimSize[i], num_threads[i])`.
     `num_threads` and `tile_size` can be either static index attributes or SSA
     values of PDL operation handle type (or a mix thereof). Operation handles
$target oilist(
                                                        $static_tile_sizes))
     (`(` `mapping` `=` $mapping^ `)`)? attr-dict
     `:` functional-type($target, results)
   }];
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure apply(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::transform::TransformResults &transformResults,
         ::mlir::transform::TransformState &state);
 
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedNumThreads();
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedTileSizes();
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TileToScfForOp
 //===----------------------------------------------------------------------===//
 
 def TileToScfForOp : Op<Transform_Dialect, "structured.tile_to_scf_for",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
-        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the given sizes.
     This transform generates a loop nest with a smaller ("tiled") target
     operation in its body. The target must implement TilingInterface.
 
     Tile sizes may be known at transformation time, in which case they are
     expected to be provided in the `static_size` attribute, or not, in which
     case the tile value must be computed by the payload IR and the handle to the

 }
 
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformEachOpTrait, TransformOpInterface]> {
+     TransformEachOpTrait, TransformOpInterface,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Indicates that the given `target` op all the ops it contains should be
     vectorized with the configuration specified by the attributes of this op.
     This vectorization only handles structured ops that operate on shaped types
     and does not vectorize loops or straight-line. Internally, it applies a
     set of rewrite patterns, some of which enable vectorization and some of
     which clean up the results. Therefore, it can only be applied to an op with
     the "isolated from above property". If finer granularity is required, it can
def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
 
   let builders = [
     OpBuilder<(ins "Value":$target,
                CArg<"bool", "false">:$vectorizePadding,
                CArg<"bool", "false">:$vectorizeNDExtract)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     TransformOpInterface]> {
+     TransformOpInterface, ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Vectorize the target ops, which must be Linalg ops, with masked vectors
     of the specified size.
 
     The vector sizes can be either static or dynamic (SSA values). In case of
     SSA values, the handle must be mapped to exactly one payload op with
     exactly one index-typed result.
 
let assemblyFormat = [{
                                               $last_vector_size_scalable)
       attr-dict
       `:` type($target)
   }];
 
   let extraClassDeclaration = [{
     // TODO: applyToOne.
     ::mlir::DiagnosedSilenceableFailure apply(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::transform::TransformResults &transformResults,
         ::mlir::transform::TransformState &state);
 
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedVectorSizes();
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // HoistRedundantVectorTransfersOp
 //===----------------------------------------------------------------------===//
 
 def HoistRedundantVectorTransfersOp :
   Op<Transform_Dialect, "structured.hoist_redundant_vector_transfers",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformEachOpTrait, TransformOpInterface]> {
+     TransformEachOpTrait, TransformOpInterface,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Hoist vector.transfer_read / vector.transfer_write pairs out of immediately
     enclosing scf::ForOp iteratively, if the following conditions are true:
        1. The 2 ops access the same memref with the same indices.
        2. All operands are invariant under the enclosing scf::ForOp.
        3. No uses of the memref either dominate the transfer_read or are
        dominated by the transfer_write (i.e. no aliasing between the write and
        the read across the loop)
def HoistRedundantVectorTransfersOp :
 
   let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
 
   let builders = [
     OpBuilder<(ins "Value":$target)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+         ::mlir::transform::TransformRewriter &rewriter,
          ::mlir::func::FuncOp target,
          ::mlir::transform::ApplyToEachResultList &results,
          ::mlir::transform::TransformState &state);
    }];
 }
 
 //===----------------------------------------------------------------------===//
 // ConvertConv2DToImg2ColOp
 //===----------------------------------------------------------------------===//
 
 def ConvertConv2DToImg2ColOp : Op<Transform_Dialect,
     "structured.convert_conv2d_to_img2col",
     [FunctionalStyleTransformOpTrait,
      MemoryEffectsOpInterface,
      TransformOpInterface,
-     TransformEachOpTrait]> {
+     TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Convert linalg.conv_2d_xxx into linalg.generic (for img2col packing)
     and linalg.matmul.
 
     A convolution operation can be written as a matrix-matrix multiplication by
     unfolding the cross-correlation between input and filter and explicitly copy
     overlapped sliding window inputs.
 
let assemblyFormat =
     "$target attr-dict `:` functional-type($target, results)";
 
   let builders = [
     OpBuilder<(ins "Value":$target)>
   ];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // HoistRedundantTensorSubsetsOp
 //===----------------------------------------------------------------------===//
 
 def HoistRedundantTensorSubsetsOp :
   Op<Transform_Dialect, "structured.hoist_redundant_tensor_subsets",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformEachOpTrait,
-     TransformOpInterface]> {
+     TransformOpInterface,
+     ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Hoists supported tensor subset extract/insert operation pairs out of 
     immediately enclosing loop iteratively, if the following conditions
     are true:
        1. The 2 ops access the same tensor subset.
        2. All operands are invariant under the enclosing loop.
     
     The supported subset extract/insert operation pairs currently comprise:
def HoistRedundantTensorSubsetsOp :
   let assemblyFormat = [{
     $target 
     attr-dict 
     `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // InsertSliceToCopyOp
def InsertSliceToCopyOp :
 
   let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
 
   let builders = [
     OpBuilder<(ins "Value":$target)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

def MemRefMakeLoopIndependentOp
 
   let arguments = (ins TransformHandleTypeInterface:$target, I64Attr:$num_loops);
   let results = (outs TransformHandleTypeInterface:$transformed);
   let assemblyFormat =
       "$target attr-dict `:` functional-type($target, $transformed)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // MEMREF_TRANSFORM_OPS

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
   // TODO: Return both the peeled loop and the remainder loop.
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::scf::ForOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
let arguments = (ins Transform_ScfForOp:$target,
                    DefaultValuedAttr<I64Attr, "10">:$read_latency);
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::scf::ForOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        ConfinedAttr<I64Attr, [IntPositive]>:$factor);
 
   let assemblyFormat = "$target attr-dict `:` type($target)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def LoopCoalesceOp : Op<Transform_Dialect, "loop.coalesce", [
   FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
def LoopCoalesceOp : Op<Transform_Dialect, "loop.coalesce", [
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
       "$target attr-dict `:` functional-type($target, $transformed)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TakeAssumedBranchOp : Op<Transform_Dialect, "scf.take_assumed_branch", [
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
let assemblyFormat = [{
       $target
       (`take_else_branch` $take_else_branch^)?
       attr-dict
        `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::scf::IfOp ifOp,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // SCF_TRANSFORM_OPS

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

def MakeLoopIndependentOp
 
   let arguments = (ins TransformHandleTypeInterface:$target, I64Attr:$num_loops);
   let results = (outs TransformHandleTypeInterface:$transformed);
   let assemblyFormat =
       "$target attr-dict `:` functional-type($target, $transformed)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // TENSOR_TRANSFORM_OPS

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

static LogicalResult verifyTrait(Operation *op) {
     if (!isa<TransformHandleTypeInterface>(operandHandle.getType())) {
       return op->emitError() << "SingleOpMatchOpTrait requires the op handle "
                                 "to be of TransformHandleTypeInterface";
     }
 
     return success();
   }
 
-  DiagnosedSilenceableFailure apply(TransformResults &results,
+  DiagnosedSilenceableFailure apply(TransformRewriter &rewriter,
+                                    TransformResults &results,
                                     TransformState &state) {
     Value operandHandle = cast<OpTy>(this->getOperation()).getOperandHandle();
     auto payload = state.getPayloadOps(operandHandle);
     if (!llvm::hasSingleElement(payload)) {
       return emitDefiniteFailure(this->getOperation()->getLoc())
              << "SingleOpMatchOpTrait requires the operand handle to point to "
                 "a single payload op";
     }
static LogicalResult verifyTrait(Operation *op) {
     if (!isa<TransformValueHandleTypeInterface>(operandHandle.getType())) {
       return op->emitError() << "SingleValueMatchOpTrait requires an operand "
                                 "of TransformValueHandleTypeInterface";
     }
 
     return success();
   }
 
-  DiagnosedSilenceableFailure apply(TransformResults &results,
+  DiagnosedSilenceableFailure apply(TransformRewriter &rewriter,
+                                    TransformResults &results,
                                     TransformState &state) {
     Value operandHandle = cast<OpTy>(this->getOperation()).getOperandHandle();
     ValueRange payload = state.getPayloadValues(operandHandle);
     if (payload.size() != 1) {
       return emitDefiniteFailure(this->getOperation()->getLoc())
              << "SingleValueMatchOpTrait requires the value handle to point to "
                 "a single payload value";
     }

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

def Transform_Dialect : Dialect {
   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
           kWithNamedSequenceAttrName = "transform.with_named_sequence";
 
-      /// Names of the attribute attachable to an operation so it can be
+      /// Name of the attribute attachable to an operation so it can be
       /// identified as root by the default interpreter pass.
       constexpr const static ::llvm::StringLiteral
           kTargetTagAttrName = "transform.target_tag";
 
+      /// Name of the attribute attachable to an operation, indicating that
+      /// TrackingListener failures should be silenced.
+      constexpr const static ::llvm::StringLiteral
+          kSilenceTrackingFailuresAttrName = "transform.silence_tracking_failures";
+
       /// Names of the attributes indicating whether an argument of an external
       /// transform dialect symbol is consumed or only read.
       constexpr const static ::llvm::StringLiteral
           kArgConsumedAttrName = "transform.consumed";
       constexpr const static ::llvm::StringLiteral
           kArgReadOnlyAttrName = "transform.readonly";
 
       template <typename DataTy>

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

 #include "mlir/Dialect/Transform/Utils/DiagnosedSilenceableFailure.h"
 #include "mlir/Dialect/Transform/Utils/RaggedArray.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
-
 namespace transform {
 
 class TransformOpInterface;
 class TransformResults;
+class TransformRewriter;
 class TransformState;
 
 using Param = Attribute;
 using MappedValue = llvm::PointerUnion<Operation *, Param, Value>;
 
 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

 
 /// A listener that updates a TransformState based on IR modifications. This
 /// listener can be used during a greedy pattern rewrite to keep the transform
 /// state up-to-date.
 class TrackingListener : public RewriterBase::Listener,
                          public TransformState::Extension {
 public:
   /// Create a new TrackingListener for usage in the specified transform op.
-  explicit TrackingListener(TransformState &state, TransformOpInterface op)
-      : TransformState::Extension(state), transformOp(op) {}
+  TrackingListener(TransformState &state, TransformOpInterface op);
 
 protected:
   /// Return a replacement payload op for the given op, which is going to be
   /// replaced with the given values. By default, if all values are defined by
   /// the same op, which also has the same type as the given op, that defining
   /// op is used as a replacement.
   ///
   /// A "failure" return value indicates that no replacement operation could be

 
 private:
   void notifyOperationRemoved(Operation *op) override;
 
   void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
 
   /// The transform op in which this TrackingListener is used.
   TransformOpInterface transformOp;
+
+  /// The handles that are consumed by the transform op.
+  DenseSet<Value> consumedHandles;
 };
 
 /// 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;
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;
 };
 
+/// This is a special rewriter to be used in transform op implementations,
+/// providing additional helper functions to update the transform state, etc.
+// TODO: Helper functions will be added in a subsequent change.
+class TransformRewriter : public RewriterBase {
+protected:
+  friend class TransformState;
+
+  /// Create a new TransformRewriter.
+  explicit TransformRewriter(MLIRContext *ctx,
+                             ErrorCheckingTrackingListener *listener);
+
+public:
+  /// Return "true" if the tracking listener had failures.
+  bool hasTrackingFailures() const;
+
+  /// Silence all tracking failures that have been encountered so far.
+  void silenceTrackingFailure();
+
+private:
+  ErrorCheckingTrackingListener *const listener;
+};
+
 /// 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 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
public:
   ///   3. All applications of `applyToOne` are checked to return a number of
   ///      results expected by the transform IR op. If not, this is a definite
   ///      failure and we return early.
   ///   4. If `applyToOne` produces ops, associate them with the result of this
   ///      transform op.
   ///   5. If any `applyToOne` return silenceableFailure, the transformation is
   ///      considered silenceable.
   ///   6. Otherwise the transformation is considered successful.
-  DiagnosedSilenceableFailure apply(TransformResults &transformResults,
+  DiagnosedSilenceableFailure apply(transform::TransformRewriter &rewriter,
+                                    TransformResults &transformResults,
                                     TransformState &state);
 
   /// Checks that the op matches the expectations of this trait.
   static LogicalResult verifyTrait(Operation *op);
 };
 
 /// Side effect resource corresponding to the mapping between Transform IR
 /// values and Payload IR operations. An Allocate effect from this resource
public:
   /// Checks that the op matches the expectation of this trait, i.e., that it
   /// implements the MemoryEffectsOpInterface and only produces parameter-typed
   /// results.
   static LogicalResult verifyTrait(Operation *op) {
     return detail::verifyParamProducerTransformOpTrait(op);
   }
 };
 
+/// `TrackingListener` failures are reported only for ops that have this trait.
+/// The purpose of this trait is to give users more time to update their custom
+/// transform ops to use the provided `TransformRewriter` for all IR
+/// modifications. This trait will eventually be removed, and failures will be
+/// reported for all transform ops.
+template <typename OpTy>
+class ReportTrackingListenerFailuresOpTrait
+    : public OpTrait::TraitBase<OpTy, ReportTrackingListenerFailuresOpTrait> {};
+
 /// A single result of applying a transform op with `ApplyEachOpTrait` to a
 /// single payload operation.
 using ApplyToEachResult = MappedValue;
 
 /// A list of results of applying a transform op with `ApplyEachOpTrait` to a
 /// single payload operation, co-indexed with the results of the transform op.
 class ApplyToEachResultList {
 public:

 ///   - DiagnosedSilenceableFailure(OpTy,
 ///       SmallVector<Operation*> &results, state)
 /// where OpTy is either
 ///   - Operation *, in which case the transform is always applied;
 ///   - a concrete Op class, in which case a check is performed whether
 ///   `targets` contains operations of the same class and a silenceable failure
 ///   is reported if it does not.
 template <typename TransformOpTy, typename Range>
-DiagnosedSilenceableFailure
-applyTransformToEach(TransformOpTy transformOp, Range &&targets,
-                     SmallVectorImpl<ApplyToEachResultList> &results,
-                     TransformState &state) {
+DiagnosedSilenceableFailure applyTransformToEach(
+    TransformOpTy transformOp, TransformRewriter &rewriter, Range &&targets,
+    SmallVectorImpl<ApplyToEachResultList> &results, TransformState &state) {
   using OpTy = typename llvm::function_traits<
-      decltype(&TransformOpTy::applyToOne)>::template arg_t<0>;
+      decltype(&TransformOpTy::applyToOne)>::template arg_t<1>;
   static_assert(std::is_convertible<OpTy, Operation *>::value,
                 "expected transform function to take an operation");
+  OpBuilder::InsertionGuard g(rewriter);
 
   SmallVector<Diagnostic> silenceableStack;
   unsigned expectedNumResults = transformOp->getNumResults();
   for (Operation *target : targets) {
     auto specificOp = dyn_cast<OpTy>(target);
     if (!specificOp) {
       Diagnostic diag(transformOp->getLoc(), DiagnosticSeverity::Error);
       diag << "transform applied to the wrong op kind";
       diag.attachNote(target->getLoc()) << "when applied to this op";
       silenceableStack.push_back(std::move(diag));
       continue;
     }
 
     ApplyToEachResultList partialResults;
     partialResults.reserve(expectedNumResults);
     Location specificOpLoc = specificOp->getLoc();
+    rewriter.setInsertionPoint(specificOp);
     DiagnosedSilenceableFailure res =
-        transformOp.applyToOne(specificOp, partialResults, state);
+        transformOp.applyToOne(rewriter, specificOp, partialResults, state);
     if (res.isDefiniteFailure())
       return DiagnosedSilenceableFailure::definiteFailure();
 
     if (res.isSilenceableFailure()) {
       res.takeDiagnostics(silenceableStack);
       continue;
     }
 

 
 } // namespace detail
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
-    TransformResults &transformResults, TransformState &state) {
+    TransformRewriter &rewriter, TransformResults &transformResults,
+    TransformState &state) {
   Value handle = this->getOperation()->getOperand(0);
   auto targets = state.getPayloadOps(handle);
 
   // If the operand is consumed, check if it is associated with operations that
   // may be erased before their nested operations are.
   if (state.getOptions().getExpensiveChecksEnabled() &&
       isHandleConsumed(handle, cast<transform::TransformOpInterface>(
                                    this->getOperation())) &&
if (std::empty(targets)) {
     }
     return DiagnosedSilenceableFailure::success();
   }
 
   // Step 2. Call applyToOne on each target and record newly produced ops in its
   // corresponding results entry.
   SmallVector<ApplyToEachResultList, 1> results;
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
-      cast<OpTy>(this->getOperation()), targets, results, state);
+      cast<OpTy>(this->getOperation()), rewriter, targets, results, state);
 
   // Step 3. Propagate the definite failure if any and bail out.
   if (result.isDefiniteFailure())
     return result;
 
   // Step 4. "Transpose" the results produced by individual applications,
   // arranging them per result value of the transform op. The number, kind and
   // nullity of per-application results have been verified by the callback

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

InterfaceMethod<
       /*desc=*/[{
         Applies the transformation represented by the current operation. This
         accepts as arguments the object that must be populated with results of
         the current transformation and a transformation state object that can be
         used for queries, e.g., to obtain the list of operations on which the
         transformation represented by the current op is targeted. Returns a
         special status object indicating whether the transformation succeeded
         or failed, and, if it failed, whether the failure is recoverable or not.
+
+        IR must be created, modified and deleted with the provided rewriter.
+        implementations are responsible for setting the insertion point of the
+        rewriter to the desired location.
       }],
       /*returnType=*/"::mlir::DiagnosedSilenceableFailure",
       /*name=*/"apply",
       /*arguments=*/(ins
+          "::mlir::transform::TransformRewriter &":$rewriter,
           "::mlir::transform::TransformResults &":$transformResults,
           "::mlir::transform::TransformState &":$state
     )>,
     InterfaceMethod<
       /*desc=*/[{
         Indicates whether the op instance allows its handle operands to be
         associated with the same payload operations.
       }],

 def NavigationTransformOpTrait : NativeOpTrait<"NavigationTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def ParamProducerTransformOpTrait : NativeOpTrait<"ParamProducerTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
+def ReportTrackingListenerFailuresOpTrait : NativeOpTrait<"ReportTrackingListenerFailuresOpTrait"> {
+  let cppNamespace = "::mlir::transform";
+}
+
 def FindPayloadReplacementOpInterface
     : OpInterface<"FindPayloadReplacementOpInterface"> {
   let description = [{
     This interface is queried by the `TrackingListener` and can be implemented
     by payload ops to indicate that the lookup should be continue with its
     operands when looking for payload op replacements.
 
     Example: Consider the case where a tracked "test.foo" payload op is replaced

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

def AnnotateOp : TransformDialectOp<"annotate",
 
   let assemblyFormat =
     "$target $name attr-dict (`=` $param^)?"
     "`:` type($target) (`,` type($param)^)?";
 }
 
 def ApplyPatternsOp : TransformDialectOp<"apply_patterns",
     [TransformOpInterface, TransformEachOpTrait,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]
         # GraphRegionNoTerminator.traits> {
   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.
 
     The patterns that should be applied are specified in the graph region of
     this op. They must implement the `PatternDescriptorOpInterface`. The order
     in which patterns are applied is unspecified; i.e., the ordering of ops in
     the region of this op is irrelevant.
 
     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`.
+    silently. 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,
-    DefaultValuedAttr<BoolAttr, "true">:$fail_on_payload_replacement_not_found);
+    TransformHandleTypeInterface:$target);
   let results = (outs);
   let regions = (region MaxSizedRegion<1>:$region);
 
   let assemblyFormat = "`to` $target $region attr-dict `:` type($target)";
   let hasVerifier = 1;
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins
         "Value":$target,
         CArg<"function_ref<void(OpBuilder &, Location)>", "nullptr">:
-            $bodyBuilder,
-        CArg<"bool", "true">:$failOnPayloadReplacementNotFound)>,
+            $bodyBuilder)>,
   ];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::transform::TransformRewriter &rewriter,
       ::mlir::Operation *target,
       ::mlir::transform::ApplyToEachResultList &results,
       ::mlir::transform::TransformState &state);
   }];
 }
 
 def ApplyCanonicalizationPatternsOp
     : TransformDialectOp<"apply_patterns.canonicalization",
let arguments = (ins TransformHandleTypeInterface:$target,
                        DefaultValuedAttr<StrAttr, "\"\"">:$options);
   let results = (outs TransformHandleTypeInterface:$result);
   let assemblyFormat = [{
     $pass_name `to` $target attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::transform::TransformRewriter &rewriter, 
       ::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)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+      ::mlir::transform::TransformRewriter &rewriter,
       ::mlir::Operation *target,
       ::mlir::transform::ApplyToEachResultList &results,
       ::mlir::transform::TransformState &state);
   }];
 }
 
 def ForeachMatchOp : TransformDialectOp<"foreach_match", [
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

LogicalResult matchAndRewrite(OpTy op,
     return success();
   }
 
   const FlatAffineValueConstraints &constraints;
 };
 } // namespace
 
 DiagnosedSilenceableFailure
-SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
+SimplifyBoundedAffineOpsOp::apply(transform::TransformRewriter &rewriter,
+                                  TransformResults &results,
                                   TransformState &state) {
   // Get constraints for bounded values.
   SmallVector<int64_t> lbs;
   SmallVector<int64_t> ubs;
   SmallVector<Value> boundedValues;
   DenseSet<Operation *> boundedOps;
   for (const auto &it : llvm::zip_equal(getBoundedValues(), getLowerBounds(),
                                         getUpperBounds())) {
SimplifyBoundedAffineOpsOp::apply(transform::TransformRewriter &rewriter,
   // Canonicalization patterns are needed so that affine.apply ops are composed
   // with the remaining affine.min/max ops.
   AffineMaxOp::getCanonicalizationPatterns(patterns, getContext());
   AffineMinOp::getCanonicalizationPatterns(patterns, getContext());
   patterns.insert<SimplifyAffineMinMaxOp<AffineMinOp>,
                   SimplifyAffineMinMaxOp<AffineMaxOp>>(getContext(), cstr);
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
   GreedyRewriteConfig config;
+  config.listener =
+      static_cast<RewriterBase::Listener *>(rewriter.getListener());
   config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
   // Apply the simplification pattern to a fixpoint.
   if (failed(applyOpPatternsAndFold(targets, frozenPatterns, config))) {
     auto diag = emitDefiniteFailure()
                 << "affine.min/max simplification did not converge";
     return diag;
   }
   return DiagnosedSilenceableFailure::success();

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

 using namespace mlir::bufferization;
 using namespace mlir::transform;
 
 //===----------------------------------------------------------------------===//
 // OneShotBufferizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
+transform::OneShotBufferizeOp::apply(transform::TransformRewriter &rewriter,
+                                     TransformResults &transformResults,
                                      TransformState &state) {
   OneShotBufferizationOptions options;
   options.allowReturnAllocs = getAllowReturnAllocs();
   options.allowUnknownOps = getAllowUnknownOps();
   options.bufferizeFunctionBoundaries = getBufferizeFunctionBoundaries();
   options.createDeallocs = getCreateDeallocs();
   options.testAnalysisOnly = getTestAnalysisOnly();
   options.printConflicts = getPrintConflicts();

 //===----------------------------------------------------------------------===//
 
 void transform::EliminateEmptyTensorsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getTarget(), effects);
   modifiesPayload(effects);
 }
 
-DiagnosedSilenceableFailure
-transform::EliminateEmptyTensorsOp::apply(TransformResults &transformResults,
-                                          TransformState &state) {
-  IRRewriter rewriter(getContext());
+DiagnosedSilenceableFailure transform::EliminateEmptyTensorsOp::apply(
+    transform::TransformRewriter &rewriter, TransformResults &transformResults,
   OneShotBufferizationOptions options;
   options.allowReturnAllocs = true;
 
   for (Operation *target : state.getPayloadOps(getTarget())) {
     OneShotAnalysisState state(target, options);
     if (failed(analyzeOp(target, state)))
       return mlir::emitSilenceableFailure(target->getLoc())
              << "failed to analyze op";
     if (failed(bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
             rewriter, target, state)))
       return mlir::emitSilenceableFailure(target->getLoc())
              << "failed to eliminate insert_slice anchored tensor.empty ops";
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // EmptyTensorToAllocTensorOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure
-EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
-                                       ApplyToEachResultList &results,
-                                       transform::TransformState &state) {
-  IRRewriter rewriter(target->getContext());
+DiagnosedSilenceableFailure EmptyTensorToAllocTensorOp::applyToOne(
+    transform::TransformRewriter &rewriter, tensor::EmptyOp target,
+    ApplyToEachResultList &results, transform::TransformState &state) {
   rewriter.setInsertionPoint(target);
   auto alloc = rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
       target, target.getType(), target.getDynamicSizes());
   results.push_back(alloc);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

launchOp = rewriter.create<LaunchOp>(loc, gridSizeX, gridSizeY, gridSizeZ,
                                        blkSizeX, blkSizeY, blkSizeZ);
   rewriter.setInsertionPointToEnd(&launchOp.getBody().front());
   rewriter.create<TerminatorOp>(loc);
   return DiagnosedSilenceableFailure::success();
 }
 
 /// Alter kernel configuration of the given kernel.
 static DiagnosedSilenceableFailure
-alterGpuLaunch(IRRewriter &rewriter, LaunchOp gpuLaunch,
+alterGpuLaunch(RewriterBase &rewriter, LaunchOp gpuLaunch,
                TransformOpInterface transformOp,
                std::optional<int64_t> gridDimX = std::nullopt,
                std::optional<int64_t> gridDimY = std::nullopt,
                std::optional<int64_t> gridDimZ = std::nullopt,
                std::optional<int64_t> blockDimX = std::nullopt,
                std::optional<int64_t> blockDimY = std::nullopt,
                std::optional<int64_t> blockDimZ = std::nullopt) {
   DiagnosedSilenceableFailure diag =
mlir::transform::gpu::findTopLevelForallOp(Operation *target,
   });
 
   if (walkResult.wasInterrupted())
     return transformOp.emitSilenceableError()
            << "could not find a unique topLevel scf.forall";
   return DiagnosedSilenceableFailure::success();
 }
 
-DiagnosedSilenceableFailure
-transform::MapForallToBlocks::applyToOne(Operation *target,
-                                         ApplyToEachResultList &results,
-                                         transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::MapForallToBlocks::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    ApplyToEachResultList &results, transform::TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
-  IRRewriter rewriter(getContext());
   auto transformOp = cast<TransformOpInterface>(getOperation());
 
   if (!getGenerateGpuLaunch() && !gpuLaunch) {
     DiagnosedSilenceableFailure diag =
         emitSilenceableError()
         << "Given target is not gpu.launch, set `generate_gpu_launch` "
            "attribute";
     diag.attachNote(target->getLoc()) << "when applied to this payload op";
DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(
   // Here, the result of mapping determines the available mapping sizes.
   replaceUnitMappingIdsHelper<ThreadIdOp>(rewriter, loc, target, zero,
                                           blockDims);
 
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(
-    Operation *target, ApplyToEachResultList &results, TransformState &state) {
+    transform::TransformRewriter &rewriter, Operation *target,
+    ApplyToEachResultList &results, TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
   auto transformOp = cast<TransformOpInterface>(getOperation());
 
   // Basic high-level verifications.
   if (!gpuLaunch)
     return emitSilenceableError() << "Given target is not a gpu.launch";
 
   // Mapping to block ids.
   SmallVector<int64_t> blockDims{getBlockDims()};
 
   DiagnosedSilenceableFailure diag =
       checkGpuLimits(transformOp, std::nullopt, std::nullopt, std::nullopt,
                      blockDims[0], blockDims[1], blockDims[2]);
   if (diag.isSilenceableFailure()) {
     diag.attachNote(getLoc()) << getBlockDimsAttrName() << " is too large";
     return diag;
   }
 
   // Set the GPU launch configuration for the block dims early, this is not
   // subject to IR inspection.
-  IRRewriter rewriter(getContext());
   diag = alterGpuLaunch(rewriter, gpuLaunch, transformOp, std::nullopt,
                         std::nullopt, std::nullopt, blockDims[0], blockDims[1],
                         blockDims[2]);
 
   rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());
   diag =
       mapNestedForallToThreadsImpl(rewriter, transformOp, gpuLaunch, blockDims,
                                    getWarpDims(), getSyncAfterDistribute());

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

void transform::ApplyTilingCanonicalizationPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   linalg::populateLinalgTilingCanonicalizationPatterns(patterns);
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure
-transform::BufferizeToAllocationOp::apply(transform::TransformResults &results,
-                                          transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::BufferizeToAllocationOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   auto transformed = llvm::to_vector(
       llvm::map_range(state.getPayloadValues(getTarget()), [&](Value v) {
         return linalg::bufferizeToAllocation(rewriter, v, memorySpace);
       }));
   results.setValues(cast<OpResult>(getTransformed()), transformed);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::BufferizeToAllocationOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   producesHandle(getTransformed(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::DecomposeOp::applyToOne(LinalgOp target,
+transform::DecomposeOp::applyToOne(transform::TransformRewriter &rewriter,
+                                   LinalgOp target,
                                    transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
 #define DOWNSCALE(trans)                                                       \
   {                                                                            \
     FailureOr<LinalgOp> res = tryApply<trans>(target);                         \
     if (succeeded(res)) {                                                      \
       results.push_back(*res);                                                 \
       return DiagnosedSilenceableFailure::success();                           \
static LogicalResult applyTilingToAll(
   transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
   for (unsigned int i = 0; i < numLoops; ++i)
     transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);
 
   return success();
 }
 
 DiagnosedSilenceableFailure
-transform::FuseOp::apply(mlir::transform::TransformResults &transformResults,
+transform::FuseOp::apply(transform::TransformRewriter &rewriter,
+                         mlir::transform::TransformResults &transformResults,
                          mlir::transform::TransformState &state) {
   SmallVector<int64_t> tileSizes = extractFromI64ArrayAttr(getTileSizes());
   SmallVector<int64_t> tileInterchange =
       extractFromI64ArrayAttr(getTileInterchange());
 
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.interchangeVector = tileInterchange;
   tilingOptions = tilingOptions.setTileSizes(tileSizes);
   scf::SCFTileAndFuseOptions tileAndFuseOptions;
   tileAndFuseOptions.tilingOptions = tilingOptions;
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   LogicalResult result = applyTilingToAll(
       rewriter, getOperation(), state.getPayloadOps(getTarget()),
       tileSizes.size() - llvm::count(tileSizes, 0), transformResults,
       [&](TilingInterface tilingInterfaceOp)
           -> FailureOr<scf::SCFTileAndFuseResult> {
         return tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
             rewriter, tilingInterfaceOp, tileAndFuseOptions);
       });
static void forciblyReplaceReferencedPayloadOperation(TransformState &state,
                                                       Operation *replacement) {
   UnsafeOpReplacementStateExtension extension(state);
   // This may return failure if the payload is not associated with any handle,
   // ignore that.
   (void)extension.doReplacePayloadOp(payload, replacement);
 }
 
 DiagnosedSilenceableFailure
-transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
+transform::FuseIntoContainingOp::apply(transform::TransformRewriter &rewriter,
+                                       transform::TransformResults &results,
                                        transform::TransformState &state) {
   SmallVector<Operation *> fusedOps;
   auto producerOps = state.getPayloadOps(getProducerOp());
   auto containingOps = state.getPayloadOps(getContainingOp());
   if (!llvm::hasSingleElement(containingOps)) {
     return emitDefiniteFailure()
            << "requires exactly one containing_op handle (got "
            << llvm::range_size(containingOps) << ")";
for (const auto &it : enumerate(remainingProducers)) {
         if (numUsesInContainingOp == 1)
           remainingProducers.erase(remainingProducers.begin() + it.index());
         return producerOp;
       }
     }
     return failure();
   };
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   while (!remainingProducers.empty()) {
     auto nextProducer = getNextProducer();
     if (failed(nextProducer)) {
       return mlir::emitSilenceableFailure(containingOp->getLoc())
              << "could not find next producer to fuse into container";
     }
 
     Operation *producerOp = *nextProducer;
void transform::FuseIntoContainingOp::getEffects(
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::GeneralizeOp::applyToOne(LinalgOp target,
+transform::GeneralizeOp::applyToOne(transform::TransformRewriter &rewriter,
+                                    LinalgOp target,
                                     transform::ApplyToEachResultList &results,
                                     transform::TransformState &state) {
   // Exit early if no transformation is needed.
   if (isa<GenericOp>(target)) {
     results.push_back(target);
     return DiagnosedSilenceableFailure::success();
   }
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<LinalgOp> generic = generalizeNamedOp(rewriter, target);
   if (succeeded(generic)) {
     results.push_back(generic->getOperation());
     return DiagnosedSilenceableFailure::success();
   }
   return emitDefaultSilenceableFailure(target);
 }
 
 //===----------------------------------------------------------------------===//
 // InterchangeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::InterchangeOp::applyToOne(GenericOp target,
+transform::InterchangeOp::applyToOne(transform::TransformRewriter &rewriter,
+                                     GenericOp target,
                                      transform::ApplyToEachResultList &results,
                                      transform::TransformState &state) {
   ArrayRef<int64_t> interchangeVector = getIteratorInterchange();
   // Exit early if no transformation is needed.
   if (interchangeVector.empty()) {
     results.push_back(target);
     return DiagnosedSilenceableFailure::success();
   }
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   FailureOr<GenericOp> res =
       interchangeGenericOp(rewriter, target,
                            SmallVector<unsigned>(interchangeVector.begin(),
                                                  interchangeVector.end()));
   if (failed(res))
     return DiagnosedSilenceableFailure::definiteFailure();
   results.push_back(res->getOperation());
   return DiagnosedSilenceableFailure::success();
LogicalResult transform::InterchangeOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LowerPackOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::LowerPackOp::applyToOne(
-    tensor::PackOp target, transform::ApplyToEachResultList &transformResults,
+    transform::TransformRewriter &rewriter, tensor::PackOp target,
+    transform::ApplyToEachResultList &transformResults,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<LowerPackResult> res = lowerPack(rewriter, target);
   if (failed(res)) {
     return mlir::emitSilenceableFailure(target->getLoc())
            << "cannot lower to pad + expand + transpose";
   }
   transformResults.push_back(res->padOp);
   transformResults.push_back(res->expandShapeOp);
   transformResults.push_back(res->transposeOp);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LowerUnPackOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::LowerUnPackOp::applyToOne(
-    tensor::UnPackOp target, transform::ApplyToEachResultList &transformResults,
+    transform::TransformRewriter &rewriter, tensor::UnPackOp target,
+    transform::ApplyToEachResultList &transformResults,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<LowerUnPackOpResult> res = lowerUnPack(rewriter, target);
   if (failed(res)) {
     return mlir::emitSilenceableFailure(target->getLoc())
            << "cannot rewrite to pad + expand + transpose";
   }
   transformResults.push_back(res->emptyOp);
   transformResults.push_back(res->transposeOp);
void transform::MatchOp::build(OpBuilder &builder, OperationState &result,
                                ArrayRef<StringRef> opNames) {
   result.addOperands(target);
   result.addAttribute(MatchOp::getOpsAttrName(result.name),
                       builder.getStrArrayAttr(opNames));
   result.addTypes(resultTypes);
 }
 
 DiagnosedSilenceableFailure
-transform::MatchOp::apply(transform::TransformResults &results,
+transform::MatchOp::apply(transform::TransformRewriter &rewriter,
+                          transform::TransformResults &results,
                           transform::TransformState &state) {
   llvm::StringSet<> strs;
   if (getOps().has_value())
     strs.insert(getOps()->getAsValueRange<StringAttr>().begin(),
                 getOps()->getAsValueRange<StringAttr>().end());
 
   auto payloadOps = state.getPayloadOps(getTarget());
   if (!llvm::hasSingleElement(payloadOps)) {
static ParseResult parseMultitileSizesTypes(OpAsmParser &parser,
   }
   targetType = funcType.getInput(0);
   lowSizeType = highSizeType = splitPointType = funcType.getResult(0);
 
   return success();
 }
 
 DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
-    LinalgOp target, transform::ApplyToEachResultList &results,
-    TransformState &state) {
+    transform::TransformRewriter &rewriter, LinalgOp target,
+    transform::ApplyToEachResultList &results, TransformState &state) {
   if (isa<TransformParamTypeInterface>(getLowSize().getType())) {
     if (target.hasDynamicShape()) {
       auto diag = emitSilenceableError()
                   << "cannot compute parametric tile sizes for dynamically "
                      "shaped payload op";
       diag.attachNote(target->getLoc()) << "payload op";
       return diag;
     }

 }
 
 SmallVector<OpFoldResult> transform::PackOp::getMixedPackedSizes() {
   Builder b(getContext());
   return getMixedValues(getStaticPackedSizes(), getPackedSizes(), b);
 }
 
 DiagnosedSilenceableFailure
-transform::PackOp::apply(transform::TransformResults &transformResults,
+transform::PackOp::apply(transform::TransformRewriter &rewriter,
+                         transform::TransformResults &transformResults,
                          transform::TransformState &state) {
   auto targetOps = state.getPayloadOps(getTarget());
   // If nothing to pack, propagate success.
   if (std::empty(targetOps)) {
     transformResults.set(cast<OpResult>(getPackedOp()),
                          ArrayRef<Operation *>({}));
     return DiagnosedSilenceableFailure::success();
   }
return emitSilenceableError()
            << ")";
   }
 
   // Unpack handles to constants or actual SSA index values.
   SmallVector<OpFoldResult> packedSizes;
   DiagnosedSilenceableFailure status = unpackSingleIndexResultPayloadOperations(
       state, *this, packedSizes, getMixedPackedSizes());
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(linalgOp);
   FailureOr<PackResult> maybeResult = pack(rewriter, linalgOp, packedSizes);
   if (failed(maybeResult))
     return emitDefiniteFailure("data tiling failed");
 
   transformResults.set(cast<OpResult>(getPackedOp()),
                        {maybeResult->packedLinalgOp.getOperation()});
   return DiagnosedSilenceableFailure::success();
packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
   // TODO: If we wanted to give the genericOp a name after packing, after
   // calling `pack` would be a good time. One would still need to check that
   // `containsMostMinorMatmul(packingRes->packedLinalgOp)` is true, since we
   // also allow degenerate matmul cases (i.e. matvec, dot).
   return linalg::pack(rewriter, genericOp, adjustedPackedSizes);
 }
 
 DiagnosedSilenceableFailure
-PackGreedilyOp::apply(transform::TransformResults &transformResults,
+PackGreedilyOp::apply(transform::TransformRewriter &rewriter,
+                      transform::TransformResults &transformResults,
                       transform::TransformState &state) {
   SmallVector<Operation *> results;
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   for (Operation *op : state.getPayloadOps(getTarget())) {
     auto linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       continue;
     // linalgOp will be replaced and the insertion point may be invalidated if
     // we set it before -> set it after.
     rewriter.setInsertionPointAfter(linalgOp);
     // Failing to pack greedily is perfectly fine.
if (std::is_same<RelayoutOpTy, tensor::PackOp>::value) {
     return permutation.size() == op.getSourceRank() &&
            isPermutationVector(permutation);
   }
   return permutation.size() == op.getDestRank() &&
          isPermutationVector(permutation);
 }
 
 DiagnosedSilenceableFailure
-transform::PackTransposeOp::apply(transform::TransformResults &transformResults,
+transform::PackTransposeOp::apply(transform::TransformRewriter &rewriter,
+                                  transform::TransformResults &transformResults,
                                   transform::TransformState &state) {
   auto packOrUnpackOps = state.getPayloadOps(getTargetPackOrUnPackOp());
   auto linalgOps = state.getPayloadOps(getTargetLinalgOp());
   // Step 1. If nothing to pack, propagate success.
   if (std::empty(packOrUnpackOps)) {
     transformResults.set(cast<OpResult>(getPackedOp()), {});
     transformResults.set(cast<OpResult>(getPackOp()), {});
     transformResults.set(cast<OpResult>(getUnPackOp()), {});
for (auto permType : {OuterOrInnerPerm::Outer, OuterOrInnerPerm::Inner}) {
     }
   }
 
   // From here on, packOp and linalgOp are always present, unPackOp may or may
   // not be present.
   assert(packOp && linalgOp && "unexpected null op");
 
   // Step 3. Actually transpose the ops.
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   FailureOr<PackTransposeResult> res = packTranspose(
       rewriter, packOp, linalgOp, unPackOp, getOuterPerm(), getInnerPerm());
   // Preconditions have been checked, it is an error to fail here.
   assert(succeeded(res) && "unexpected packTranspose failure");
 
   // Step 4. Return results.
   transformResults.set(cast<OpResult>(getPackOp()), {res->transposedPackOp});
   transformResults.set(cast<OpResult>(getPackedOp()),
                        {res->transposedLinalgOp});
   if (unPackOp) {
     transformResults.set(cast<OpResult>(getUnPackOp()),
                          {res->transposedUnPackOp});
   } else {
     transformResults.set(cast<OpResult>(getUnPackOp()), {});
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::PadOp::applyToOne(LinalgOp target,
+transform::PadOp::applyToOne(transform::TransformRewriter &rewriter,
+                             LinalgOp target,
                              transform::ApplyToEachResultList &results,
                              transform::TransformState &state) {
   // Convert the integer packing flags to booleans.
   SmallVector<bool> packPaddings;
   for (int64_t packPadding : extractFromI64ArrayAttr(getPackPaddings()))
     packPaddings.push_back(static_cast<bool>(packPadding));
 
   // Convert the padding values to attributes.
transform::PadOp::applyToOne(transform::TransformRewriter &rewriter,
   }
 
   // Extract the transpose vectors.
   SmallVector<SmallVector<int64_t>> transposePaddings;
   for (Attribute transposeVector : cast<ArrayAttr>(getTransposePaddings()))
     transposePaddings.push_back(
         extractFromI64ArrayAttr(cast<ArrayAttr>(transposeVector)));
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   LinalgOp paddedOp;
   SmallVector<int64_t> paddingDimensions =
       extractFromI64ArrayAttr(getPaddingDimensions());
   SmallVector<int64_t> padToMultipleOf(paddingDimensions.size(), 1);
   if (getPadToMultipleOf().has_value())
     padToMultipleOf = extractFromI64ArrayAttr(*getPadToMultipleOf());
   FailureOr<SmallVector<Value>> result =
       rewriteAsPaddedOp(rewriter, target, paddingDimensions, padToMultipleOf,
LogicalResult transform::PadOp::verify() {
   return success();
 }
 
 //===---------------------------------------------------------------------===//
 // HoistPadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::HoistPadBuildPackingLoopNestOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
   auto targetOps = state.getPayloadOps(getTarget());
   auto loopOps = state.getPayloadOps(getLoop());
   if (!llvm::hasSingleElement(targetOps) || !llvm::hasSingleElement(loopOps)) {
     return emitDefiniteFailure()
            << "requires exactly one target and one loop handle (got "
            << llvm::range_size(targetOps) << " and "
            << llvm::range_size(loopOps) << ")";
   }
 
   auto padOp = dyn_cast_or_null<tensor::PadOp>(*targetOps.begin());
   auto loopOp = dyn_cast_or_null<scf::ForOp>(*loopOps.begin());
   if (!padOp || !loopOp)
     return emitDefiniteFailure() << "requires exactly 2 non-null handles";
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   FailureOr<linalg::detail::PackingResult> result =
       linalg::detail::buildPackingLoopNest(rewriter, padOp, loopOp,
                                            getTranspose());
   if (failed(result))
     return emitDefiniteFailure() << "could not build packing loop nest";
 
   if (result->clonedLoopIvs.empty()) {
     transformResults.set(cast<OpResult>(getPackingLoop()),
void transform::HoistPadBuildPackingLoopNestOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getTarget(), effects);
   transform::onlyReadsHandle(getLoop(), effects);
   transform::producesHandle(getPackingLoop(), effects);
   transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure
-transform::HoistPadOp::applyToOne(tensor::PadOp target,
+transform::HoistPadOp::applyToOne(transform::TransformRewriter &rewriter,
+                                  tensor::PadOp target,
                                   transform::ApplyToEachResultList &results,
                                   transform::TransformState &state) {
   tensor::PadOp hoistedPadOp;
   SmallVector<GenericOp> transposeOps;
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   FailureOr<Value> result =
       hoistPaddingOnTensors(rewriter, target, getNumLoops(), getTranspose(),
                             hoistedPadOp, transposeOps);
   if (succeeded(result)) {
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
     // TODO: clean this up and stop "pad and hoist" behavior more globally now
LogicalResult transform::HoistPadOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::PromoteOp::applyToOne(LinalgOp target,
+transform::PromoteOp::applyToOne(transform::TransformRewriter &rewriter,
+                                 LinalgOp target,
                                  transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {
   LinalgPromotionOptions promotionOptions;
   if (!getOperandsToPromote().empty())
     promotionOptions = promotionOptions.setOperandsToPromote(
         extractFromI64ArrayAttr(getOperandsToPromote()));
   if (getUseFullTilesByDefault())
     promotionOptions = promotionOptions.setUseFullTileBuffersByDefault(
if (getMapping().has_value()) {
     } else {
       return emitDefaultDefiniteFailure(target);
     }
   }
 
   if (failed(promoteSubviewsPrecondition(target, promotionOptions)))
     return emitDefaultDefiniteFailure(target);
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<LinalgOp> res = promoteSubViews(rewriter, target, promotionOptions);
   if (failed(res))
     return emitDefaultDefiniteFailure(target);
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::ReplaceOp::apply(TransformResults &transformResults,
+transform::ReplaceOp::apply(transform::TransformRewriter &rewriter,
+                            TransformResults &transformResults,
                             TransformState &state) {
   auto payload = state.getPayloadOps(getTarget());
 
   // Check for invalid targets.
   for (Operation *target : payload) {
     if (target->getNumOperands() > 0)
       return emitDefiniteFailure() << "expected target without operands";
     if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
         target->getNumRegions() > 0)
       return emitDefiniteFailure()
              << "expected target that is isolated from above";
   }
 
   // Clone and replace.
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   Operation *pattern = &getBodyRegion().front().front();
   SmallVector<Operation *> replacements;
   for (Operation *target : payload) {
     if (getOperation()->isAncestor(target))
       continue;
     rewriter.setInsertionPoint(target);
     Operation *replacement = rewriter.clone(*pattern);
     rewriter.replaceOp(target, replacement->getResults());
LogicalResult transform::ReplaceOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ScalarizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::ScalarizeOp::applyToOne(LinalgOp target,
+transform::ScalarizeOp::applyToOne(transform::TransformRewriter &rewriter,
+                                   LinalgOp target,
                                    transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.setTileSizeComputationFunction([&](OpBuilder &b, Operation *) {
     SmallVector<Value, 4> tileSizes;
     Location loc = target.getLoc();
     SmallVector<OpFoldResult> allShapeSizes =
         target.createFlatListOfOperandDims(b, loc);
     AffineMap map = target.getShapesToLoopsMap();
     if (!map)
       return tileSizes;
-    TrackingListener listener(state, *this);
-    IRRewriter rewriter(getContext(), &listener);
     SmallVector<OpFoldResult> shapeSizes =
         affine::makeComposedFoldedMultiResultAffineApply(rewriter, loc, map,
                                                          allShapeSizes);
     // If the shape size is dynamic, tile by 1.
     // Otherwise, do not tile (i.e. tile size 0).
     for (OpFoldResult shapeSize : shapeSizes) {
       tileSizes.push_back(getConstantIntValue(shapeSize)
                               ? b.create<arith::ConstantIndexOp>(loc, 0)
                               : b.create<arith::ConstantIndexOp>(loc, 1));
     }
     return tileSizes;
   });
   SmallVector<int64_t> emptyTileSizes;
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<scf::SCFTilingResult> maybeTilingResult = tileUsingSCFForOp(
       rewriter, cast<TilingInterface>(target.getOperation()), tilingOptions);
   if (failed(maybeTilingResult))
     return emitDefaultDefiniteFailure(target);
 
   if (target->getNumResults())
     rewriter.replaceOp(target, maybeTilingResult->replacements);
   else
     rewriter.eraseOp(target);
 
   results.reserve(maybeTilingResult->tiledOps.size());
   for (Operation *tiled : maybeTilingResult->tiledOps)
     results.push_back(tiled);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // RewriteInDestinationPassingStyleOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::RewriteInDestinationPassingStyleOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   SmallVector<Operation *> res;
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<Operation *> maybeResult =
       TypeSwitch<Operation *, FailureOr<Operation *>>(target)
           .Case<tensor::FromElementsOp, tensor::GenerateOp, tensor::PadOp>(
               [&rewriter](auto op) {
                 return rewriteInDestinationPassingStyle(rewriter, op);
               });
   if (failed(maybeResult))
     return emitDefaultSilenceableFailure(target);
   results.push_back(*maybeResult);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
-                                           TransformState &state) {
+DiagnosedSilenceableFailure
+SplitOp::apply(transform::TransformRewriter &rewriter,
+               TransformResults &results, TransformState &state) {
   // Collect the dynamic split points if provided.
   SmallVector<Operation *> payload =
       llvm::to_vector(state.getPayloadOps(getTarget()));
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   SmallVector<OpFoldResult> splitPoints;
   splitPoints.reserve(payload.size());
   if (getDynamicSplitPoint()) {
     auto diag = DiagnosedSilenceableFailure::success();
     if (isa<TransformHandleTypeInterface>(getDynamicSplitPoint().getType())) {
       splitPoints = llvm::to_vector(llvm::map_range(
           state.getPayloadOps(getDynamicSplitPoint()), [&](Operation *op) {
             if (op->getNumResults() != 1 ||
if (useAlloc) {
     result.addAttribute(SplitReductionOp::getUseAllocAttrName(result.name),
                         builder.getUnitAttr());
   }
   auto resultType = transform::AnyOpType::get(ctx);
   result.addTypes({resultType, resultType, resultType, resultType});
 }
 
 DiagnosedSilenceableFailure transform::SplitReductionOp::applyToOne(
-    LinalgOp target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, LinalgOp target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   ControlSplitReductionFn splitFn = [&](LinalgOp) {
     return linalg::SplitReductionOptions{int64_t(getSplitFactor()),
                                          unsigned(getInsertSplitDimension()),
                                          bool(getInnerParallel())};
   };
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<SplitReductionResult> splitResult =
       (getUseScalingAlgorithm())
           ? splitReductionByScaling(rewriter, target, splitFn, getUseAlloc())
           : splitReduction(rewriter, target, splitFn, getUseAlloc());
   if (failed(splitResult))
     return emitDefaultDefiniteFailure(target);
 
void transform::TileReductionUsingScfOp::build(
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
   build(builder, result,
         /*resultTypes=*/TypeRange{opTy, opTy, opTy, opTy},
         /*target=*/target,
         /*tile_sizes=*/staticTileSizesAttr);
 }
 
 DiagnosedSilenceableFailure transform::TileReductionUsingScfOp::applyToOne(
-    LinalgOp target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, LinalgOp target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   FailureOr<scf::SCFReductionTilingResult> result = scf::tileReductionUsingScf(
       rewriter, cast<PartialReductionOpInterface>(target.getOperation()),
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getTileSizes())));
 
   if (failed(result))
     return emitDefaultSilenceableFailure(target);
   results.push_back(result->loops.front());
build(builder, result,
         /*resultTypes=*/TypeRange{opTy, opTy, opTy, opTy},
         /*target=*/target,
         /*num_threads=*/staticNumThreadsAttr,
         /*tile_sizes=*/staticTileSizesAttr,
         /*mapping=*/mapping);
 }
 
 DiagnosedSilenceableFailure transform::TileReductionUsingForallOp::applyToOne(
-    LinalgOp target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, LinalgOp target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   SmallVector<OpFoldResult> numThreads =
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getNumThreads()));
   SmallVector<OpFoldResult> tileSizes =
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getTileSizes()));
   FailureOr<linalg::ForallReductionTilingResult> result =
       linalg::tileReductionUsingForall(
           rewriter, cast<PartialReductionOpInterface>(target.getOperation()),
build(builder, result, /*tiled_linalg_op=*/target.getType(),
         /*loops=*/resultTypes,
         /*target=*/target,
         /*dynamic_sizes=*/dynamicTileSizes,
         /*static_sizes=*/staticTileSizesAttr,
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange));
 }
 
 DiagnosedSilenceableFailure
-transform::TileOp::apply(TransformResults &transformResults,
+transform::TileOp::apply(transform::TransformRewriter &rewriter,
+                         TransformResults &transformResults,
                          TransformState &state) {
   ArrayRef<int64_t> tileSizes = getStaticSizes();
 
   SmallVector<Operation *> targets =
       llvm::to_vector(state.getPayloadOps(getTarget()));
   SmallVector<SmallVector<Operation *>> dynamicSizeProducers;
   SmallVector<SmallVector<int64_t>> paramSizes;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
if (!tileSizes.empty()) {
             sizes.push_back(dynamicSizes[index]->getResult(0));
           }
         }
         return sizes;
       });
     }
 
     tilingOptions.setInterchange(getInterchange());
-    TrackingListener listener(state, *this);
-    IRRewriter rewriter(getContext(), &listener);
     FailureOr<scf::SCFTilingResult> maybeTilingResult =
         tileUsingSCFForOp(rewriter, tilingInterface, tilingOptions);
     if (failed(maybeTilingResult))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     if (dpsInterface.hasBufferSemantics())
       rewriter.eraseOp(op);
     else
if (failed(maybeTilingResult))
       return transformOp.emitDefaultSilenceableFailure(tileableOp);
     rewriter.replaceOp(tileableOp, maybeTilingResult->tileOp->getResults());
 
     tilingResult = *maybeTilingResult;
     return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
-transform::TileToForallOp::apply(transform::TransformResults &transformResults,
+transform::TileToForallOp::apply(transform::TransformRewriter &rewriter,
+                                 transform::TransformResults &transformResults,
                                  transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
     auto transformOp = cast<TransformOpInterface>(getOperation());
 
     // Result payload ops.
     SmallVector<Operation *> tileOps;
     SmallVector<Operation *> tiledOps;
 
     // Unpack handles.
     SmallVector<OpFoldResult> mixedNumThreads;
     DiagnosedSilenceableFailure status =
         getPackedNumThreads()
             ? unpackSingleIndexResultPayloadOperations(
                   state, transformOp, mixedNumThreads, getPackedNumThreads())
             : unpackSingleIndexResultPayloadOperations(
                   state, transformOp, mixedNumThreads, getMixedNumThreads());
     if (!status.succeeded())
       return status;
     SmallVector<OpFoldResult> mixedTileSizes;
     status = getPackedTileSizes()
                  ? unpackSingleIndexResultPayloadOperations(
                        state, transformOp, mixedTileSizes, getPackedTileSizes())
                  : unpackSingleIndexResultPayloadOperations(
                        state, transformOp, mixedTileSizes, getMixedTileSizes());
     if (!status.succeeded())
       return status;
 
     for (Operation *target : state.getPayloadOps(getTarget())) {
       linalg::ForallTilingResult tilingResult;
       DiagnosedSilenceableFailure diag = tileToForallOpImpl(
           rewriter, state, transformOp, target, mixedNumThreads, mixedTileSizes,
           getMapping(), tilingResult);
       if (!diag.succeeded())
       return diag;
     tileOps.push_back(tilingResult.tileOp);
     tiledOps.push_back(tilingResult.tiledOp);
     }
 
   transformResults.set(cast<OpResult>(getForallOp()), tileOps);
   transformResults.set(cast<OpResult>(getTiledOp()), tiledOps);
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::TileToForallOp::getEffects(
build(builder, result,
         /*loops=*/resultTypes,
         /*target=*/target,
         /*dynamic_sizes=*/dynamicTileSizes,
         /*static_sizes=*/staticTileSizesAttr,
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange));
 }
 
 DiagnosedSilenceableFailure
-transform::TileToScfForOp::apply(TransformResults &transformResults,
+transform::TileToScfForOp::apply(transform::TransformRewriter &rewriter,
+                                 TransformResults &transformResults,
                                  TransformState &state) {
   ArrayRef<int64_t> tileSizes = getStaticSizes();
 
   SmallVector<Operation *> targets =
       llvm::to_vector(state.getPayloadOps(getTarget()));
   SmallVector<SmallVector<Operation *>> dynamicSizeProducers;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
   for (Value dynamicSizeProducerHandle : getDynamicSizes()) {
if (!tileSizes.empty()) {
                     dynamicSizeProducers[dynamicIdx++][index]->getResult(0));
               }
             }
             return sizes;
           });
     }
 
     tilingOptions.setInterchange(getInterchange());
-    TrackingListener listener(state, *this);
-    IRRewriter rewriter(getContext(), &listener);
     FailureOr<scf::SCFTilingResult> tilingResult =
         tileUsingSCFForOp(rewriter, tilingInterfaceOp, tilingOptions);
     if (failed(tilingResult))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     rewriter.replaceOp(tilingInterfaceOp, tilingResult->replacements);
 
     tiled.append(tilingResult->tiledOps);

 private:
   /// Controls whether to vectorize `tensor.extract` when the input tensor is
   /// rank >= 2.
   bool vectorizeNDExtract = false;
 };
 } // namespace
 
 DiagnosedSilenceableFailure
-transform::VectorizeOp::applyToOne(Operation *target,
+transform::VectorizeOp::applyToOne(transform::TransformRewriter &rewriter,
+                                   Operation *target,
                                    transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     auto diag = this->emitOpError("requires isolated-from-above targets");
     diag.attachNote(target->getLoc()) << "non-isolated target";
     return DiagnosedSilenceableFailure::definiteFailure();
   }
 
transform::VectorizeOp::applyToOne(transform::TransformRewriter &rewriter,
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // MaskedVectorizeOp
 //===----------------------------------------------------------------------===//
 DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
+    transform::TransformRewriter &rewriter,
     mlir::transform::TransformResults &transformResults,
     mlir::transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   auto targets = state.getPayloadOps(getTarget());
   if (std::empty(targets))
     return DiagnosedSilenceableFailure::success();
 
   SmallVector<int64_t> vectorSizes;
   for (OpFoldResult sz : getMixedVectorSizes()) {
     if (sz.is<Attribute>()) {
       auto attr = sz.get<Attribute>();

 }
 
 //===----------------------------------------------------------------------===//
 // HoistRedundantVectorTransfersOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::HoistRedundantVectorTransfersOp::applyToOne(
-    func::FuncOp target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, func::FuncOp target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   // WARNING: This hoisting does not model parallelism and is generally
   // incorrect when used on distributed loops with memref semantics!
   // TODO: obsolete and should be retired.
   linalg::hoistRedundantVectorTransfers(target);
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // ConvertConv2DToImg2ColOp.
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::ConvertConv2DToImg2ColOp::applyToOne(
-    linalg::LinalgOp target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, linalg::LinalgOp target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   rewriter.setInsertionPoint(target);
   auto maybeTransformed =
       TypeSwitch<Operation *, FailureOr<std::pair<Operation *, Operation *>>>(
           target)
           .Case([&](linalg::Conv2DNhwcHwcfOp op) {
             return rewriteInIm2Col(rewriter, op);
           })
           .Case([&](linalg::DepthwiseConv2DNhwcHwcOp op) {

 }
 
 //===----------------------------------------------------------------------===//
 // HoistRedundantTensorSubsetsOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::HoistRedundantTensorSubsetsOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   auto forOp = dyn_cast<scf::ForOp>(target);
   if (forOp) {
     linalg::hoistRedundantSubsetExtractInsert(rewriter, forOp);
     return DiagnosedSilenceableFailure::success();
   }
 
   // TODO: walking in some reverse / inside-out order would be more efficient
   // and would capture more cases.
rewriter.replaceOpWithNewOp<OpTy>(
       target, copied, target.getDest(), target.getMixedOffsets(),
       target.getMixedSizes(), target.getMixedStrides());
 
   results.push_back(copied.getDefiningOp());
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure transform::InsertSliceToCopyOp::applyToOne(
-    Operation *targetOp, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *targetOp,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
 
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(targetOp->getContext(), &listener);
   rewriter.setInsertionPoint(targetOp);
   if (auto target = dyn_cast<tensor::InsertSliceOp>(targetOp))
     return doit(rewriter, target, results, state);
   if (auto target = dyn_cast<tensor::ParallelInsertSliceOp>(targetOp))
     return doit(rewriter, target, results, state);
 
   DiagnosedSilenceableFailure diag =
       emitSilenceableError()

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

void transform::ApplyResolveRankedShapedTypeResultDimsPatternsOp::
   memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefMultiBufferOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
   SmallVector<Operation *> results;
-  IRRewriter rewriter(getContext());
   for (Operation *op : state.getPayloadOps(getTarget())) {
     bool canApplyMultiBuffer = true;
     auto target = cast<memref::AllocOp>(op);
     LLVM_DEBUG(DBGS() << "Start multibuffer transform op: " << target << "\n";);
     // Skip allocations not used in a loop.
     for (Operation *user : target->getUsers()) {
       if (isa<memref::DeallocOp>(user))
         continue;
DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefMakeLoopIndependentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMakeLoopIndependentOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   // Gather IVs.
   SmallVector<Value> ivs;
   Operation *nextOp = target;
   for (uint64_t i = 0, e = getNumLoops(); i < e; ++i) {
     nextOp = nextOp->getParentOfType<scf::ForOp>();
     if (!nextOp) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "could not find " << i
                                          << "-th enclosing loop";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     ivs.push_back(cast<scf::ForOp>(nextOp).getInductionVar());
   }
 
   // Rewrite IR.
-  IRRewriter rewriter(target->getContext());
   FailureOr<Value> replacement = failure();
   if (auto allocaOp = dyn_cast<memref::AllocaOp>(target)) {
     replacement = memref::replaceWithIndependentOp(rewriter, allocaOp, ivs);
   } else {
     DiagnosedSilenceableFailure diag = emitSilenceableError()
                                        << "unsupported target op";
     diag.attachNote(target->getLoc()) << "target op";
     return diag;

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

void transform::ApplyForLoopCanonicalizationPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   scf::populateSCFForLoopCanonicalizationPatterns(patterns);
 }
 
 //===----------------------------------------------------------------------===//
 // GetParentForOp
 //===----------------------------------------------------------------------===//
 DiagnosedSilenceableFailure
-transform::GetParentForOp::apply(transform::TransformResults &results,
+transform::GetParentForOp::apply(transform::TransformRewriter &rewriter,
+                                 transform::TransformResults &results,
                                  transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Operation *loop, *current = target;
     for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
       loop = getAffine()
                  ? current->getParentOfType<AffineForOp>().getOperation()
                  : current->getParentOfType<scf::ForOp>().getOperation();
b.inlineRegionBefore(op->getRegions().front(), clonedRegion,
                          clonedRegion.end());
     b.create<scf::YieldOp>(op->getLoc(), clonedOp->getResults());
   }
   b.replaceOp(op, executeRegionOp.getResults());
   return executeRegionOp;
 }
 
 DiagnosedSilenceableFailure
-transform::LoopOutlineOp::apply(transform::TransformResults &results,
+transform::LoopOutlineOp::apply(transform::TransformRewriter &rewriter,
+                                transform::TransformResults &results,
                                 transform::TransformState &state) {
   SmallVector<Operation *> functions;
   SmallVector<Operation *> calls;
   DenseMap<Operation *, SymbolTable> symbolTables;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Location location = target->getLoc();
     Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(target);
-    IRRewriter rewriter(getContext());
     scf::ExecuteRegionOp exec = wrapInExecuteRegion(rewriter, target);
     if (!exec) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "failed to outline";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     func::CallOp call;
transform::LoopOutlineOp::apply(transform::TransformRewriter &rewriter,
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopPeelOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::LoopPeelOp::applyToOne(scf::ForOp target,
+transform::LoopPeelOp::applyToOne(transform::TransformRewriter &rewriter,
+                                  scf::ForOp target,
                                   transform::ApplyToEachResultList &results,
                                   transform::TransformState &state) {
   scf::ForOp result;
-  IRRewriter rewriter(target->getContext());
   // This helper returns failure when peeling does not occur (i.e. when the IR
   // is not modified). This is not a failure for the op as the postcondition:
   //    "the loop trip count is divisible by the step"
   // is valid.
   LogicalResult status =
       scf::peelForLoopAndSimplifyBounds(rewriter, target, result);
   // TODO: Return both the peeled loop and the remainder loop.
   results.push_back(failed(status) ? target : result);
for (const auto &it : wrappedSchedule) {
     for (Operation *op : it.second) {
       unsigned cycle = opCycles[op];
       schedule.emplace_back(op, cycle / iterationInterval);
     }
   }
 }
 
 DiagnosedSilenceableFailure
-transform::LoopPipelineOp::applyToOne(scf::ForOp target,
+transform::LoopPipelineOp::applyToOne(transform::TransformRewriter &rewriter,
+                                      scf::ForOp target,
                                       transform::ApplyToEachResultList &results,
                                       transform::TransformState &state) {
   scf::PipeliningOption options;
   options.getScheduleFn =
       [this](scf::ForOp forOp,
              std::vector<std::pair<Operation *, unsigned>> &schedule) mutable {
         loopScheduling(forOp, schedule, getIterationInterval(),
                        getReadLatency());
       };
   scf::ForLoopPipeliningPattern pattern(options, target->getContext());
-  IRRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<scf::ForOp> patternResult =
       scf::pipelineForLoop(rewriter, target, options);
   if (succeeded(patternResult)) {
     results.push_back(*patternResult);
     return DiagnosedSilenceableFailure::success();
   }
   return emitDefaultSilenceableFailure(target);
 }
 
 //===----------------------------------------------------------------------===//
 // LoopUnrollOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::LoopUnrollOp::applyToOne(Operation *op,
+transform::LoopUnrollOp::applyToOne(transform::TransformRewriter &rewriter,
+                                    Operation *op,
                                     transform::ApplyToEachResultList &results,
                                     transform::TransformState &state) {
   LogicalResult result(failure());
   if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))
     result = loopUnrollByFactor(scfFor, getFactor());
   else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
     result = loopUnrollByFactor(affineFor, getFactor());
 
   if (failed(result)) {
     DiagnosedSilenceableFailure diag = emitSilenceableError()
                                        << "failed to unroll";
     return diag;
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopCoalesceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::LoopCoalesceOp::applyToOne(Operation *op,
+transform::LoopCoalesceOp::applyToOne(transform::TransformRewriter &rewriter,
+                                      Operation *op,
                                       transform::ApplyToEachResultList &results,
                                       transform::TransformState &state) {
   LogicalResult result(failure());
   if (scf::ForOp scfForOp = dyn_cast<scf::ForOp>(op))
     result = coalescePerfectlyNestedLoops(scfForOp);
   else if (AffineForOp affineForOp = dyn_cast<AffineForOp>(op))
     result = coalescePerfectlyNestedLoops(affineForOp);
 
static void replaceOpWithRegion(RewriterBase &rewriter, Operation *op,
   Operation *terminator = block->getTerminator();
   ValueRange results = terminator->getOperands();
   rewriter.inlineBlockBefore(block, op, /*blockArgs=*/{});
   rewriter.replaceOp(op, results);
   rewriter.eraseOp(terminator);
 }
 
 DiagnosedSilenceableFailure transform::TakeAssumedBranchOp::applyToOne(
-    scf::IfOp ifOp, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, scf::IfOp ifOp,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  TrackingListener listener(state, *this);
-  IRRewriter rewriter(ifOp->getContext(), &listener);
   rewriter.setInsertionPoint(ifOp);
-
   Region &region =
       getTakeElseBranch() ? ifOp.getElseRegion() : ifOp.getThenRegion();
   if (!llvm::hasSingleElement(region)) {
     return emitDefiniteFailure()
            << "requires an scf.if op with a single-block "
            << ((getTakeElseBranch()) ? "`else`" : "`then`") << " region";
   }
   replaceOpWithRegion(rewriter, ifOp, region);

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

void transform::ApplyRewriteTensorOpsAsConstantPatternsOp::populatePatterns(
   tensor::populateRewriteAsConstantPatterns(patterns);
 }
 
 //===----------------------------------------------------------------------===//
 // MakeLoopIndependentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MakeLoopIndependentOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   // Gather IVs.
   SmallVector<Value> ivs;
   Operation *nextOp = target;
   for (uint64_t i = 0, e = getNumLoops(); i < e; ++i) {
     nextOp = nextOp->getParentOfType<scf::ForOp>();
     if (!nextOp) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "could not find " << i
                                          << "-th enclosing loop";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     ivs.push_back(cast<scf::ForOp>(nextOp).getInductionVar());
   }
 
   // Rewrite IR.
-  IRRewriter rewriter(target->getContext());
   FailureOr<Value> replacement = failure();
   if (auto padOp = dyn_cast<tensor::PadOp>(target)) {
     replacement = tensor::buildIndependentOp(rewriter, padOp, ivs);
   } else if (auto emptyOp = dyn_cast<tensor::EmptyOp>(target)) {
     replacement = tensor::buildIndependentOp(rewriter, emptyOp, ivs);
   } else {
     DiagnosedSilenceableFailure diag = emitSilenceableError()
                                        << "unsupported target op";

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

LogicalResult transform::TransformDialect::verifyOperationAttribute(
   if (attribute.getName().getValue() == kArgConsumedAttrName ||
       attribute.getName().getValue() == kArgReadOnlyAttrName) {
     if (!llvm::isa<UnitAttr>(attribute.getValue())) {
       return op->emitError()
              << attribute.getName() << " must be a unit attribute";
     }
     return success();
   }
+  if (attribute.getName().getValue() == kSilenceTrackingFailuresAttrName) {
+    if (!llvm::isa<UnitAttr>(attribute.getValue())) {
+      return op->emitError()
+             << attribute.getName() << " must be a unit attribute";
+    }
+    return success();
+  }
   return emitError(op->getLoc())
          << "unknown attribute: " << attribute.getName();
 }

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/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/STLExtras.h"
DiagnosedDefiniteFailure diag =
         emitDefiniteFailure(transform->getLoc())
         << "unexpectedly consumed a value that is not a handle as operand #"
         << opOperand->getOperandNumber();
     diag.attachNote(operand.getLoc())
         << "value defined here with type " << operand.getType();
     return diag;
   }
 
+  // Prepare rewriter and listener.
+  transform::ErrorCheckingTrackingListener trackingListener(*this, transform);
+  transform::TransformRewriter rewriter(transform->getContext(),
+                                        &trackingListener);
+
   // Compute the result but do not short-circuit the silenceable failure case as
   // we still want the handles to propagate properly so the "suppress" mode can
   // proceed on a best effort basis.
   transform::TransformResults results(transform->getNumResults());
-  DiagnosedSilenceableFailure result(transform.apply(results, *this));
+  DiagnosedSilenceableFailure result(transform.apply(rewriter, results, *this));
   compactOpHandles();
+
+  // Error handling: fail if transform or listener failed.
+  DiagnosedSilenceableFailure trackingFailure =
+      trackingListener.checkAndResetError();
+  if (!transform->hasTrait<ReportTrackingListenerFailuresOpTrait>() ||
+      transform->hasAttr(
+          transform::TransformDialect::kSilenceTrackingFailuresAttrName)) {
+    // Only report failures for ReportTrackingListenerFailuresOpTrait ops. Also
+    // do not report failures if the above mentioned attribute is set.
+    if (trackingFailure.isSilenceableFailure())
+      (void)trackingFailure.silence();
+    trackingFailure = DiagnosedSilenceableFailure::success();
+  }
+  if (!trackingFailure.succeeded()) {
+    if (result.succeeded()) {
+      result = std::move(trackingFailure);
+    } else {
+      // Transform op errors have precedence, report those first.
+      if (result.isSilenceableFailure())
+        result.attachNote() << "tracking listener also failed: "
+                            << trackingFailure.getMessage();
+      (void)trackingFailure.silence();
+    }
+  }
   if (result.isDefiniteFailure())
     return result;
 
   // If a silenceable failure was produced, some results may be unset, set them
   // to empty lists.
   if (result.isSilenceableFailure())
     results.setRemainingToEmpty(transform);
 
return params[resultNumber].data() != nullptr ||
          operations[resultNumber].data() != nullptr ||
          values[resultNumber].data() != nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // TrackingListener
 //===----------------------------------------------------------------------===//
 
+transform::TrackingListener::TrackingListener(TransformState &state,
+                                              TransformOpInterface op)
+    : TransformState::Extension(state), transformOp(op) {
+  if (op) {
+    for (OpOperand *opOperand : transformOp.getConsumedHandleOpOperands()) {
+      consumedHandles.insert(opOperand->get());
+    }
+  }
+}
+
 Operation *transform::TrackingListener::getCommonDefiningOp(ValueRange values) {
   Operation *defOp = nullptr;
   for (Value v : values) {
     // Skip empty values.
     if (!v)
       continue;
     if (!defOp) {
       defOp = v.getDefiningOp();
for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
     (void)replacePayloadValue(oldValue, newValue);
 
   // Replace op handle.
   SmallVector<Value> opHandles;
   if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles))) {
     // Op is not tracked.
     return;
   }
+
+  // Helper function to check if the current transform op consumes any handle
+  // that is mapped to `op`.
+  //
+  // Note: If a handle was consumed, there shouldn't be any alive users, so it
+  // is not really necessary to check for consumed handles. However, in case
+  // there are indeed alive handles that were consumed (which is undefined
+  // behavior) and a replacement op could not be found, we want to fail with a
+  // nicer error message: "op uses a handle invalidated..." instead of "could
+  // not find replacement op". This nicer error is produced later.
+  auto handleWasConsumed = [&] {
+    return llvm::any_of(opHandles,
+                        [&](Value h) { return consumedHandles.contains(h); });
+  };
+
+  // Helper function to check if the handle is alive.
   auto hasAliveUser = [&]() {
-    for (Value v : opHandles)
+    for (Value v : opHandles) {
       for (Operation *user : v.getUsers())
-        if (!happensBefore(user, transformOp))
+        if (user != transformOp && !happensBefore(user, transformOp))
           return true;
+    }
     return false;
   };
-  if (!hasAliveUser()) {
-    // The op is tracked but the corresponding handles are dead.
+
+  if (!hasAliveUser() || handleWasConsumed()) {
+    // The op is tracked but the corresponding handles are dead or were
+    // consumed. Drop the op form the mapping.
     (void)replacePayloadOp(op, nullptr);
     return;
   }
 
   FailureOr<Operation *> replacement = findReplacementOp(op, newValues);
   // If the op is tracked but no replacement op was found, send a
   // notification.
   if (failed(replacement)) {
void transform::ErrorCheckingTrackingListener::notifyPayloadReplacementNotFound(
   for (auto &&[index, value] : llvm::enumerate(values))
     status.attachNote(value.getLoc())
         << "[" << errorCounter << "] replacement value " << index;
 
   ++errorCounter;
 }
 
 //===----------------------------------------------------------------------===//
+// TransformRewriter
+//===----------------------------------------------------------------------===//
+
+transform::TransformRewriter::TransformRewriter(
+    MLIRContext *ctx, ErrorCheckingTrackingListener *listener)
+    : RewriterBase(ctx), listener(listener) {
+  setListener(listener);
+}
+
+bool transform::TransformRewriter::hasTrackingFailures() const {
+  return listener->failed();
+}
+
+/// Silence all tracking failures that have been encountered so far.
+void transform::TransformRewriter::silenceTrackingFailure() {
+  if (hasTrackingFailures()) {
+    DiagnosedSilenceableFailure status = listener->checkAndResetError();
+    (void)status.silence();
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // Utilities for TransformEachOpTrait.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 transform::detail::checkNestedConsumption(Location loc,
                                           ArrayRef<Operation *> targets) {
   for (auto &&[position, parent] : llvm::enumerate(targets)) {
     for (Operation *child : targets.drop_front(position + 1)) {

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

-//===- TransformDialect.cpp - Transform dialect operations ----------------===//
+//===- TransformOps.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/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"

 
 static void forwardEmptyOperands(Block *block, transform::TransformState &state,
                                  transform::TransformResults &results) {
   for (const auto &res : block->getParentOp()->getOpResults())
     results.set(res, {});
 }
 
 DiagnosedSilenceableFailure
-transform::AlternativesOp::apply(transform::TransformResults &results,
+transform::AlternativesOp::apply(transform::TransformRewriter &rewriter,
+                                 transform::TransformResults &results,
                                  transform::TransformState &state) {
   SmallVector<Operation *> originals;
   if (Value scopeHandle = getScope())
     llvm::append_range(originals, state.getPayloadOps(scopeHandle));
   else
     originals.push_back(state.getTopLevel());
 
   for (Operation *original : originals) {
LogicalResult transform::AlternativesOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AnnotateOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::AnnotateOp::apply(transform::TransformResults &results,
+transform::AnnotateOp::apply(transform::TransformRewriter &rewriter,
+                             transform::TransformResults &results,
                              transform::TransformState &state) {
   SmallVector<Operation *> targets =
       llvm::to_vector(state.getPayloadOps(getTarget()));
 
   Attribute attr = UnitAttr::get(getContext());
   if (auto paramH = getParam()) {
     ArrayRef<Attribute> params = state.getParams(paramH);
     if (params.size() != 1) {
void transform::AnnotateOp::getEffects(
   onlyReadsHandle(getParam(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // ApplyPatternsOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure
-transform::ApplyPatternsOp::applyToOne(Operation *target,
-                                       ApplyToEachResultList &results,
-                                       transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::ApplyPatternsOp::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    ApplyToEachResultList &results, transform::TransformState &state) {
   // Make sure that this transform is not applied to itself. Modifying the
   // transform IR while it is being interpreted is generally dangerous. Even
   // more so for the ApplyPatternsOp because the GreedyPatternRewriteDriver
   // performs many additional simplifications such as dead code elimination.
   DiagnosedSilenceableFailure payloadCheck =
       ensurePayloadIsSeparateFromTransform(*this, target);
   if (!payloadCheck.succeeded())
     return payloadCheck;
 
   // Gather all specified patterns.
   MLIRContext *ctx = target->getContext();
   RewritePatternSet patterns(ctx);
   if (!getRegion().empty()) {
     for (Operation &op : getRegion().front()) {
       cast<transform::PatternDescriptorOpInterface>(&op).populatePatterns(
           patterns);
     }
   }
 
   // Configure the GreedyPatternRewriteDriver.
-  ErrorCheckingTrackingListener listener(state, *this);
   GreedyRewriteConfig config;
-  config.listener = &listener;
+  config.listener =
+      static_cast<RewriterBase::Listener *>(rewriter.getListener());
 
   LogicalResult result = failure();
   if (target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     // Op is isolated from above. Apply patterns and also perform region
     // simplification.
     result = applyPatternsAndFoldGreedily(target, std::move(patterns), config);
   } else {
     // Manually gather list of ops because the other GreedyPatternRewriteDriver
DiagnosedSilenceableFailure transform::ApplyPatternsOp::applyToOne(
 
   // 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() {
   if (!getRegion().empty()) {
     for (Operation &op : getRegion().front()) {
       if (!isa<transform::PatternDescriptorOpInterface>(&op)) {
         InFlightDiagnostic diag = emitOpError()

 void transform::ApplyPatternsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getTarget(), effects);
   transform::modifiesPayload(effects);
 }
 
 void transform::ApplyPatternsOp::build(
     OpBuilder &builder, OperationState &result, Value target,
-    function_ref<void(OpBuilder &, Location)> bodyBuilder,
-    bool failOnPayloadReplacementNotFound) {
+    function_ref<void(OpBuilder &, Location)> bodyBuilder) {
   result.addOperands(target);
-  result.getOrAddProperties<Properties>()
-      .fail_on_payload_replacement_not_found =
-      builder.getBoolAttr(failOnPayloadReplacementNotFound);
 
   OpBuilder::InsertionGuard g(builder);
   Region *region = result.addRegion();
   builder.createBlock(region);
   if (bodyBuilder)
     bodyBuilder(builder, result.location);
 }
 
void transform::ApplyCanonicalizationPatternsOp::populatePatterns(
   for (RegisteredOperationName op : ctx->getRegisteredOperations())
     op.getCanonicalizationPatterns(patterns, ctx);
 }
 
 //===----------------------------------------------------------------------===//
 // ApplyRegisteredPassOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure
-transform::ApplyRegisteredPassOp::applyToOne(Operation *target,
-                                             ApplyToEachResultList &results,
-                                             transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    ApplyToEachResultList &results, transform::TransformState &state) {
   // Make sure that this transform is not applied to itself. Modifying the
   // transform IR while it is being interpreted is generally dangerous. Even
   // more so when applying passes because they may perform a wide range of IR
   // modifications.
   DiagnosedSilenceableFailure payloadCheck =
       ensurePayloadIsSeparateFromTransform(*this, target);
   if (!payloadCheck.succeeded())
     return payloadCheck;
DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne(
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
+transform::CastOp::applyToOne(transform::TransformRewriter &rewriter,
+                              Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::CastOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsPayload(effects);
matchBlock(Block &block, Operation *op, transform::TransformState &state,
   // Remember the values mapped to the terminator operands so we can
   // forward them to the action.
   ValueRange yieldedValues = block.getTerminator()->getOperands();
   transform::detail::prepareValueMappings(mappings, yieldedValues, state);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
-transform::ForeachMatchOp::apply(transform::TransformResults &results,
+transform::ForeachMatchOp::apply(transform::TransformRewriter &rewriter,
+                                 transform::TransformResults &results,
                                  transform::TransformState &state) {
   SmallVector<std::pair<FunctionOpInterface, FunctionOpInterface>>
       matchActionPairs;
   matchActionPairs.reserve(getMatchers().size());
   SymbolTableCollection symbolTable;
   for (auto &&[matcher, action] :
        llvm::zip_equal(getMatchers(), getActions())) {
     auto matcherSymbol =
LogicalResult transform::ForeachMatchOp::verifySymbolUses(
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ForeachOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::ForeachOp::apply(transform::TransformResults &results,
+transform::ForeachOp::apply(transform::TransformRewriter &rewriter,
+                            transform::TransformResults &results,
                             transform::TransformState &state) {
   SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});
 
   for (Operation *op : state.getPayloadOps(getTarget())) {
     auto scope = state.make_region_scope(getBody());
     if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
       return DiagnosedSilenceableFailure::definiteFailure();
 
LogicalResult transform::ForeachOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetClosestIsolatedParentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::GetClosestIsolatedParentOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Operation *parent =
         target->getParentWithTrait<OpTrait::IsIsolatedFromAbove>();
     if (!parent) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError()
           << "could not find an isolated-from-above parent op";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     parents.insert(parent);
   }
   results.set(llvm::cast<OpResult>(getResult()), parents.getArrayRef());
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetConsumersOfResult
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::GetConsumersOfResult::apply(transform::TransformResults &results,
+transform::GetConsumersOfResult::apply(transform::TransformRewriter &rewriter,
+                                       transform::TransformResults &results,
                                        transform::TransformState &state) {
   int64_t resultNumber = getResultNumber();
   auto payloadOps = state.getPayloadOps(getTarget());
   if (std::empty(payloadOps)) {
     results.set(cast<OpResult>(getResult()), {});
     return DiagnosedSilenceableFailure::success();
   }
   if (!llvm::hasSingleElement(payloadOps))
     return emitDefiniteFailure()
            << "handle must be mapped to exactly one payload op";
 
   Operation *target = *payloadOps.begin();
   if (target->getNumResults() <= resultNumber)
     return emitDefiniteFailure() << "result number overflow";
   results.set(llvm::cast<OpResult>(getResult()),
               llvm::to_vector(target->getResult(resultNumber).getUsers()));
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetDefiningOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::GetDefiningOp::apply(transform::TransformResults &results,
+transform::GetDefiningOp::apply(transform::TransformRewriter &rewriter,
+                                transform::TransformResults &results,
                                 transform::TransformState &state) {
   SmallVector<Operation *> definingOps;
   for (Value v : state.getPayloadValues(getTarget())) {
     if (llvm::isa<BlockArgument>(v)) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "cannot get defining op of block argument";
       diag.attachNote(v.getLoc()) << "target value";
       return diag;
     }
     definingOps.push_back(v.getDefiningOp());
   }
   results.set(llvm::cast<OpResult>(getResult()), definingOps);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetProducerOfOperand
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::GetProducerOfOperand::apply(transform::TransformResults &results,
+transform::GetProducerOfOperand::apply(transform::TransformRewriter &rewriter,
+                                       transform::TransformResults &results,
                                        transform::TransformState &state) {
   int64_t operandNumber = getOperandNumber();
   SmallVector<Operation *> producers;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Operation *producer =
         target->getNumOperands() <= operandNumber
             ? nullptr
             : target->getOperand(operandNumber).getDefiningOp();
transform::GetProducerOfOperand::apply(transform::TransformRewriter &rewriter,
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetResultOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::GetResultOp::apply(transform::TransformResults &results,
+transform::GetResultOp::apply(transform::TransformRewriter &rewriter,
+                              transform::TransformResults &results,
                               transform::TransformState &state) {
   int64_t resultNumber = getResultNumber();
   SmallVector<Value> opResults;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     if (resultNumber >= target->getNumResults()) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "targeted op does not have enough results";
       diag.attachNote(target->getLoc()) << "target op";
applySequenceBlock(Block &block, transform::FailurePropagationMode mode,
 
   // Forward the operation mapping for values yielded from the sequence to the
   // values produced by the sequence op.
   transform::detail::forwardTerminatorOperands(&block, state, results);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
-transform::IncludeOp::apply(transform::TransformResults &results,
+transform::IncludeOp::apply(transform::TransformRewriter &rewriter,
+                            transform::TransformResults &results,
                             transform::TransformState &state) {
   auto callee = SymbolTable::lookupNearestSymbolFrom<NamedSequenceOp>(
       getOperation(), getTarget());
   assert(callee && "unverified reference to unknown symbol");
 
   if (callee.isExternal())
     return emitDefiniteFailure() << "unresolved external named sequence";
 
DiagnosedSilenceableFailure transform::MatchOperationNameOp::matchOperation(
   return emitSilenceableError() << "wrong operation name";
 }
 
 //===----------------------------------------------------------------------===//
 // MatchParamCmpIOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::MatchParamCmpIOp::apply(transform::TransformResults &results,
+transform::MatchParamCmpIOp::apply(transform::TransformRewriter &rewriter,
+                                   transform::TransformResults &results,
                                    transform::TransformState &state) {
   auto signedAPIntAsString = [&](APInt value) {
     std::string str;
     llvm::raw_string_ostream os(str);
     value.print(os, /*isSigned=*/true);
     return os.str();
   };
 
void transform::MatchParamCmpIOp::getEffects(
   onlyReadsHandle(getReference(), effects);
 }
 
 //===----------------------------------------------------------------------===//
 // ParamConstantOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::ParamConstantOp::apply(transform::TransformResults &results,
+transform::ParamConstantOp::apply(transform::TransformRewriter &rewriter,
+                                  transform::TransformResults &results,
                                   transform::TransformState &state) {
   results.setParams(cast<OpResult>(getParam()), {getValue()});
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // MergeHandlesOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::MergeHandlesOp::apply(transform::TransformResults &results,
+transform::MergeHandlesOp::apply(transform::TransformRewriter &rewriter,
+                                 transform::TransformResults &results,
                                  transform::TransformState &state) {
   SmallVector<Operation *> operations;
   for (Value operand : getHandles())
     llvm::append_range(operations, state.getPayloadOps(operand));
   if (!getDeduplicate()) {
     results.set(llvm::cast<OpResult>(getResult()), operations);
     return DiagnosedSilenceableFailure::success();
   }
OpFoldResult transform::MergeHandlesOp::fold(FoldAdaptor adaptor) {
   return getHandles().front();
 }
 
 //===----------------------------------------------------------------------===//
 // NamedSequenceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::NamedSequenceOp::apply(transform::TransformResults &results,
+transform::NamedSequenceOp::apply(transform::TransformRewriter &rewriter,
+                                  transform::TransformResults &results,
                                   transform::TransformState &state) {
   // Nothing to do here.
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::NamedSequenceOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 

 
 void transform::SplitHandleOp::build(OpBuilder &builder, OperationState &result,
                                      Value target, int64_t numResultHandles) {
   result.addOperands(target);
   result.addTypes(SmallVector<Type>(numResultHandles, target.getType()));
 }
 
 DiagnosedSilenceableFailure
-transform::SplitHandleOp::apply(transform::TransformResults &results,
+transform::SplitHandleOp::apply(transform::TransformRewriter &rewriter,
+                                transform::TransformResults &results,
                                 transform::TransformState &state) {
   int64_t numPayloadOps = llvm::range_size(state.getPayloadOps(getHandle()));
   auto produceNumOpsError = [&]() {
     return emitSilenceableError()
            << getHandle() << " expected to contain " << this->getNumResults()
            << " payload ops but it contains " << numPayloadOps
            << " payload ops";
   };
LogicalResult transform::SplitHandleOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReplicateOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::ReplicateOp::apply(transform::TransformResults &results,
+transform::ReplicateOp::apply(transform::TransformRewriter &rewriter,
+                              transform::TransformResults &results,
                               transform::TransformState &state) {
   unsigned numRepetitions = llvm::range_size(state.getPayloadOps(getPattern()));
   for (const auto &en : llvm::enumerate(getHandles())) {
     Value handle = en.value();
     if (isa<TransformHandleTypeInterface>(handle.getType())) {
       SmallVector<Operation *> current =
           llvm::to_vector(state.getPayloadOps(handle));
       SmallVector<Operation *> payload;
void transform::ReplicateOp::getEffects(
   producesHandle(getReplicated(), effects);
 }
 
 //===----------------------------------------------------------------------===//
 // SequenceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::SequenceOp::apply(transform::TransformResults &results,
+transform::SequenceOp::apply(transform::TransformRewriter &rewriter,
+                             transform::TransformResults &results,
                              transform::TransformState &state) {
   // Map the entry block argument to the list of operations.
   auto scope = state.make_region_scope(*getBodyBlock()->getParent());
   if (failed(mapBlockArguments(state)))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   return applySequenceBlock(*getBodyBlock(), getFailurePropagationMode(), state,
                             results);

 
 void transform::PrintOp::build(OpBuilder &builder, OperationState &result,
                                Value target, StringRef name) {
   result.addOperands({target});
   build(builder, result, name);
 }
 
 DiagnosedSilenceableFailure
-transform::PrintOp::apply(transform::TransformResults &results,
+transform::PrintOp::apply(transform::TransformRewriter &rewriter,
+                          transform::TransformResults &results,
                           transform::TransformState &state) {
   llvm::outs() << "[[[ IR printer: ";
   if (getName().has_value())
     llvm::outs() << *getName() << " ";
 
   if (!getTarget()) {
     llvm::outs() << "top-level ]]]\n" << *state.getTopLevel() << "\n";
     return DiagnosedSilenceableFailure::success();

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

transform::PDLMatchHooks::getPDLConstraintHooks() const {
   return pdlMatchHooks.getConstraintFunctions();
 }
 
 //===----------------------------------------------------------------------===//
 // PDLMatchOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::PDLMatchOp::apply(transform::TransformResults &results,
+transform::PDLMatchOp::apply(transform::TransformRewriter &rewriter,
+                             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))) {
void transform::PDLMatchOp::getEffects(
   onlyReadsPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // WithPDLPatternsOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::WithPDLPatternsOp::apply(transform::TransformResults &results,
+transform::WithPDLPatternsOp::apply(transform::TransformRewriter &rewriter,
+                                    transform::TransformResults &results,
                                     transform::TransformState &state) {
   TransformOpInterface transformOp = nullptr;
   for (Operation &nested : getBody().front()) {
     if (!isa<pdl::PatternOp>(nested)) {
       transformOp = cast<TransformOpInterface>(nested);
       break;
     }
   }

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

 }
 
 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 to %0 {
     transform.apply_patterns.transform.test_patterns
-  } {fail_on_payload_replacement_not_found = false} : !transform.any_op
+  } {transform.silence_tracking_failures} : !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() {

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

public:
   using Op::Op;
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static constexpr llvm::StringLiteral getOperationName() {
     return llvm::StringLiteral("transform.test_transform_op");
   }
 
-  DiagnosedSilenceableFailure apply(transform::TransformResults &results,
+  DiagnosedSilenceableFailure apply(transform::TransformRewriter &rewriter,
+                                    transform::TransformResults &results,
                                     transform::TransformState &state) {
     InFlightDiagnostic remark = emitRemark() << "applying transformation";
     if (Attribute message = getMessage())
       remark << " " << message;
 
     return DiagnosedSilenceableFailure::success();
   }
 
public:
 
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static constexpr llvm::StringLiteral getOperationName() {
     return llvm::StringLiteral(
         "transform.test_transform_unrestricted_op_no_interface");
   }
 
-  DiagnosedSilenceableFailure apply(transform::TransformResults &results,
+  DiagnosedSilenceableFailure apply(transform::TransformRewriter &rewriter,
+                                    transform::TransformResults &results,
                                     transform::TransformState &state) {
     return DiagnosedSilenceableFailure::success();
   }
 
   // No side effects.
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceSelfHandleOrForwardOperandOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   if (getOperation()->getNumOperands() != 0) {
     results.set(cast<OpResult>(getResult()),
                 {getOperation()->getOperand(0).getDefiningOp()});
   } else {
     results.set(cast<OpResult>(getResult()), {getOperation()});
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceSelfHandleOrForwardOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   if (getOperand())
     transform::onlyReadsHandle(getOperand(), effects);
   transform::producesHandle(getRes(), effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceValueHandleToSelfOperand::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   results.setValues(llvm::cast<OpResult>(getOut()), getIn());
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceValueHandleToSelfOperand::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::producesHandle(getOut(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceValueHandleToResult::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   if (target->getNumResults() <= getNumber())
     return emitSilenceableError() << "payload has no result #" << getNumber();
   results.push_back(target->getResult(getNumber()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceValueHandleToResult::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::producesHandle(getOut(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceValueHandleToArgumentOfParentBlock::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   if (!target->getBlock())
     return emitSilenceableError() << "payload has no parent block";
   if (target->getBlock()->getNumArguments() <= getNumber())
     return emitSilenceableError()
            << "parent of the payload has no argument #" << getNumber();
   results.push_back(target->getBlock()->getArgument(getNumber()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceValueHandleToArgumentOfParentBlock::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::producesHandle(getOut(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 bool mlir::test::TestConsumeOperand::allowsRepeatedHandleOperands() {
   return getAllowRepeatedHandles();
 }
 
 DiagnosedSilenceableFailure
-mlir::test::TestConsumeOperand::apply(transform::TransformResults &results,
+mlir::test::TestConsumeOperand::apply(transform::TransformRewriter &rewriter,
+                                      transform::TransformResults &results,
                                       transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestConsumeOperand::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getOperand(), effects);
   if (getSecondOperand())
     transform::consumesHandle(getSecondOperand(), effects);
   transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestConsumeOperandOfOpKindOrFail::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   auto payload = state.getPayloadOps(getOperand());
   assert(llvm::hasSingleElement(payload) && "expected a single target op");
   if ((*payload.begin())->getName().getStringRef() != getOpKind()) {
     return emitSilenceableError()
            << "op expected the operand to be associated a payload op of kind "
            << getOpKind() << " got "
            << (*payload.begin())->getName().getStringRef();

 
 void mlir::test::TestSucceedIfOperandOfOpKind::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestPrintRemarkAtOperandOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   auto payload = state.getPayloadOps(getOperand());
   for (Operation *op : payload)
     op->emitRemark() << getMessage();
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintRemarkAtOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestPrintRemarkAtOperandValue::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   ArrayRef<Value> values = state.getPayloadValues(getIn());
   for (Value value : values) {
     std::string note;
     llvm::raw_string_ostream os(note);
     if (auto arg = llvm::dyn_cast<BlockArgument>(value)) {
       os << "a block argument #" << arg.getArgNumber() << " in block #"
          << std::distance(arg.getOwner()->getParent()->begin(),

 }
 
 void mlir::test::TestPrintRemarkAtOperandValue::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::onlyReadsPayload(effects);
 }
 
-DiagnosedSilenceableFailure
-mlir::test::TestAddTestExtensionOp::apply(transform::TransformResults &results,
-                                          transform::TransformState &state) {
+DiagnosedSilenceableFailure mlir::test::TestAddTestExtensionOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
   state.addExtension<TestTransformStateExtension>(getMessageAttr());
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestCheckIfTestExtensionPresentOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   auto *extension = state.getExtension<TestTransformStateExtension>();
   if (!extension) {
     emitRemark() << "extension absent";
     return DiagnosedSilenceableFailure::success();
   }
 
   InFlightDiagnostic diag = emitRemark()

 
 void mlir::test::TestCheckIfTestExtensionPresentOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemapOperandPayloadToSelfOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   auto *extension = state.getExtension<TestTransformStateExtension>();
   if (!extension)
     return emitDefiniteFailure("TestTransformStateExtension missing");
 
   if (failed(extension->updateMapping(
           *state.getPayloadOps(getOperand()).begin(), getOperation())))
     return DiagnosedSilenceableFailure::definiteFailure();
   if (getNumResults() > 0)
     results.set(cast<OpResult>(getResult(0)), {getOperation()});
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestRemapOperandPayloadToSelfOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::producesHandle(getOut(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemoveTestExtensionOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   state.removeExtension<TestTransformStateExtension>();
   return DiagnosedSilenceableFailure::success();
 }
 
-DiagnosedSilenceableFailure
-mlir::test::TestReversePayloadOpsOp::apply(transform::TransformResults &results,
-                                           transform::TransformState &state) {
+DiagnosedSilenceableFailure mlir::test::TestReversePayloadOpsOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
   auto payloadOps = state.getPayloadOps(getTarget());
   auto reversedOps = llvm::to_vector(llvm::reverse(payloadOps));
   results.set(llvm::cast<OpResult>(getResult()), reversedOps);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestTransformOpWithRegions::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestTransformOpWithRegions::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 
 DiagnosedSilenceableFailure
 mlir::test::TestBranchingTransformOpTerminator::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestBranchingTransformOpTerminator::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 
 DiagnosedSilenceableFailure mlir::test::TestEmitRemarkAndEraseOperandOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   emitRemark() << getRemark();
   for (Operation *op : state.getPayloadOps(getTarget()))
     op->erase();
 
   if (getFailAfterErase())
     return emitSilenceableError() << "silenceable error";
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestEmitRemarkAndEraseOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getTarget(), effects);
   transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestWrongNumberOfResultsOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestWrongNumberOfMultiResultsOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   static int count = 0;
   if (count++ == 0) {
     OperationState opState(target->getLoc(), "foo");
     results.push_back(OpBuilder(target).create(opState));
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestCorrectNumberOfMultiResultsOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(OpBuilder(target).create(opState));
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestMixedNullAndNonNullResultsOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(nullptr);
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
-    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   if (target->hasAttr("target_me"))
     return DiagnosedSilenceableFailure::success();
   return emitDefaultSilenceableFailure(target);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestPrintNumberOfAssociatedPayloadIROps::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   if (!getHandle())
     emitRemark() << 0;
   emitRemark() << llvm::range_size(state.getPayloadOps(getHandle()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintNumberOfAssociatedPayloadIROps::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getHandle(), effects);
 }
 
 DiagnosedSilenceableFailure
-mlir::test::TestCopyPayloadOp::apply(transform::TransformResults &results,
+mlir::test::TestCopyPayloadOp::apply(transform::TransformRewriter &rewriter,
+                                     transform::TransformResults &results,
                                      transform::TransformState &state) {
   results.set(llvm::cast<OpResult>(getCopy()),
               state.getPayloadOps(getHandle()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestCopyPayloadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {

 
 void mlir::test::TestReportNumberOfTrackedHandlesNestedUnder::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getTarget(), effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestReportNumberOfTrackedHandlesNestedUnder::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   int64_t count = 0;
   for (Operation *op : state.getPayloadOps(getTarget())) {
     op->walk([&](Operation *nested) {
       SmallVector<Value> handles;
       (void)state.getHandlesForPayloadOp(nested, handles);
       count += handles.size();
     });
   }
   emitRemark() << count << " handles nested under";
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintParamOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getParam(), effects);
   if (getAnchor())
     transform::onlyReadsHandle(getAnchor(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure
-mlir::test::TestPrintParamOp::apply(transform::TransformResults &results,
+mlir::test::TestPrintParamOp::apply(transform::TransformRewriter &rewriter,
+                                    transform::TransformResults &results,
                                     transform::TransformState &state) {
   std::string str;
   llvm::raw_string_ostream os(str);
   if (getMessage())
     os << *getMessage() << " ";
   llvm::interleaveComma(state.getParams(getParam()), os);
   if (!getAnchor()) {
     emitRemark() << os.str();
     return DiagnosedSilenceableFailure::success();
   }
   for (Operation *payload : state.getPayloadOps(getAnchor()))
     ::mlir::emitRemark(payload->getLoc()) << os.str();
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
-mlir::test::TestAddToParamOp::apply(transform::TransformResults &results,
+mlir::test::TestAddToParamOp::apply(transform::TransformRewriter &rewriter,
+                                    transform::TransformResults &results,
                                     transform::TransformState &state) {
   SmallVector<uint32_t> values(/*Size=*/1, /*Value=*/0);
   if (Value param = getParam()) {
     values = llvm::to_vector(
         llvm::map_range(state.getParams(param), [](Attribute attr) -> uint32_t {
           return llvm::cast<IntegerAttr>(attr).getValue().getLimitedValue(
               UINT32_MAX);
         }));
   }
 
   Builder builder(getContext());
   SmallVector<Attribute> result = llvm::to_vector(
       llvm::map_range(values, [this, &builder](uint32_t value) -> Attribute {
         return builder.getI32IntegerAttr(value + getAddendum());
       }));
   results.setParams(llvm::cast<OpResult>(getResult()), result);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceParamWithNumberOfTestOps::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   Builder builder(getContext());
   SmallVector<Attribute> result = llvm::to_vector(
       llvm::map_range(state.getPayloadOps(getHandle()),
                       [&builder](Operation *payload) -> Attribute {
                         int32_t count = 0;
                         payload->walk([&count](Operation *op) {
                           if (op->getName().getDialectNamespace() == "test")
                             ++count;
                         });
                         return builder.getI32IntegerAttr(count);
                       }));
   results.setParams(llvm::cast<OpResult>(getResult()), result);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceIntegerParamWithTypeOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   Attribute zero = IntegerAttr::get(getType(), 0);
   results.setParams(llvm::cast<OpResult>(getResult()), zero);
   return DiagnosedSilenceableFailure::success();
 }
 
 LogicalResult mlir::test::TestProduceIntegerParamWithTypeOp::verify() {
   if (!llvm::isa<IntegerType>(getType())) {
     return emitOpError() << "expects an integer type";
   }
   return success();
 }
 
 void mlir::test::TestProduceTransformParamOrForwardOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::producesHandle(getOut(), effects);
   transform::producesHandle(getParam(), effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceTransformParamOrForwardOperandOp::applyToOne(
-    Operation *target, ::transform::ApplyToEachResultList &results,
+    transform::TransformRewriter &rewriter, Operation *target,
+    ::transform::ApplyToEachResultList &results,
     ::transform::TransformState &state) {
   Builder builder(getContext());
   if (getFirstResultIsParam()) {
     results.push_back(builder.getI64IntegerAttr(0));
   } else if (getFirstResultIsNull()) {
     results.push_back(nullptr);
   } else {
     results.push_back(*state.getPayloadOps(getIn()).begin());

 }
 
 void mlir::test::TestProduceNullPayloadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::producesHandle(getOut(), effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestProduceNullPayloadOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   SmallVector<Operation *, 1> null({nullptr});
   results.set(llvm::cast<OpResult>(getOut()), null);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestProduceEmptyPayloadOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   results.set(cast<OpResult>(getOut()), {});
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceNullParamOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::producesHandle(getOut(), effects);
 }
 
-DiagnosedSilenceableFailure
-mlir::test::TestProduceNullParamOp::apply(transform::TransformResults &results,
-                                          transform::TransformState &state) {
+DiagnosedSilenceableFailure mlir::test::TestProduceNullParamOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
   results.setParams(llvm::cast<OpResult>(getOut()), Attribute());
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceNullValueOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::producesHandle(getOut(), effects);
 }
 
-DiagnosedSilenceableFailure
-mlir::test::TestProduceNullValueOp::apply(transform::TransformResults &results,
-                                          transform::TransformState &state) {
+DiagnosedSilenceableFailure mlir::test::TestProduceNullValueOp::apply(
+    transform::TransformRewriter &rewriter,
+    transform::TransformResults &results, transform::TransformState &state) {
   results.setValues(llvm::cast<OpResult>(getOut()), Value());
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestRequiredMemoryEffectsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   if (getHasOperandEffect())
     transform::consumesHandle(getIn(), effects);
 
   if (getHasResultEffect())
     transform::producesHandle(getOut(), effects);
   else
     transform::onlyReadsHandle(getOut(), effects);
 
   if (getModifiesPayload())
     transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRequiredMemoryEffectsOp::apply(
+    transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
   results.set(llvm::cast<OpResult>(getOut()), state.getPayloadOps(getIn()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestTrackedRewriteOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getIn(), effects);
   transform::modifiesPayload(effects);
 }
 
 void mlir::test::TestDummyPayloadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   for (OpResult result : getResults())
     transform::producesHandle(result, effects);
 }
 
 DiagnosedSilenceableFailure
-mlir::test::TestTrackedRewriteOp::apply(transform::TransformResults &results,
+mlir::test::TestTrackedRewriteOp::apply(transform::TransformRewriter &rewriter,
+                                        transform::TransformResults &results,
                                         transform::TransformState &state) {
-  transform::ErrorCheckingTrackingListener listener(state, *this);
-  IRRewriter rewriter(getContext(), &listener);
   int64_t numIterations = 0;
 
   // `getPayloadOps` returns an iterator that skips ops that are erased in the
   // loop body. Replacement ops are not enumerated.
   for (Operation *op : state.getPayloadOps(getIn())) {
     ++numIterations;
     (void)op;
 

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

: Op<Transform_Dialect, "test_produce_value_handle_to_result",
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in,
                        I64Attr:$number);
   let results = (outs TransformValueHandleTypeInterface:$out);
   let assemblyFormat = "$in `,` $number attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
     
 def TestProduceValueHandleToArgumentOfParentBlock
   : Op<Transform_Dialect, "test_produce_value_handle_to_argument_of_parent_block",
        [TransformEachOpTrait, TransformOpInterface,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in,
                        I64Attr:$number);
   let results = (outs TransformValueHandleTypeInterface:$out);
   let assemblyFormat = "$in `,` $number attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestConsumeOperand : Op<Transform_Dialect, "test_consume_operand",
      [DeclareOpInterfaceMethods<TransformOpInterface, ["allowsRepeatedHandleOperands"]>,

 def TestConsumeOperandEach : Op<Transform_Dialect, "test_consume_operand_each",
      [TransformOpInterface, TransformEachOpTrait,
       MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]> {
   let arguments = (ins TransformHandleTypeInterface:$target);
   let assemblyFormat = "$target attr-dict `:` type($target)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state) {
       return ::mlir::DiagnosedSilenceableFailure::success();
     }
   }];
 }
 
def TestWrongNumberOfResultsOp
   let results = (outs TransformHandleTypeInterface:$a,
                       TransformHandleTypeInterface:$b,
                       TransformHandleTypeInterface:$c);
   let assemblyFormat =
       "$target attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation * target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestWrongNumberOfMultiResultsOp
   : Op<Transform_Dialect, "test_wrong_number_of_multi_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$result);
   let assemblyFormat =
       "$target attr-dict `:` functional-type($target, $result)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation * target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestCorrectNumberOfMultiResultsOp
   : Op<Transform_Dialect, "test_correct_number_of_multi_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$result1,
                       TransformHandleTypeInterface:$result2);
   let assemblyFormat =
       "$target attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation * target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestMixedNullAndNonNullResultsOp
   : Op<Transform_Dialect, "test_mixed_null_and_non_null_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs TransformHandleTypeInterface:$null,
                       TransformHandleTypeInterface:$non_null);
   let assemblyFormat =
       "$target attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation * target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestMixedSuccessAndSilenceableOp
   : Op<Transform_Dialect, "test_mixed_success_and_silenceable",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs);
   let assemblyFormat = "$target attr-dict `:` type($target)";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation * target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestPrintNumberOfAssociatedPayloadIROps
   : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_ops",
let arguments = (ins TransformHandleTypeInterface:$in,
                        UnitAttr:$second_result_is_handle);
   let results = (outs AnyType:$out,
                       TransformParamTypeInterface:$param);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestProduceNullPayloadOp
   : Op<Transform_Dialect, "test_produce_null_payload",
: Op<Transform_Dialect, "test_dummy_payload_op",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        TransformOpInterface]> {
   let arguments = (ins Variadic<AnyType>:$args);
   let results = (outs Variadic<AnyType>:$outs);
   let assemblyFormat = "$args attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 
   let extraClassDeclaration = [{
-    DiagnosedSilenceableFailure apply(transform::TransformResults &results,
+    DiagnosedSilenceableFailure apply(transform::TransformRewriter &rewriter,
+                                      transform::TransformResults &results,
                                       transform::TransformState &state) {
       llvm_unreachable("op should not be used as a transform");
       return DiagnosedSilenceableFailure::definiteFailure();
     }
   }];
 }
 
 def TestTrackedRewriteOp