[mlir][transform] Use TrackingListener-aware iterator for getPayloadOps

[springerm]

Summary

Instead of returning an ArrayRef<Operation *>, return at iterator that skips ops that were erased/replaced while iterating over the payload ops.

This fixes an issue in conjuction with TrackingListener, where a tracked op was erased during the iteration. Elements may not be removed from an array while iterating over it; this invalidates the iterator.

When ops are erased/removed via replacePayloadOp, they are not immediately removed from the mappings data structure. Instead, they are set to nullptr. nullptrs are not enumerated by getPayloadOps. At the end of each transformation, nullptrs are removed from the mapping data structure.

Depends On: D149937

[springerm]

rebase

[ftynse]

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

592

Nit: please document.

645–663

Nit: let's have Range && here to avoid unnecessary copying.

654

Nit: std::forward<Range>(ops)

665

Isn't this covered by the template case? (Maybe it wasn't because of missing &&?)

1113

Same as above.

1181

Nit: std::empty(targets)

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

213–214

Same as above: Range &&.

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

376–378

Not preserving the order may be problematic. I've seen cases where two handles are co-iterated. We can preserve the order here, I believe. Specifically, when compacting the associations of a handle, we can iterate over the vector backwards and swap nullptr with the last non-null element, which is the replacement value due to the order of addition. We can then just resize the vector down to the size with non-null elements, which is also more efficient than erasing nulls one by one and copying the tail every time,

399

This is problematic. Appending to the vector invalidates iterators because it may trigger reallocation, so if we are iterating over the exact same container elsewhere, it may segfault. We need to either double the size of the vector before the iterator is produced so we can just write to the second half, or store these replacements separately.

[springerm]

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

399

replacement ops are no longer added to the end of the list. they are immediately put at the right place in this function (same functionality as before this revision). the only special handling is for op erasures: nullptrs are put in the direct mapping (and removed later).

[springerm]

address comments

[springerm]

improve test case

[nicolasvasilache]

Great, thanks for pushing on this!

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

static LogicalResult verifyTrait(Operation *op) {
     }
 
     return success();
   }
 
   DiagnosedSilenceableFailure apply(TransformResults &results,
                                     TransformState &state) {
     Value operandHandle = cast<OpTy>(this->getOperation()).getOperandHandle();
-    ArrayRef<Operation *> payload = state.getPayloadOps(operandHandle);
-    if (payload.size() != 1) {
+    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";
     }
 
     return cast<OpTy>(this->getOperation())
-        .matchOperation(payload[0], results, state);
+        .matchOperation(*payload.begin(), results, state);
   }
 
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
     onlyReadsHandle(this->getOperation()->getOperands(), effects);
     producesHandle(this->getOperation()->getResults(), effects);
     onlyReadsPayload(effects);
   }
 };

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

 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 namespace transform {
 
 class TransformOpInterface;
 class TransformResults;
+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
+/// the payload IR, depending on what is available in the context.
+LogicalResult
+mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
+                                             Operation *op, Region &region);
+
+/// Verification hook for PossibleTopLevelTransformOpTrait.
+LogicalResult verifyPossibleTopLevelTransformOpTrait(Operation *op);
+
+/// Verification hook for TransformOpInterface.
+LogicalResult verifyTransformOpInterface(Operation *op);
+
+/// Populates `mappings` with mapped values associated with the given transform
+/// IR values in the given `state`.
+void prepareValueMappings(
+    SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings,
+    ValueRange values, const transform::TransformState &state);
+
+/// Populates `results` with payload associations that match exactly those of
+/// the operands to `block`'s terminator.
+void forwardTerminatorOperands(Block *block, transform::TransformState &state,
+                               transform::TransformResults &results);
+} // namespace detail
 
 /// Options controlling the application of transform operations by the
 /// TransformState.
 class TransformOptions {
 public:
   TransformOptions() = default;
   TransformOptions(const TransformOptions &) = default;
   TransformOptions &operator=(const TransformOptions &) = default;
 
   /// Requests computationally expensive checks of the transform and payload IR
   /// well-formedness to be performed before each transformation. In particular,
   /// these ensure that the handles still point to valid operations when used.
   TransformOptions &enableExpensiveChecks(bool enable = true) {
     expensiveChecksEnabled = enable;
     return *this;
   }
 
   /// Returns true if the expensive checks are requested.
   bool getExpensiveChecksEnabled() const { return expensiveChecksEnabled; }
 
 private:
   bool expensiveChecksEnabled = true;
 };
 
