[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

590

Nit: please document.

647–665

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

656

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

667

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

1115

Same as above.

1183

Nit: std::empty(targets)

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

212–213

Same as above: Range &&.

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

377–379

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,
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. 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:
   /// Checks that the operation does not use invalidated handles as operands.
   /// Reports errors and returns failure if it does. Otherwise, invalidates the
   /// handles consumed by the operation as well as any handles pointing to
   /// payload IR operations nested in the operations associated with the
   /// consumed handles.
   LogicalResult
   checkAndRecordHandleInvalidation(TransformOpInterface transform);
 
+  /// Remove all nullptrs from op handles that were added by `replacePayloadOp`.
+  void compactOpHandles();
+
   /// 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;
 
+  /// Op handles may be temporarily mapped to nullptr to avoid invalidating
+  /// payload op iterators. This set contains all op handles with nullptrs.
+  /// These handles are "compacted" (i.e., nullptrs removed) at the end of each
+  /// transform.
+  DenseSet<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, std::forward<Range>(ops));
+  }
+
+  /// 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, std::initializer_list<Operation *> ops) {
+    set(value, ArrayRef<Operation *>(ops));
+  }
 
   /// 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::empty(targets)) {
     SmallVector<Operation *> emptyPayload;
     SmallVector<Attribute> emptyParams;
     for (OpResult r : this->getOperation()->getResults()) {
       if (isa<TransformParamTypeInterface>(r.getType()))
         transformResults.setParams(r, emptyParams);
       else if (isa<TransformValueHandleTypeInterface>(r.getType()))
         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);
-    for (Operation *op : boundedValueOps) {
+    for (Operation *op : state.getPayloadOps(handle)) {
       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());
-  for (Operation *target : payloadOps) {
+  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";

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

DiagnosedSilenceableFailure transform::MatchStructuredResultOp::matchOperation(
   }
 
   if (result.getUsers().empty()) {
     return emitSilenceableError()
            << "no users of the result #" << getPosition();
   }
   Operation *firstUser = *result.getUsers().begin();
   if (getAny()) {
-    results.set(cast<OpResult>(getResult()), firstUser);
+    results.set(cast<OpResult>(getResult()), {firstUser});
     return DiagnosedSilenceableFailure::success();
   }
   if (getSingle()) {
     if (!llvm::hasSingleElement(result.getUsers())) {
       return emitSilenceableError()
              << "more than one result user with single user requested";
     }
-    results.set(cast<OpResult>(getResult()), firstUser);
+    results.set(cast<OpResult>(getResult()), {firstUser});
     return DiagnosedSilenceableFailure::success();
   }
 
   return emitDefiniteFailure() << "unknown sub-predicate";
 }
 
 DiagnosedSilenceableFailure
 transform::MatchStructuredResultOp::getPositionFor(linalg::LinalgOp op,

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 (!isa<IntegerAttr>(ofr.get<Attribute>()))
         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);
-  for (Operation *op : payloadOps) {
+  for (Operation *op : state.getPayloadOps(packedHandle)) {
     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 (std::empty(producerOps)) {
     results.set(cast<OpResult>(getFusedOp()), 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(cast<OpResult>(getResult()), 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(cast<OpResult>(getPackedOp()), {});
+  if (std::empty(targetOps)) {
+    transformResults.set(cast<OpResult>(getPackedOp()),
+                         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()
            << ")";
   }
 
   // Unpack handles to constants or actual SSA index values.
   SmallVector<OpFoldResult> packedSizes;
   DiagnosedSilenceableFailure status = unpackSingleIndexResultPDLOperations(
       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());
+                       {maybeResult->packedLinalgOp.getOperation()});
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::PackOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getTarget(), effects);
   transform::onlyReadsHandle(getPackedSizes(), effects);
   transform::producesHandle(getPackedOp(), effects);
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());
-  // 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);
-  for (Operation *op : targetOps) {
+  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.
     // In the future we will want to order packings according to some metric.
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()) {
+  if (std::empty(packOrUnpackOps)) {
     transformResults.set(cast<OpResult>(getPackedOp()), {});
     transformResults.set(cast<OpResult>(getPackOp()), {});
     transformResults.set(cast<OpResult>(getUnPackOp()), {});
     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)

 
 //===---------------------------------------------------------------------===//
 // 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());
   if (failed(result))
     return emitDefiniteFailure() << "could not build packing loop nest";
 
   if (result->clonedLoopIvs.empty()) {
     transformResults.set(cast<OpResult>(getPackingLoop()),
-                         result->hoistedPadOp.getOperation());
+                         {result->hoistedPadOp.getOperation()});
     return DiagnosedSilenceableFailure::success();
   }
   auto outerPackedLoop =
       scf::getForInductionVarOwner(result->clonedLoopIvs.front());
   transformResults.set(cast<OpResult>(getPackingLoop()),
-                       outerPackedLoop.getOperation());
+                       {outerPackedLoop.getOperation()});
   return DiagnosedSilenceableFailure::success();
 }
 
 LogicalResult transform::HoistPadBuildPackingLoopNestOp::verify() {
   ArrayRef<int64_t> transpose = getTranspose();
   auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, transpose.size()));
   if (!std::is_permutation(sequence.begin(), sequence.end(), transpose.begin(),
                            transpose.end())) {

 
 //===----------------------------------------------------------------------===//
 // 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 (isa<TransformHandleTypeInterface>(getDynamicSplitPoint().getType())) {
       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 (isa<ParamType>(transformValue.getType())) {
       dynamicSizeProducers.push_back({});
       ArrayRef<Attribute> params = state.getParams(transformValue);
       paramSizes.push_back(
if (isa<ParamType>(transformValue.getType())) {
             << 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());
-  // 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;
 
   // Unpack handles.
   SmallVector<OpFoldResult> mixedNumThreads;
   DiagnosedSilenceableFailure status =
       getPackedNumThreads()
           ? unpackSingleIndexResultPDLOperations(
                 state, transformOp, mixedNumThreads, getPackedNumThreads())
           : unpackSingleIndexResultPDLOperations(
                 state, transformOp, mixedNumThreads, getMixedNumThreads());
   if (!status.succeeded())
     return status;
   SmallVector<OpFoldResult> mixedTileSizes;
   status = getPackedTileSizes()
                ? unpackSingleIndexResultPDLOperations(
                      state, transformOp, mixedTileSizes, getPackedTileSizes())
                : unpackSingleIndexResultPDLOperations(
                      state, transformOp, mixedTileSizes, getMixedTileSizes());
   if (!status.succeeded())
     return status;
 
-  for (Operation *target : targets) {
+  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);
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 (std::empty(targets))
     return DiagnosedSilenceableFailure::success();
 
   SmallVector<int64_t> vectorSizes;
   for (OpFoldResult sz : getMixedVectorSizes()) {
     if (sz.is<Attribute>()) {
       auto attr = sz.get<Attribute>();
       vectorSizes.push_back(cast<IntegerAttr>(attr).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());
   IRRewriter rewriter(getContext());
-  for (auto *op : payloadOps) {
+  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;
       auto loop = user->getParentOfType<LoopLikeOpInterface>();

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 the replacement object.
+  // In case a payload op was erased (replacement object is nullptr), a nullptr
+  // is stored in the mapping. These nullptrs are removed after each transform.
+  // Furthermore, nullptrs are not enumerated by payload op iterators. The
+  // relative order of ops is preserved.
+  //
+  // 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.
   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.
     for (Operation *&mapped : association) {
       if (mapped == op)
         mapped = replacement;
     }
+
+    if (replacement) {
       mappings.reverse[replacement].push_back(handle);
+    } else {
+      opHandlesToCompact.insert(handle);
+    }
   }
 
   return success();
 }
 
 LogicalResult
 transform::TransformState::replacePayloadValue(Value value, Value replacement) {
   SmallVector<Value> valueHandles;
   (void)getHandlesForPayloadValue(value, valueHandles);
 
   for (Value handle : 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 (llvm::isa<TransformHandleTypeInterface>(target.get().getType())) {
         FULL_LDBG("----recordOpHandleInvalidation\n");
-        ArrayRef<Operation *> payloadOps = getPayloadOps(target.get());
+        ArrayRef<Operation *> payloadOps = getPayloadOpsView(target.get());
         recordOpHandleInvalidation(target, payloadOps);
       } else if (llvm::isa<TransformValueHandleTypeInterface>(
                      target.get().getType())) {
         FULL_LDBG("----recordValueHandleInvalidation\n");
         recordValueHandleInvalidation(target);
       } else {
         FULL_LDBG("----not a TransformHandle -> SKIP AND DROP ON THE FLOOR\n");
       }
if (!seen.insert(p).second) {
       else
         diag.attachNote(p.getLoc()) << "repeated target value";
       return diag;
     }
   }
   return DiagnosedSilenceableFailure::success();
 }
 
+void transform::TransformState::compactOpHandles() {
+  for (Value handle : opHandlesToCompact) {
+    Mappings &mappings = getMapping(handle);
+    llvm::erase_value(mappings.direct[handle], nullptr);
+  }
+  opHandlesToCompact.clear();
+}
+
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG({
     DBGS() << "applying: ";
     transform->print(llvm::dbgs(), OpPrintingFlags().skipRegions());
     llvm::dbgs() << "\n";
   });
   DEBUG_WITH_TYPE(DEBUG_TYPE_FULL,
for (OpOperand &operand : transform->getOpOperands()) {
       }
       FULL_LDBG("--handle is consumed\n");
 
       Type operandType = operand.get().getType();
       if (llvm::isa<TransformHandleTypeInterface>(operandType)) {
         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 (llvm::isa<TransformValueHandleTypeInterface>(operandType)) {
         FULL_LDBG("--checkRepeatedConsumptionInOperand For Value\n");
         DiagnosedSilenceableFailure check =
#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
     return diag;
   }
 
   // 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));
+  compactOpHandles();
   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);
 

 //===----------------------------------------------------------------------===//
 
 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");

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

 
 //===----------------------------------------------------------------------===//
 // ForeachOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ForeachOp::apply(transform::TransformResults &results,
                             transform::TransformState &state) {
-  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
   SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});
 
-  for (Operation *op : payloadOps) {
+  for (Operation *op : state.getPayloadOps(getTarget())) {
     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(llvm::cast<OpResult>(getResult(i)), 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(llvm::cast<OpResult>(getResult()), {});
+  auto payloadOps = state.getPayloadOps(getTarget());
+  if (std::empty(payloadOps)) {
+    results.set(cast<OpResult>(getResult()), {});
     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(llvm::cast<OpResult>(getResult()),
               llvm::to_vector(target->getResult(resultNumber).getUsers()));
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
void transform::SplitHandleOp::build(OpBuilder &builder, OperationState &result,
   result.addOperands(target);
   auto pdlOpType = pdl::OperationType::get(builder.getContext());
   result.addTypes(SmallVector<pdl::OperationType>(numResultHandles, pdlOpType));
 }
 
 DiagnosedSilenceableFailure
 transform::SplitHandleOp::apply(transform::TransformResults &results,
                                 transform::TransformState &state) {
-  int64_t numPayloadOps = state.getPayloadOps(getHandle()).size();
+  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";
   };
 
   // Fail if there are more payload ops than results and no overflow result was

 
 //===----------------------------------------------------------------------===//
 // 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 (llvm::isa<TransformHandleTypeInterface>(handle.getType())) {
-      ArrayRef<Operation *> current = state.getPayloadOps(handle);
+    if (isa<TransformHandleTypeInterface>(handle.getType())) {
+      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(llvm::cast<OpResult>(getReplicated()[en.index()]), payload);
     } else {
       assert(llvm::isa<TransformParamTypeInterface>(handle.getType()) &&
              "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());
-  for (Operation *target : targets)
+  for (Operation *target : state.getPayloadOps(getTarget()))
     llvm::outs() << *target << "\n";
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::PrintOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getTarget(), 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.update_mapping"() {original_op = "test.replace_me"}
+//  CHECK-NEXT:   "test.drop_mapping"() {original_op = "test.replace_me"}
+//  CHECK-NEXT:   "test.update_mapping"() {original_op = "test.replace_me"}
+//  CHECK-NEXT: }
+func.func @test_tracked_rewrite() {
+  %0 = "test.replace_me"() {replacement = "test.update_mapping"} : () -> (i1)
+  %1 = "test.replace_me"() {replacement = "test.drop_mapping"} : () -> (i1)
+  %2 = "test.replace_me"() {replacement = "test.update_mapping"} : () -> (i1)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["test.replace_me"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  // expected-remark @below {{2 iterations}}
+  transform.test_tracked_rewrite %0 : (!pdl.operation) -> ()
+  // One replacement op (test.drop_mapping) is dropped from the mapping.
+  // expected-remark @below {{2}}
+  test_print_number_of_associated_payload_ir_ops %0
+}

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(llvm::cast<OpResult>(getResult()),
-                getOperation()->getOperand(0).getDefiningOp());
+    results.set(cast<OpResult>(getResult()),
+                {getOperation()->getOperand(0).getDefiningOp()});
   } else {
-    results.set(llvm::cast<OpResult>(getResult()), getOperation());
+    results.set(cast<OpResult>(getResult()), {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(llvm::cast<OpResult>(getResult(0)), getOperation());
+    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::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(llvm::cast<OpResult>(getResult()), 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(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);
+}
+
+namespace {
+/// A TrackingListener for test cases. When the replacement op is
+/// "test.update_mapping", it is considered as a replacement op in the transform
+/// state mapping. Otherwise, it is not and the original op is simply removed
+/// from the mapping.
+class TestTrackingListener : public transform::TrackingListener {
+  using transform::TrackingListener::TrackingListener;
+
+protected:
+  Operation *findReplacementOp(Operation *op,
+                               ValueRange newValues) const override {
+    if (newValues.size() != 1)
+      return nullptr;
+    Operation *replacement = newValues[0].getDefiningOp();
+    if (!replacement)
+      return nullptr;
+    if (replacement->getName().getStringRef() != "test.update_mapping")
+      return nullptr;
+    return replacement;
+  }
+};
+} // namespace
+
+DiagnosedSilenceableFailure
+mlir::test::TestTrackedRewriteOp::apply(transform::TransformResults &results,
+                                        transform::TransformState &state) {
+  TestTrackingListener 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())) {
+      if (op->getName().getStringRef() != "test.replace_me")
+        continue;
+      auto replacementName = op->getAttrOfType<StringAttr>("replacement");
+      if (!replacementName)
+        continue;
+      SmallVector<NamedAttribute> attributes;
+      attributes.emplace_back(rewriter.getStringAttr("original_op"),
+                              op->getName().getIdentifier());
+      OperationState opState(op->getLoc(), replacementName,
+                             /*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