-using Param = Attribute;
-using MappedValue = llvm::PointerUnion<Operation *, Param, Value>;
-
 /// Entry point to the Transform dialect infrastructure. Applies the
 /// transformation specified by `transform` to payload IR contained in
 /// `payloadRoot`. The `transform` operation may contain other operations that
 /// will be executed following the internal logic of the operation. It must
 /// have the `PossibleTopLevelTransformOp` trait and not have any operands.
 /// This function internally keeps track of the transformation state.
 LogicalResult
 applyTransforms(Operation *payloadRoot, TransformOpInterface transform,

 /// end of the `apply` function of the parent operation. If a region contains
 /// blocks with arguments, the client can map those arguments to payload IR ops
 /// using `mapBlockArguments`.
 class TransformState {
 public:
   using Param = transform::Param;
 
 private:
+  friend void
+  detail::forwardTerminatorOperands(Block *block,
+                                    transform::TransformState &state,
+                                    transform::TransformResults &results);
+
   /// Mapping between a Value in the transform IR and the corresponding set of
   /// operations in the payload IR.
   using TransformOpMapping = DenseMap<Value, SmallVector<Operation *, 2>>;
 
   /// Mapping between a payload IR operation and the transform IR values it is
   /// associated with.
   using TransformOpReverseMapping =
       DenseMap<Operation *, SmallVector<Value, 2>>;
public:
   /// Returns the number of extra mappings for the top-level operation.
   size_t getNumTopLevelMappings() const { return topLevelMappedValues.size(); }
 
   /// Returns the position-th extra mapping for the top-level operation.
   ArrayRef<MappedValue> getTopLevelMapping(size_t position) const {
     return topLevelMappedValues[position];
   }
 
-  /// Returns the list of ops that the given transform IR value corresponds to.
-  /// This is helpful for transformations that apply to a particular handle.
-  ArrayRef<Operation *> getPayloadOps(Value value) const;
+  /// Returns an iterator that enumerates all ops that the given transform IR
+  /// value corresponds to at the time when this function is called. Ops may be
+  /// erased while iterating; erased ops are not enumerated. This function is
+  /// helpful for transformations that apply to a particular handle.
+  auto getPayloadOps(Value value) const {
+    // When ops are replaced/erased, they are replaced with nullptr (until
+    // the data structure is compacted). Do not enumerate these ops.
+    return llvm::make_filter_range(getPayloadOpsView(value),
+                                   [](Operation *op) { return op != nullptr; });
+  }
 
   /// Returns the list of parameters that the given transform IR value
   /// corresponds to.
   ArrayRef<Attribute> getParams(Value value) const;
 
   /// Returns the list of payload IR values that the given transform IR value
   /// corresponds to.
   ArrayRef<Value> getPayloadValues(Value handleValue) const;
Mappings &getMapping(Operation *operation) {
     return it->second;
   }
 
   /// Updates the state to include the associations between op results and the
   /// provided result of applying a transform op.
   LogicalResult updateStateFromResults(const TransformResults &results,
                                        ResultRange opResults);
 
+  /// Returns a list of all ops that the given transform IR value corresponds to
+  /// at the time when this function is called. In case an op was erased, the
+  /// returned list contains nullptr. This function is helpful for
+  /// transformations that apply to a particular handle.
+  ArrayRef<Operation *> getPayloadOpsView(Value value) const;
+
   /// Sets the payload IR ops associated with the given transform IR value
   /// (handle). A payload op may be associated multiple handles as long as
   /// at most one of them gets consumed by further transformations.
   /// For example, a hypothetical "find function by name" may be called twice in
   /// a row to produce two handles pointing to the same function:
   ///
   ///   %0 = transform.find_func_by_name { name = "myfunc" }
   ///   %1 = transform.find_func_by_name { name = "myfunc" }
private:
   /// consumed handles.
   LogicalResult
   checkAndRecordHandleInvalidation(TransformOpInterface transform);
 
   /// The mappings between transform IR values and payload IR ops, aggregated by
   /// the region in which the transform IR values are defined.
   llvm::SmallDenseMap<Region *, Mappings> mappings;
 
+  SmallVector<Value> opHandlesToCompact;
+
   /// Extensions attached to the TransformState, identified by the TypeID of
   /// their type. Only one extension of any given type is allowed.
   DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
 
   /// The top-level operation that contains all payload IR, typically a module.
   Operation *topLevel;
 
   /// Extra mapped values (payload operations, values or parameters) to be

 class TransformResults {
   friend class TransformState;
 
 public:
   /// Indicates that the result of the transform IR op at the given position
   /// corresponds to the given list of payload IR ops. Each result must be set
   /// by the transformation exactly once in case of transformation succeeding.
   /// The value must have a type implementing TransformHandleTypeInterface.
-  void set(OpResult value, ArrayRef<Operation *> ops);
+  template <typename Range> void set(OpResult value, Range ops) {
+    int64_t position = value.getResultNumber();
+    assert(position < static_cast<int64_t>(operations.size()) &&
+           "setting results for a non-existent handle");
+    assert(operations[position].data() == nullptr && "results already set");
+    assert(params[position].data() == nullptr &&
+           "another kind of results already set");
+    assert(values[position].data() == nullptr &&
+           "another kind of results already set");
+    operations.replace(position, ops);
+  }
+
+  /// Indicates that the result of the transform IR op at the given position
+  /// corresponds to the given payload IR op. Each result must be set by the
+  /// transformation exactly once in case of transformation succeeding. The
+  /// value must have a type implementing TransformHandleTypeInterface.
+  void set(OpResult value, Operation *op) {
+    set(value, ArrayRef<Operation *>(op));
+  }
 
   /// Indicates that the result of the transform IR op at the given position
   /// corresponds to the given list of parameters. Each result must be set by
   /// the transformation exactly once in case of transformation succeeding. The
   /// value must have a type implementing TransformParamTypeInterface.
   void setParams(OpResult value, ArrayRef<TransformState::Param> params);
 
   /// Indicates that the result of the transform IR op at the given position
   /// corresponds to the given range of payload IR values. Each result must be
   /// set by the transformation exactly once in case of transformation

 /// Creates a RAII object the lifetime of which corresponds to the new mapping
 /// for transform IR values defined in the given isolated-from-above region.
 /// Values defined in surrounding regions cannot be accessed.
 TransformState::RegionScope
 TransformState::make_isolated_region_scope(Region &region) {
   return RegionScope(*this, region, RegionScope::Isolated());
 }
 
-namespace detail {
-/// Maps the only block argument of the op with PossibleTopLevelTransformOpTrait
-/// to either the list of operations associated with its operand or the root of
-/// the payload IR, depending on what is available in the context.
-LogicalResult
-mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
-                                             Operation *op, Region &region);
-
-/// Verification hook for PossibleTopLevelTransformOpTrait.
-LogicalResult verifyPossibleTopLevelTransformOpTrait(Operation *op);
-
-/// Verification hook for TransformOpInterface.
-LogicalResult verifyTransformOpInterface(Operation *op);
-
-/// Populates `mappings` with mapped values associated with the given transform
-/// IR values in the given `state`.
-void prepareValueMappings(
-    SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings,
-    ValueRange values, const transform::TransformState &state);
-
-/// Populates `results` with payload associations that match exactly those of
-/// the operands to `block`'s terminator.
-void forwardTerminatorOperands(Block *block, transform::TransformState &state,
-                               transform::TransformResults &results);
-} // namespace detail
-
 /// This trait is supposed to be attached to Transform dialect operations that
 /// can be standalone top-level transforms. Such operations typically contain
 /// other Transform dialect operations that can be executed following some
 /// control flow logic specific to the current operation. The operations with
 /// this trait are expected to have at least one single-block region with one
 /// argument of PDL Operation type. The operations are also expected to be valid
 /// without operands, in which case they are considered top-level, and with one
 /// or more arguments, in which case they are considered nested. Top-level

 /// callable with the following signature:
 ///   - 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>
+template <typename TransformOpTy, typename Range>
 DiagnosedSilenceableFailure
-applyTransformToEach(TransformOpTy transformOp, ArrayRef<Operation *> targets,
+applyTransformToEach(TransformOpTy transformOp, Range targets,
                      SmallVectorImpl<ApplyToEachResultList> &results,
                      TransformState &state) {
   using OpTy = typename llvm::function_traits<
       decltype(&TransformOpTy::applyToOne)>::template arg_t<0>;
   static_assert(std::is_convertible<OpTy, Operation *>::value,
                 "expected transform function to take an operation");
 
   SmallVector<Diagnostic> silenceableStack;

 } // namespace detail
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
     TransformResults &transformResults, TransformState &state) {
-  ArrayRef<Operation *> targets =
-      state.getPayloadOps(this->getOperation()->getOperand(0));
+  auto targets = state.getPayloadOps(this->getOperation()->getOperand(0));
 
   // Step 1. Handle the corner case where no target is specified.
   // This is typically the case when the matcher fails to apply and we need to
   // propagate gracefully.
   // In this case, we fill all results with an empty vector.
-  if (targets.empty()) {
+  if (std::begin(targets) == std::end(targets)) {
     SmallVector<Operation *> emptyPayload;
     SmallVector<Attribute> emptyParams;
     for (OpResult r : this->getOperation()->getResults()) {
       if (r.getType().isa<TransformParamTypeInterface>())
         transformResults.setParams(r, emptyParams);
       else if (r.getType().isa<TransformValueHandleTypeInterface>())
         transformResults.setValues(r, ValueRange());
       else
         transformResults.set(r, emptyPayload);
     }
     return DiagnosedSilenceableFailure::success();
   }
 
   // Step 2. Call applyToOne on each target and record newly produced ops in its
   // corresponding results entry.
   SmallVector<ApplyToEachResultList, 1> results;
-  results.reserve(targets.size());
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
       cast<OpTy>(this->getOperation()), 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,

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

SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
   // 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())) {
     Value handle = std::get<0>(it);
-    ArrayRef<Operation *> boundedValueOps = state.getPayloadOps(handle);
+    auto boundedValueOps = state.getPayloadOps(handle);
     for (Operation *op : boundedValueOps) {
       if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
         auto diag =
             emitDefiniteFailure()
             << "expected bounded value handle to point to one or multiple "
                "single-result index-typed ops";
         diag.attachNote(op->getLoc()) << "multiple/non-index result";
         return diag;
for (const auto &it : llvm::zip(boundedValues, lbs, ubs)) {
     if (!cstr.findVar(std::get<0>(it), &pos))
       pos = cstr.appendSymbolVar(std::get<0>(it));
     cstr.addBound(presburger::BoundType::LB, pos, std::get<1>(it));
     // Note: addBound bounds are inclusive, but specified UB is exclusive.
     cstr.addBound(presburger::BoundType::UB, pos, std::get<2>(it) - 1);
   }
 
   // Transform all targets.
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
-  for (Operation *target : targets) {
+  SmallVector<Operation *> targets;
+  for (Operation *target : state.getPayloadOps(getTarget())) {
     if (!isa<AffineMinOp, AffineMaxOp>(target)) {
       auto diag = emitDefiniteFailure()
                   << "target must be affine.min or affine.max";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     if (boundedOps.contains(target)) {
       auto diag = emitDefiniteFailure()
                   << "target op result must not be constrainted";
       diag.attachNote(target->getLoc()) << "target/constrained op";
       return diag;
     }
+    targets.push_back(target);
   }
   SmallVector<Operation *> transformed;
   RewritePatternSet patterns(getContext());
   // 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>,

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

transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
   options.bufferizeFunctionBoundaries = getBufferizeFunctionBoundaries();
   options.createDeallocs = getCreateDeallocs();
   options.testAnalysisOnly = getTestAnalysisOnly();
   options.printConflicts = getPrintConflicts();
   if (getFunctionBoundaryTypeConversion().has_value())
     options.setFunctionBoundaryTypeConversion(
         *getFunctionBoundaryTypeConversion());
 
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto payloadOps = state.getPayloadOps(getTarget());
   for (Operation *target : payloadOps) {
     if (!isa<ModuleOp, FunctionOpInterface>(target))
       return emitSilenceableError() << "expected module or function target";
     auto moduleOp = dyn_cast<ModuleOp>(target);
     if (options.bufferizeFunctionBoundaries) {
       if (!moduleOp)
         return emitSilenceableError() << "expected module target";
       if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options)))

 
 DiagnosedSilenceableFailure
 transform::EliminateEmptyTensorsOp::apply(TransformResults &transformResults,
                                           TransformState &state) {
   IRRewriter rewriter(getContext());
   OneShotBufferizationOptions options;
   options.allowReturnAllocs = true;
 
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto payloadOps = state.getPayloadOps(getTarget());
   for (Operation *target : payloadOps) {
     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())

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

static DiagnosedSilenceableFailure unpackSingleIndexResultPDLOperations(
     SmallVector<OpFoldResult> &result, ArrayRef<OpFoldResult> ofrs) {
   for (OpFoldResult ofr : ofrs) {
     if (ofr.is<Attribute>()) {
       if (!ofr.get<Attribute>().isa<IntegerAttr>())
         return transformOp.emitDefiniteFailure() << "expected IntegerAttr";
       result.push_back(ofr);
       continue;
     }
-    ArrayRef<Operation *> payloadOps = state.getPayloadOps(ofr.get<Value>());
-    if (payloadOps.size() != 1) {
+    auto payloadOps = state.getPayloadOps(ofr.get<Value>());
+    if (!llvm::hasSingleElement(payloadOps)) {
       DiagnosedSilenceableFailure diag =
           transformOp.emitSilenceableError()
           << "handle must be mapped to exactly one payload op";
       diag.attachNote(ofr.get<Value>().getLoc())
-          << "mapped to " << payloadOps.size() << " payload ops";
+          << "mapped to " << llvm::range_size(payloadOps) << " payload ops";
       return diag;
     }
 
-    Operation *op = payloadOps[0];
+    Operation *op = *payloadOps.begin();
     if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
       DiagnosedSilenceableFailure diag =
           transformOp.emitSilenceableError()
           << "payload op must have exactly 1 index result";
       diag.attachNote(op->getLoc())
           << "has " << op->getNumResults() << " results";
       return diag;
     }
     result.push_back(op->getResult(0));
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 // Given a list of OpFoldResults that are either index attrs or op
 // handles, return a list of OpFoldResults where all op handles are
 // replaced with the first (and only) OpResult of that payload op. (There
 // must be exactly one mapped payload op and it must have exactly one
 // index result.)
 static DiagnosedSilenceableFailure unpackSingleIndexResultPDLOperations(
     transform::TransformState &state, TransformOpInterface transformOp,
     SmallVector<OpFoldResult> &result, Value packedHandle) {
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(packedHandle);
+  auto payloadOps = state.getPayloadOps(packedHandle);
   for (Operation *op : payloadOps) {
     if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
       DiagnosedSilenceableFailure diag =
           transformOp.emitSilenceableError()
           << "payload op must have exactly 1 index result";
       diag.attachNote(op->getLoc())
           << "has " << op->getNumResults() << " results";
       return diag;
#undef DOWNSCALE
   return emitDefaultSilenceableFailure(target);
 }
 //===----------------------------------------------------------------------===//
 // FuseOp
 //===----------------------------------------------------------------------===//
 
 /// Apply a tiling transformation to all payload ops and store both the
 /// tiled operation as well as the created tile loops.
+template <typename Range>
 static LogicalResult applyTilingToAll(
-    RewriterBase &rewriter, Operation *transformOp,
-    ArrayRef<Operation *> payloadOps, unsigned numLoops,
-    transform::TransformResults &transformResults,
+    RewriterBase &rewriter, Operation *transformOp, Range payloadOps,
+    unsigned numLoops, transform::TransformResults &transformResults,
     function_ref<FailureOr<scf::SCFTileAndFuseResult>(TilingInterface)>
         applyFn) {
   SmallVector<Operation *> tiledLinalgOps;
   SmallVector<SmallVector<Operation *>> loopOps(numLoops);
-  for (unsigned int i = 0; i < numLoops; ++i)
-    loopOps[i].reserve(payloadOps.size());
 
   for (Operation *target : payloadOps) {
     auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
     if (!tilingInterfaceOp)
       return transformOp->emitError("only TilingInterface ops are supported");
 
     rewriter.setInsertionPoint(target);
     FailureOr<scf::SCFTileAndFuseResult> tiledResults =
static Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
 
   return fusedOp;
 }
 
 DiagnosedSilenceableFailure
 transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
                                        transform::TransformState &state) {
   SmallVector<Operation *> fusedOps;
-  ArrayRef<Operation *> producerOps = state.getPayloadOps(getProducerOp());
+  auto producerOps = state.getPayloadOps(getProducerOp());
   // If nothing to fuse, propagate success.
-  if (producerOps.empty()) {
+  if (producerOps.begin() == producerOps.end()) {
     results.set(getFusedOp().cast<OpResult>(),
                 SmallVector<mlir::Operation *>{});
     return DiagnosedSilenceableFailure::success();
   }
-  ArrayRef<Operation *> containingOps = state.getPayloadOps(getContainingOp());
-  if (containingOps.size() != 1) {
+  auto containingOps = state.getPayloadOps(getContainingOp());
+  if (!llvm::hasSingleElement(containingOps)) {
     return emitDefiniteFailure()
            << "requires exactly one containing_op handle (got "
-           << containingOps.size() << ")";
+           << llvm::range_size(containingOps) << ")";
   }
-  Operation *containingOp = containingOps.front();
+  Operation *containingOp = *containingOps.begin();
 
   // Helper function to find the next producer that should be fused. Take any
   // producer that has a use inside the containing op.
   SmallVector<Operation *> remainingProducers(producerOps.begin(),
                                               producerOps.end());
   auto getNextProducer = [&]() -> FailureOr<Operation *> {
     for (const auto &it : enumerate(remainingProducers)) {
       Operation *producerOp = it.value();

 DiagnosedSilenceableFailure
 transform::MatchOp::apply(transform::TransformResults &results,
                           transform::TransformState &state) {
   llvm::StringSet<> strs;
   if (getOps().has_value())
     strs.insert(getOps()->getAsValueRange<StringAttr>().begin(),
                 getOps()->getAsValueRange<StringAttr>().end());
 
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
-  if (payloadOps.size() != 1) {
+  auto payloadOps = state.getPayloadOps(getTarget());
+  if (!llvm::hasSingleElement(payloadOps)) {
     return emitDefiniteFailure("requires exactly one target handle");
   }
 
   SmallVector<Operation *> res;
   auto matchFun = [&](Operation *op) {
     if (getOps().has_value() && !strs.contains(op->getName().getStringRef()))
       return;
 
if (getFilterResultType().has_value()) {
         return;
     }
 
     // All constraints are satisfied.
     res.push_back(op);
     return;
   };
 
-  payloadOps.front()->walk(matchFun);
+  (*payloadOps.begin())->walk(matchFun);
   results.set(getResult().cast<OpResult>(), res);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // MultiTileSizesOp
 //===---------------------------------------------------------------------===//
 

 SmallVector<OpFoldResult> transform::PackOp::getMixedPackedSizes() {
   Builder b(getContext());
   return getMixedValues(getStaticPackedSizes(), getPackedSizes(), b);
 }
 
 DiagnosedSilenceableFailure
 transform::PackOp::apply(transform::TransformResults &transformResults,
                          transform::TransformState &state) {
-  ArrayRef<Operation *> targetOps = state.getPayloadOps(getTarget());
+  auto targetOps = state.getPayloadOps(getTarget());
   // If nothing to pack, propagate success.
-  if (targetOps.empty()) {
-    transformResults.set(getPackedOp().cast<OpResult>(), {});
+  if (targetOps.begin() == targetOps.end()) {
+    transformResults.set(getPackedOp().cast<OpResult>(),
+                         ArrayRef<Operation *>({}));
     return DiagnosedSilenceableFailure::success();
   }
   // Fail on multi-op handles.
-  auto linalgOp = dyn_cast<LinalgOp>(targetOps.front());
-  if (targetOps.size() != 1 || !linalgOp) {
+  auto linalgOp = dyn_cast<LinalgOp>(*targetOps.begin());
+  if (!llvm::hasSingleElement(targetOps) || !linalgOp) {
     return emitSilenceableError()
            << "requires target to map to exactly 1 LinalgOp (got "
-           << targetOps.size() << ")";
+           << llvm::range_size(targetOps) << ")";
   }
   // Fail on mismatched number of pack sizes.
   if (getMixedPackedSizes().size() != linalgOp.getNumLoops()) {
     return emitSilenceableError()
            << "requires number of packed sizes match the number of loops ("
            << getMixedPackedSizes().size() << " vs " << linalgOp.getNumLoops()
            << ")";
   }
packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
   // `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,
                       transform::TransformState &state) {
-  ArrayRef<Operation *> targetOpsView = state.getPayloadOps(getTarget());
+  auto targetOpsView = state.getPayloadOps(getTarget());
   // Store payload ops into a separate SmallVector because the TrackingListener
   // removes erased ops from the transform state.
   SmallVector<Operation *> targetOps(targetOpsView.begin(),
                                      targetOpsView.end());
 
   SmallVector<Operation *> results;
   TrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
bool isValidPackingPermutation(
   }
   return permutation.size() == op.getDestRank() &&
          isPermutationVector(permutation);
 }
 
 DiagnosedSilenceableFailure
 transform::PackTransposeOp::apply(transform::TransformResults &transformResults,
                                   transform::TransformState &state) {
-  ArrayRef<Operation *> packOrUnpackOps =
-      state.getPayloadOps(getTargetPackOrUnPackOp());
-  ArrayRef<Operation *> linalgOps = state.getPayloadOps(getTargetLinalgOp());
+  auto packOrUnpackOps = state.getPayloadOps(getTargetPackOrUnPackOp());
+  auto linalgOps = state.getPayloadOps(getTargetLinalgOp());
   // Step 1. If nothing to pack, propagate success.
-  if (packOrUnpackOps.empty()) {
-    transformResults.set(getPackedOp().cast<OpResult>(), {});
-    transformResults.set(getPackOp().cast<OpResult>(), {});
-    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+  if (packOrUnpackOps.begin() == packOrUnpackOps.end()) {
+    transformResults.set(getPackedOp().cast<OpResult>(),
+                         ArrayRef<Operation *>({}));
+    transformResults.set(getPackOp().cast<OpResult>(),
+                         ArrayRef<Operation *>({}));
+    transformResults.set(getUnPackOp().cast<OpResult>(),
+                         ArrayRef<Operation *>({}));
     return DiagnosedSilenceableFailure::success();
   }
 
   // Step 2. Bunch of runtime sanity check and error messages.
   // Step 2.1. Fail on multi-op handles.
-  if (packOrUnpackOps.size() != 1 || linalgOps.size() != 1) {
-    return emitSilenceableError() << "requires target to map to exactly 1 "
-                                     "packing op and 1 packed op ("
-                                  << "got " << packOrUnpackOps.size() << " and "
-                                  << linalgOps.size() << ")";
+  if (!llvm::hasSingleElement(packOrUnpackOps) ||
+      !llvm::hasSingleElement(linalgOps)) {
+    return emitSilenceableError()
+           << "requires target to map to exactly 1 "
+           << "got " << llvm::range_size(packOrUnpackOps) << " and "
+           << llvm::range_size(linalgOps) << ")";
   }
 
   // Step 2.2. Fail on wrong type.
-  auto packOp = dyn_cast<tensor::PackOp>(packOrUnpackOps.front());
-  auto unPackOp = dyn_cast<tensor::UnPackOp>(packOrUnpackOps.front());
+  auto packOp = dyn_cast<tensor::PackOp>(*packOrUnpackOps.begin());
+  auto unPackOp = dyn_cast<tensor::UnPackOp>(*packOrUnpackOps.begin());
   if ((!packOp && !unPackOp)) {
     return emitSilenceableError() << "requires target to map to a "
                                      "tensor.pack or tensor.unpack";
   }
-  LinalgOp linalgOpTarget = dyn_cast<LinalgOp>(linalgOps.front());
+  LinalgOp linalgOpTarget = dyn_cast<LinalgOp>(*linalgOps.begin());
   if (!linalgOpTarget)
     return emitSilenceableError() << "requires a LinalgOp target";
 
   // Step 2.3. Fail if we can't get the producer / consumer Linalg op.
   LinalgOp linalgOp;
   if (packOp && packOp.getResult().hasOneUse())
     linalgOp = dyn_cast<LinalgOp>(*(packOp.getResult().getUsers().begin()));
   else if (unPackOp)
transform::PackTransposeOp::apply(transform::TransformResults &transformResults,
   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(getPackOp().cast<OpResult>(), {res->transposedPackOp});
+  transformResults.set(getPackOp().cast<OpResult>(),
+                       ArrayRef<Operation *>({res->transposedPackOp}));
   transformResults.set(getPackedOp().cast<OpResult>(),
-                       {res->transposedLinalgOp});
+                       ArrayRef<Operation *>({res->transposedLinalgOp}));
   if (unPackOp) {
     transformResults.set(getUnPackOp().cast<OpResult>(),
-                         {res->transposedUnPackOp});
+                         ArrayRef<Operation *>({res->transposedUnPackOp}));
   } else {
-    transformResults.set(getUnPackOp().cast<OpResult>(), {});
+    transformResults.set(getUnPackOp().cast<OpResult>(),
+                         ArrayRef<Operation *>({}));
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//

 
 //===---------------------------------------------------------------------===//
 // HoistPadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::HoistPadBuildPackingLoopNestOp::apply(
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
-  ArrayRef<Operation *> targetOps = state.getPayloadOps(getTarget());
-  ArrayRef<Operation *> loopOps = state.getPayloadOps(getLoop());
-  if (targetOps.size() != 1 || loopOps.size() != 1) {
+  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 "
-           << targetOps.size() << " and " << loopOps.size() << ")";
+           << llvm::range_size(targetOps) << " and "
+           << llvm::range_size(loopOps) << ")";
   }
 
-  auto padOp = dyn_cast_or_null<tensor::PadOp>(targetOps.front());
-  auto loopOp = dyn_cast_or_null<scf::ForOp>(loopOps.front());
+  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());

 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ReplaceOp::apply(TransformResults &transformResults,
                             TransformState &state) {
-  ArrayRef<Operation *> payload = state.getPayloadOps(getTarget());
+  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()

 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
                                            TransformState &state) {
   // Collect the dynamic split points if provided.
-  ArrayRef<Operation *> payload = state.getPayloadOps(getTarget());
+  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 (getDynamicSplitPoint().getType().isa<TransformHandleTypeInterface>()) {
       splitPoints = llvm::to_vector(llvm::map_range(
build(builder, result, /*tiled_linalg_op=*/target.getType(),
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange));
 }
 
 DiagnosedSilenceableFailure
 transform::TileOp::apply(TransformResults &transformResults,
                          TransformState &state) {
   ArrayRef<int64_t> tileSizes = getStaticSizes();
 
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
-  SmallVector<ArrayRef<Operation *>> dynamicSizeProducers;
+  SmallVector<Operation *> targets =
+      llvm::to_vector(state.getPayloadOps(getTarget()));
+  SmallVector<SmallVector<Operation *>> dynamicSizeProducers;
   SmallVector<SmallVector<int64_t>> paramSizes;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
   paramSizes.reserve(getDynamicSizes().size());
   for (Value transformValue : getDynamicSizes()) {
     if (transformValue.getType().isa<ParamType>()) {
       dynamicSizeProducers.push_back({});
       ArrayRef<Attribute> params = state.getParams(transformValue);
       paramSizes.push_back(
if (transformValue.getType().isa<ParamType>()) {
             << targets.size() << ")";
         diag.attachNote(transformValue.getLoc()) << "for this parameter";
         return diag;
       }
 
       continue;
     }
     paramSizes.push_back({});
-    dynamicSizeProducers.push_back(state.getPayloadOps(transformValue));
+    dynamicSizeProducers.push_back(
+        llvm::to_vector(state.getPayloadOps(transformValue)));
 
     if (dynamicSizeProducers.back().size() != targets.size()) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError()
           << "expected as many dynamic size-producing operations ("
           << dynamicSizeProducers.back().size() << ") as target ops ("
           << targets.size() << ")";
       diag.attachNote(transformValue.getLoc()) << "for this handle";

 }
 
 DiagnosedSilenceableFailure
 transform::TileToForallOp::apply(transform::TransformResults &transformResults,
                                  transform::TransformState &state) {
   TrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
   auto transformOp = cast<TransformOpInterface>(getOperation());
-  ArrayRef<Operation *> targetsView = state.getPayloadOps(getTarget());
+  auto targetsView = state.getPayloadOps(getTarget());
   // Store payload ops into a separate SmallVector because the TrackingListener
   // removes erased ops from the transform state.
   SmallVector<Operation *> targets(targetsView.begin(), targetsView.end());
 
   // Result payload ops.
   SmallVector<Operation *> tileOps;
   SmallVector<Operation *> tiledOps;
 
transform::TileToForallOp::apply(transform::TransformResults &transformResults,
   status = getPackedTileSizes()
                ? unpackSingleIndexResultPDLOperations(
                      state, transformOp, mixedTileSizes, getPackedTileSizes())
                : unpackSingleIndexResultPDLOperations(
                      state, transformOp, mixedTileSizes, getMixedTileSizes());
   if (!status.succeeded())
     return status;
 
+  // TODO: Put loop into this function.
   DiagnosedSilenceableFailure diag =
       tileToForallOpImpl(rewriter, state, transformOp, targets, mixedNumThreads,
                          mixedTileSizes, getMapping(), tileOps, tiledOps);
 
   if (!diag.succeeded())
     return diag;
 
   transformResults.set(getForallOp().cast<OpResult>(), tileOps);
build(builder, result,
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange));
 }
 
 DiagnosedSilenceableFailure
 transform::TileToScfForOp::apply(TransformResults &transformResults,
                                  TransformState &state) {
   ArrayRef<int64_t> tileSizes = getStaticSizes();
 
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
-  SmallVector<ArrayRef<Operation *>> dynamicSizeProducers;
+  SmallVector<Operation *> targets =
+      llvm::to_vector(state.getPayloadOps(getTarget()));
+  SmallVector<SmallVector<Operation *>> dynamicSizeProducers;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
   for (Value dynamicSizeProducerHandle : getDynamicSizes()) {
     dynamicSizeProducers.push_back(
-        state.getPayloadOps(dynamicSizeProducerHandle));
+        llvm::to_vector(state.getPayloadOps(dynamicSizeProducerHandle)));
 
     if (dynamicSizeProducers.back().size() != targets.size()) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError()
           << "expected as many dynamic size-producing operations ("
           << dynamicSizeProducers.back().size() << ") as target ops ("
           << targets.size() << ")";
       diag.attachNote(dynamicSizeProducerHandle.getLoc()) << "for this handle";

 // MaskedVectorizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
     mlir::transform::TransformResults &transformResults,
     mlir::transform::TransformState &state) {
   TrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
-  if (targets.empty())
+  auto targets = state.getPayloadOps(getTarget());
+  if (targets.begin() == targets.end())
     return DiagnosedSilenceableFailure::success();
 
   SmallVector<int64_t> vectorSizes;
   for (OpFoldResult sz : getMixedVectorSizes()) {
     if (sz.is<Attribute>()) {
       auto attr = sz.get<Attribute>();
       vectorSizes.push_back(attr.cast<IntegerAttr>().getInt());
       continue;
     }
 
-    ArrayRef<Operation *> szPayloads = state.getPayloadOps(sz.get<Value>());
-    if (szPayloads.size() != 1) {
+    auto szPayloads = state.getPayloadOps(sz.get<Value>());
+    if (!llvm::hasSingleElement(szPayloads)) {
       auto diag = this->emitOpError(
           "requires vector size handle that is mapped to 1 payload op");
       diag.attachNote(sz.get<Value>().getLoc())
-          << "mapped to " << szPayloads.size() << " payload ops";
+          << "mapped to " << llvm::range_size(szPayloads) << " payload ops";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
 
-    Operation *szPayloadOp = szPayloads[0];
+    Operation *szPayloadOp = *szPayloads.begin();
     if (szPayloadOp->getNumResults() != 1 ||
         !szPayloadOp->getResult(0).getType().isIndex()) {
       auto diag = this->emitOpError(
           "requires vector size payload op with 1 index result");
       diag.attachNote(szPayloadOp->getLoc()) << "vector size payload op";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
 

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

 //===----------------------------------------------------------------------===//
 // MemRefMultiBufferOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::apply(
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
   SmallVector<Operation *> results;
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto payloadOps = state.getPayloadOps(getTarget());
   IRRewriter rewriter(getContext());
   for (auto *op : payloadOps) {
     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))

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

 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   regionStack.push_back(region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 }
 
 Operation *transform::TransformState::getTopLevel() const { return topLevel; }
 
 ArrayRef<Operation *>
-transform::TransformState::getPayloadOps(Value value) const {
+transform::TransformState::getPayloadOpsView(Value value) const {
   const TransformOpMapping &operationMapping = getMapping(value).direct;
   auto iter = operationMapping.find(value);
   assert(
       iter != operationMapping.end() &&
       "cannot find mapping for payload handle (param/value handle provided?)");
   return iter->getSecond();
 }
 
if (options.getExpensiveChecksEnabled()) {
     assert(it != cachedNames.end() && "entry not found");
     assert(it->second == op->getName() && "operation name mismatch");
     cachedNames.erase(it);
   }
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
   // TODO: consider invalidating the handles to nested objects here.
 
-  // If replacing with null, that is erasing the mapping, drop the mapping
-  // between the handles and the IR objects and return.
-  if (!replacement) {
-    for (Value handle : opHandles) {
-      Mappings &mappings = getMapping(handle);
-      dropMappingEntry(mappings.direct, handle, op);
-    }
-    return success();
-  }
-
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
-  if (options.getExpensiveChecksEnabled()) {
+  if (replacement && options.getExpensiveChecksEnabled()) {
     auto insertion = cachedNames.insert({replacement, replacement->getName()});
     if (!insertion.second) {
       assert(insertion.first->second == replacement->getName() &&
              "operation is already cached with a different name");
     }
   }
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
-  // Otherwise, replace the pointed-to object of all handles while preserving
-  // their relative order. First, replace the mapped operation if present.
+  // Replace the pointed-to object of all handles with nullptr. This ensures
+  // payload op iterators stay in a consistent state. (Removing an op from the
+  // mapping would be problematic because removing an element from an array
+  // invalidates iterators; merely changing the value of elements does not.)
+  // nullptrs are not enumerated by `getPayloadOps`. They are removed at the end
+  // of each transform op application.
+  //
+  // Replacement ops are appended to the end of the list. Note that, therefore,
+  // the relative order of ops is not preserved. I.e., the replacement op is
+  // always at the end of the list of payload ops.
   for (Value handle : opHandles) {
     Mappings &mappings = getMapping(handle);
     auto it = mappings.direct.find(handle);
     if (it == mappings.direct.end())
       continue;
 
     SmallVector<Operation *, 2> &association = it->getSecond();
     // Note that an operation may be associated with the handle more than once.
+    int64_t counter = 0;
     for (Operation *&mapped : association) {
-      if (mapped == op)
-        mapped = replacement;
+      if (mapped == op) {
+        mapped = nullptr;
+        ++counter;
+      }
     }
+    assert(counter > 0 && "inconsistent mapping state");
+
+    if (replacement) {
+      // Add replacement op to the end of the list.
+      for (int64_t i = 0; i < counter; ++i)
+        association.push_back(replacement);
       mappings.reverse[replacement].push_back(handle);
     }
 
+    // nullptr was added to the mapping data structure. For efficiency reasons,
+    // compact the list of payload ops at the end of the transformation by
+    // removing all nullptrs.
+    opHandlesToCompact.push_back(handle);
+  }
+
   return success();
 }
 
 LogicalResult
 transform::TransformState::replacePayloadValue(Value value, Value replacement) {
   SmallVector<Value> valueHandles;
   (void)getHandlesForPayloadValue(value, valueHandles);
 
for (OpOperand &target : transform->getOpOperands()) {
     auto consumesTarget = [&](const MemoryEffects::EffectInstance &effect) {
       return isa<MemoryEffects::Free>(effect.getEffect()) &&
              effect.getValue() == target.get();
     };
     if (llvm::any_of(effects, consumesTarget)) {
       FULL_LDBG("----found consume effect -> SKIP\n");
       if (target.get().getType().isa<TransformHandleTypeInterface>()) {
         FULL_LDBG("----recordOpHandleInvalidation\n");
-        ArrayRef<Operation *> payloadOps = getPayloadOps(target.get());
+        ArrayRef<Operation *> payloadOps = getPayloadOpsView(target.get());
         recordOpHandleInvalidation(target, payloadOps);
       } else if (target.get()
                      .getType()
                      .isa<TransformValueHandleTypeInterface>()) {
         FULL_LDBG("----recordValueHandleInvalidation\n");
         recordValueHandleInvalidation(target);
       } else {
         FULL_LDBG("----not a TransformHandle -> SKIP AND DROP ON THE FLOOR\n");
transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG({
     DBGS() << "applying: ";
     transform->print(llvm::dbgs(), OpPrintingFlags().skipRegions());
     llvm::dbgs() << "\n";
   });
   DEBUG_WITH_TYPE(DEBUG_TYPE_FULL,
                   DBGS() << "Top-level payload before application:\n"
                          << *getTopLevel() << "\n");
+  auto compactOpHandles = llvm::make_scope_exit([this] {
+    // Remove all nullptr ops from handles that had replacements/erasures.
+    for (Value handle : opHandlesToCompact) {
+      Mappings &mappings = getMapping(handle);
+      dropMappingEntry(mappings.direct, handle, nullptr);
+    }
+    opHandlesToCompact.clear();
+  });
   auto printOnFailureRAII = llvm::make_scope_exit([this] {
     (void)this;
     LLVM_DEBUG(DBGS() << "Failing Top-level payload:\n"; getTopLevel()->print(
         llvm::dbgs(), mlir::OpPrintingFlags().printGenericOpForm()););
   });
   if (options.getExpensiveChecksEnabled()) {
     FULL_LDBG("ExpensiveChecksEnabled\n");
     if (failed(checkAndRecordHandleInvalidation(transform)))
for (OpOperand &operand : transform->getOpOperands()) {
       }
       FULL_LDBG("--handle is consumed\n");
 
       Type operandType = operand.get().getType();
       if (operandType.isa<TransformHandleTypeInterface>()) {
         FULL_LDBG("--checkRepeatedConsumptionInOperand for Operation*\n");
         DiagnosedSilenceableFailure check =
             checkRepeatedConsumptionInOperand<Operation *>(
-                getPayloadOps(operand.get()), transform,
+                getPayloadOpsView(operand.get()), transform,
                 operand.getOperandNumber());
         if (!check.succeeded()) {
           FULL_LDBG("----FAILED\n");
           return check;
         }
       } else if (operandType.isa<TransformValueHandleTypeInterface>()) {
         FULL_LDBG("--checkRepeatedConsumptionInOperand For Value\n");
         DiagnosedSilenceableFailure check =

 //===----------------------------------------------------------------------===//
 
 transform::TransformResults::TransformResults(unsigned numSegments) {
   operations.appendEmptyRows(numSegments);
   params.appendEmptyRows(numSegments);
   values.appendEmptyRows(numSegments);
 }
 
-void transform::TransformResults::set(OpResult value,
-                                      ArrayRef<Operation *> ops) {
-  int64_t position = value.getResultNumber();
-  assert(position < static_cast<int64_t>(operations.size()) &&
-         "setting results for a non-existent handle");
-  assert(operations[position].data() == nullptr && "results already set");
-  assert(params[position].data() == nullptr &&
-         "another kind of results already set");
-  assert(values[position].data() == nullptr &&
-         "another kind of results already set");
-  operations.replace(position, ops);
-}
-
 void transform::TransformResults::setParams(
     OpResult value, ArrayRef<transform::TransformState::Param> params) {
   int64_t position = value.getResultNumber();
   assert(position < static_cast<int64_t>(this->params.size()) &&
          "setting params for a non-existent handle");
   assert(this->params[position].data() == nullptr && "params already set");
   assert(operations[position].data() == nullptr &&
          "another kind of results already set");

 
 void transform::detail::forwardTerminatorOperands(
     Block *block, transform::TransformState &state,
     transform::TransformResults &results) {
   for (auto &&[terminatorOperand, result] :
        llvm::zip(block->getTerminator()->getOperands(),
                  block->getParentOp()->getOpResults())) {
     if (result.getType().isa<transform::TransformHandleTypeInterface>()) {
-      results.set(result, state.getPayloadOps(terminatorOperand));
+      results.set(result, state.getPayloadOpsView(terminatorOperand));
     } else if (result.getType()
                    .isa<transform::TransformValueHandleTypeInterface>()) {
       results.setValues(result, state.getPayloadValues(terminatorOperand));
     } else {
       assert(result.getType().isa<transform::TransformParamTypeInterface>() &&
              "unhandled transform type interface");
       results.setParams(result, state.getParams(terminatorOperand));
     }

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

void transform::AlternativesOp::getRegionInvocationBounds(
   bounds.reserve(getNumRegions());
   bounds.emplace_back(1, 1);
   bounds.resize(getNumRegions(), InvocationBounds(0, 1));
 }
 
 static void forwardEmptyOperands(Block *block, transform::TransformState &state,
                                  transform::TransformResults &results) {
   for (const auto &res : block->getParentOp()->getOpResults())
-    results.set(res, {});
+    results.set(res, ArrayRef<Operation *>({}));
 }
 
 DiagnosedSilenceableFailure
 transform::AlternativesOp::apply(transform::TransformResults &results,
                                  transform::TransformState &state) {
   SmallVector<Operation *> originals;
   if (Value scopeHandle = getScope())
     llvm::append_range(originals, state.getPayloadOps(scopeHandle));

 
 //===----------------------------------------------------------------------===//
 // ForeachOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ForeachOp::apply(transform::TransformResults &results,
                             transform::TransformState &state) {
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto payloadOps = state.getPayloadOps(getTarget());
   SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});
 
   for (Operation *op : payloadOps) {
     auto scope = state.make_region_scope(getBody());
     if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     // Execute loop body.
     for (Operation &transform : getBody().front().without_terminator()) {
       DiagnosedSilenceableFailure result = state.applyTransform(
           cast<transform::TransformOpInterface>(transform));
       if (!result.succeeded())
         return result;
     }
 
     // Append yielded payload ops to result list (if any).
     for (unsigned i = 0; i < getNumResults(); ++i) {
-      ArrayRef<Operation *> yieldedOps =
-          state.getPayloadOps(getYieldOp().getOperand(i));
+      auto yieldedOps = state.getPayloadOps(getYieldOp().getOperand(i));
       resultOps[i].append(yieldedOps.begin(), yieldedOps.end());
     }
   }
 
   for (unsigned i = 0; i < getNumResults(); ++i)
     results.set(getResult(i).cast<OpResult>(), resultOps[i]);
 
   return DiagnosedSilenceableFailure::success();

 //===----------------------------------------------------------------------===//
 // GetConsumersOfResult
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::GetConsumersOfResult::apply(transform::TransformResults &results,
                                        transform::TransformState &state) {
   int64_t resultNumber = getResultNumber();
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
-  if (payloadOps.empty()) {
-    results.set(getResult().cast<OpResult>(), {});
+  auto payloadOps = state.getPayloadOps(getTarget());
+  if (payloadOps.begin() == payloadOps.end()) {
+    results.set(getResult().cast<OpResult>(), ArrayRef<Operation *>({}));
     return DiagnosedSilenceableFailure::success();
   }
-  if (payloadOps.size() != 1)
+  if (!llvm::hasSingleElement(payloadOps))
     return emitDefiniteFailure()
            << "handle must be mapped to exactly one payload op";
 
-  Operation *target = payloadOps.front();
+  Operation *target = *payloadOps.begin();
   if (target->getNumResults() <= resultNumber)
     return emitDefiniteFailure() << "result number overflow";
   results.set(getResult().cast<OpResult>(),
               llvm::to_vector(target->getResult(resultNumber).getUsers()));
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
result.addAttribute(SplitHandlesOp::getNumResultHandlesAttrName(result.name),
                       builder.getI64IntegerAttr(numResultHandles));
   auto pdlOpType = pdl::OperationType::get(builder.getContext());
   result.addTypes(SmallVector<pdl::OperationType>(numResultHandles, pdlOpType));
 }
 
 DiagnosedSilenceableFailure
 transform::SplitHandlesOp::apply(transform::TransformResults &results,
                                  transform::TransformState &state) {
-  int64_t numResultHandles =
-      getHandle() ? state.getPayloadOps(getHandle()).size() : 0;
+  SmallVector<Operation *> payloadOps =
+      llvm::to_vector(state.getPayloadOps(getHandle()));
+  int64_t numResultHandles = getHandle() ? payloadOps.size() : 0;
   int64_t expectedNumResultHandles = getNumResultHandles();
   if (numResultHandles != expectedNumResultHandles) {
     // Empty input handle corner case: always propagates empty handles in both
     // suppress and propagate modes.
     if (numResultHandles == 0)
       return DiagnosedSilenceableFailure::success();
     // If the input handle was not empty and the number of result handles does
     // not match, this is a legit silenceable error.
     return emitSilenceableError()
            << getHandle() << " expected to contain " << expectedNumResultHandles
            << " operation handles but it only contains " << numResultHandles
            << " handles";
   }
   // Normal successful case.
-  for (const auto &en : llvm::enumerate(state.getPayloadOps(getHandle())))
+  for (const auto &en : llvm::enumerate(payloadOps))
     results.set(getResults()[en.index()].cast<OpResult>(), en.value());
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::SplitHandlesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getHandle(), effects);
   producesHandle(getResults(), effects);

 
 //===----------------------------------------------------------------------===//
 // ReplicateOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ReplicateOp::apply(transform::TransformResults &results,
                               transform::TransformState &state) {
-  unsigned numRepetitions = state.getPayloadOps(getPattern()).size();
+  unsigned numRepetitions = llvm::range_size(state.getPayloadOps(getPattern()));
   for (const auto &en : llvm::enumerate(getHandles())) {
     Value handle = en.value();
     if (handle.getType().isa<TransformHandleTypeInterface>()) {
-      ArrayRef<Operation *> current = state.getPayloadOps(handle);
+      SmallVector<Operation *> current =
+          llvm::to_vector(state.getPayloadOps(handle));
       SmallVector<Operation *> payload;
       payload.reserve(numRepetitions * current.size());
       for (unsigned i = 0; i < numRepetitions; ++i)
         llvm::append_range(payload, current);
       results.set(getReplicated()[en.index()].cast<OpResult>(), payload);
     } else {
       assert(handle.getType().isa<TransformParamTypeInterface>() &&
              "expected param type");
if (getName().has_value())
     llvm::outs() << *getName() << " ";
 
   if (!getTarget()) {
     llvm::outs() << "top-level ]]]\n" << *state.getTopLevel() << "\n";
     return DiagnosedSilenceableFailure::success();
   }
 
   llvm::outs() << "]]]\n";
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  auto targets = state.getPayloadOps(getTarget());
   for (Operation *target : targets)
     llvm::outs() << *target << "\n";
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::PrintOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {

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

-// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics | FileCheck %s
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-remark @below {{applying transformation}}
   transform.test_transform_op
 }
 
 // -----
module attributes { transform.with_named_sequence } {
   // expected-remark @below {{matched > -1}}
   // expected-remark @below {{matched >= 1}}
   // expected-remark @below {{matched <= 1}}
   func.func @definition() {
     "test.something"() : () -> ()
     return
   }
 }
+
+// -----
+
+// CHECK-LABEL: func @test_tracked_rewrite() {
+//  CHECK-NEXT:   "test.foo"() {original_op = "test.bar"}
+//  CHECK-NEXT:   "test.foo"() {original_op = "test.bar"}
+//  CHECK-NEXT:   "test.foo"() {original_op = "test.bar"}
+//  CHECK-NEXT: }
+func.func @test_tracked_rewrite() {
+  %0 = "test.bar"() : () -> (i1)
+  %1 = "test.bar"() : () -> (i1)
+  %2 = "test.bar"() : () -> (i1)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %bar = transform.structured.match ops{["test.bar"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below {{1 iterations}}
+  transform.test_tracked_rewrite %bar : (!transform.any_op) -> ()
+}

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "TestTransformDialectExtension.h"
 #include "TestTransformStateExtension.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
public:
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceSelfHandleOrForwardOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   if (getOperation()->getNumOperands() != 0) {
-    results.set(getResult().cast<OpResult>(),
-                getOperation()->getOperand(0).getDefiningOp());
+    results.set(
+        getResult().cast<OpResult>(),
+        ArrayRef<Operation *>(getOperation()->getOperand(0).getDefiningOp()));
   } else {
-    results.set(getResult().cast<OpResult>(), getOperation());
+    results.set(getResult().cast<OpResult>(),
+                ArrayRef<Operation *>(getOperation()));
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceSelfHandleOrForwardOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   if (getOperand())
     transform::onlyReadsHandle(getOperand(), effects);
void mlir::test::TestConsumeOperand::getEffects(
   transform::consumesHandle(getOperand(), effects);
   if (getSecondOperand())
     transform::consumesHandle(getSecondOperand(), effects);
   transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestConsumeOperandOfOpKindOrFail::apply(
     transform::TransformResults &results, transform::TransformState &state) {
-  ArrayRef<Operation *> payload = state.getPayloadOps(getOperand());
-  assert(payload.size() == 1 && "expected a single target op");
-  if (payload[0]->getName().getStringRef() != getOpKind()) {
+  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[0]->getName().getStringRef();
+           << getOpKind() << " got "
+           << (*payload.begin())->getName().getStringRef();
   }
 
   emitRemark() << "succeeded";
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestConsumeOperandOfOpKindOrFail::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {

 void mlir::test::TestSucceedIfOperandOfOpKind::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestPrintRemarkAtOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
-  ArrayRef<Operation *> payload = state.getPayloadOps(getOperand());
+  auto payload = state.getPayloadOps(getOperand());
   for (Operation *op : payload)
     op->emitRemark() << getMessage();
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintRemarkAtOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {

 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemapOperandPayloadToSelfOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   auto *extension = state.getExtension<TestTransformStateExtension>();
   if (!extension)
     return emitDefiniteFailure("TestTransformStateExtension missing");
 
-  if (failed(extension->updateMapping(state.getPayloadOps(getOperand()).front(),
-                                      getOperation())))
+  if (failed(extension->updateMapping(
+          *state.getPayloadOps(getOperand()).begin(), getOperation())))
     return DiagnosedSilenceableFailure::definiteFailure();
   if (getNumResults() > 0)
-    results.set(getResult(0).cast<OpResult>(), getOperation());
+    results.set(getResult(0).cast<OpResult>(),
+                ArrayRef<Operation *>(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::TransformResults &results, transform::TransformState &state) {
   state.removeExtension<TestTransformStateExtension>();
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestReversePayloadOpsOp::apply(transform::TransformResults &results,
                                            transform::TransformState &state) {
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto payloadOps = state.getPayloadOps(getTarget());
   auto reversedOps = llvm::to_vector(llvm::reverse(payloadOps));
   results.set(getResult().cast<OpResult>(), reversedOps);
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestTransformOpWithRegions::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
   return emitDefaultSilenceableFailure(target);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestPrintNumberOfAssociatedPayloadIROps::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   if (!getHandle())
     emitRemark() << 0;
-  emitRemark() << state.getPayloadOps(getHandle()).size();
+  emitRemark() << llvm::range_size(state.getPayloadOps(getHandle()));
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintNumberOfAssociatedPayloadIROps::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getHandle(), effects);
 }
 
mlir::test::TestProduceTransformParamOrForwardOperandOp::applyToOne(
     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()).front());
+    results.push_back(*state.getPayloadOps(getIn()).begin());
   }
 
   if (getSecondResultIsHandle()) {
-    results.push_back(state.getPayloadOps(getIn()).front());
+    results.push_back(*state.getPayloadOps(getIn()).begin());
   } else {
     results.push_back(builder.getI64IntegerAttr(42));
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceNullPayloadOp::getEffects(

 }
 
 DiagnosedSilenceableFailure mlir::test::TestRequiredMemoryEffectsOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   results.set(getOut().cast<OpResult>(), state.getPayloadOps(getIn()));
   return DiagnosedSilenceableFailure::success();
 }
 
+void mlir::test::TestTrackedRewriteOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestTrackedRewriteOp::apply(transform::TransformResults &results,
+                                        transform::TransformState &state) {
+  transform::TrackingListener 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;
+    rewriter.setInsertionPointToEnd(op->getBlock());
+
+    // Erase all payload ops. The outer loop should have only one iteration.
+    for (Operation *op : state.getPayloadOps(getIn())) {
+      SmallVector<NamedAttribute> attributes;
+      attributes.emplace_back(rewriter.getStringAttr("original_op"),
+                              op->getName().getIdentifier());
+      OperationState opState(op->getLoc(), "test.foo",
+                             /*operands=*/ValueRange(),
+                             /*types=*/op->getResultTypes(), attributes);
+      Operation *newOp = rewriter.create(opState);
+      rewriter.replaceOp(op, newOp->getResults());
+    }
+  }
+
+  emitRemark() << numIterations << " iterations";
+  return DiagnosedSilenceableFailure::success();
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

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

let arguments = (ins TransformHandleTypeInterface:$in,
                        UnitAttr:$has_operand_effect,
                        UnitAttr:$has_result_effect,
                        UnitAttr:$modifies_payload);
   let results = (outs TransformHandleTypeInterface:$out);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestTrackedRewriteOp
+  : Op<Transform_Dialect, "test_tracked_rewrite",
+      [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$in);
+  let results = (outs);
+  let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD