[mlir] add support for transform dialect value handles

[ftynse]

Summary

Introduce support for the third kind of values in the transform dialect:
value handles. Similarly to operation handles, value handles are
pointing to a set of values in the payload IR. This enables
transformation to be targeted at specific values, such as individual
results of a multi-result payload operation without indirecting through
the producing op or block arguments that previously could not be easily
addressed. This is expected to support a broad class of memory-oriented
transformations such as selective bufferization, buffer assignment, and
memory transfer management.

Value handles are functionally similar to operation handles and require
similar implementation logic. The most important change concerns the
handle invalidation mechanism where operation and value handles can
affect each other.

This patch includes two cleanups that make it easier to introduce value
handles:

• RaggedArray structure that encapsulates the SmallVector of ArrayRef backed by flat SmallVector logic, frequently used in the transform interfaces implementation;
• rewrite the tests that associated payload handles with an integer value reinterpret_casted as a pointer, which were a frequent source of confusion and crashes when adding more debugging facilities that can inspect the payload.

[ftynse]

Remove spurious commented-out code.

[springerm]

mlir/docs/Dialects/Transform.md

48

Is this a new op introduced by this change? Is there also a transform.get_defining_op?

Is there a best practice when adding designing transforms ops: Is it better to return OpResults, operations or both?

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

158–160

Op handle interfaces are used to support types such operation<"scf.if">, right? What's the benefit of the value handle interface? (As opposed to a single value handle type.)

[ftynse]

mlir/docs/Dialects/Transform.md

48

This is a fake op used for an example, so are all the other ops in this snippet. We may end up having such a transform, but I rather expect us to match specific values directly rather than extract them from operations.

Whether a transformation returns a result or an entire operation depends on the kind of the transformation. Data transformations will return (and likely accept) value handles.

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

158–160

So we can do typed_value<i32> or value<shaped>, for example.

mlir/docs/Dialects/Transform.md

 
 ## Overview
 
 This dialect provides operations that can be used to control transformation
 of the IR using a different portion of the IR. It refers to the IR being
 transformed as payload IR, and to the IR guiding the transformation as
 transform IR.
 
-The main use case for this dialect is orchestrating fine-grain
-transformations on individual operations or sets thereof. For example, it
-may involve finding loop-like operations with specific properties (e.g.,
-large size) in the payload IR, applying loop tiling to those and only those
-operations, and then applying loop unrolling to the inner loops produced
-by the previous transformations. As such, it is not intended as a
-replacement for the pass infrastructure, nor for the pattern rewriting
-infrastructure. In the most common case, the transform IR will be processed
-and applied to the payload IR by a pass. Transformations expressed by the
-transform dialect may be implemented using the pattern infrastructure or any
-other relevant MLIR component.
+The main use case for this dialect is orchestrating fine-grain transformations
+on individual IR objects (operations or values) or sets thereof. For example, it
+may involve finding loop-like operations with specific properties (e.g., large
+size) in the payload IR, applying loop tiling to those and only those
+operations, and then applying loop unrolling to the inner loops produced by the
+previous transformations. As such, it is not intended as a replacement for the
+pass infrastructure, nor for the pattern rewriting infrastructure. In the most
+common case, the transform IR will be processed and applied to the payload IR by
+a pass. Transformations expressed by the transform dialect may be implemented
+using the pattern infrastructure or any other relevant MLIR component.
 
 The following IR gives a rough idea of what the operations in this dialect
-may look like:
+may look like without using actually existing operations:
 
 ```mlir
 %0 = transform.loop.find { size > 42 } : !transform.interface<tileable>
 %1 = transform.compute_trailing_tile_size %0 : !transform.param<index>
 %2:2 = transform.loop.tile %0 tile_sizes(1, 4, %1)
       : (!transform.interface<tileable>)
      -> (!transform.op<loop>, !transform.op<loop>)
+%3 = transform.get_op_result [0] %2#0 : !transform.any_value
+transform.assign_to_fast_memory %3
 transform.loop.unroll %1#1 : !transform.op<loop>
 ```
 
-The values used in the Transform dialect may correspond to either:
+The values used in the Transform dialect may correspond to:
 
   * sets of operations in the payload IR;
 
+  * sets of values in the payload IR;
+
   * sets of parameters (attributes) known at the execution time of the
     transform dialect.
 
-The former kind of values is also referred to as *handles*. In the example
-above, `%0` corresponds to the set of loops found in the payload IR that
-satisfy the condition, and `%2` correspond to groups of outer and inner
-loops, respectively, produced by the tiling transformation, whereas `%1`
-corresponds to a list of tile sizes selected for each of the operations
-that `%0` corresponds to.
+The former two kinds of values are also referred to as operation and value
+*handles*, respectively. In the example above, `%0` corresponds to the set of
+loops found in the payload IR that satisfy the condition, and `%2` correspond to
+groups of outer and inner loops, respectively, produced by the tiling
+transformation. `%3` corresponds to a set of values that are produced by the
+outer loops after tiling. `%1` corresponds to a list of tile sizes selected for
+each of the operations that `%0` corresponds to.
 
-A transform handle such as `%0` may be associated with multiple payload
+An operation handle such as `%0` may be associated with multiple payload
 operations. This is conceptually a set of operations and no assumptions should
 be made about the order of ops unless specified otherwise by the operation.
-Operations may take as operands and produce an arbitrary combination of values
-representing handles and parameters. Most Transform IR ops support operand
-values that are mapped to multiple operations. They usually apply the respective
-transformation for every mapped op ("batched execution"). Deviations from this
-convention are described in the documentation of Transform IR ops.
-
-The transform IR values have transform IR types, which implement either
-[TransformHandleTypeInterface](Transform.md#transformhandletypeinterface-transformhandletypeinterface)
-or
-[TransformParamTypeInterface](Transform.md##transformparamtypeinterface-transformparamtypeinterface).
-The former interface verifiers properties of payload IR operations associated
-with the value that are known to the transform dialect, for example, all
-associated payload operations implement a "TileableOp" interface, or have a
-specific "loop" kind. Similarly, the latter interface verifies properties of
-attributes associated with the parameter value. These properties are used to
-statically indicate pre- and post-conditions of a transformation connected to a
-Transform dialect operation. The conditions are verified when attributes or
-payload IR operations are first associated with a transform handle. By
-convention, Transform dialect operations are expected to indicate narrow
-preconditions for their operands by enforcing operand type constraints in the
-their definitions and verifiers. On the contrary, operations are expected to
-have few constraints on their results. Specific instances of a transform
-operation can then be created with a more restricted result type than the
-constraint in the operation (e.g., the "find" operation only constrains the
-result type to be a transform IR type while its concrete instance can have a
-type with stricter constraints such as implementing the "tilable" interface).
-The verification will then happen at transform execution time. This approach
-allows one to capture payload IR operation properties in the transform IR
-without resorting to excessive use of type casts or coupling dialect extensions
-between themselves. It is a trade-off between verbosity/complexity and static
-hardening, which can be revised in the future.
+Similarly, a value handle such as `%3` may be associated with a set of payload
+IR values. Transform dialect operations may take as operands and produce an
+arbitrary combination of values representing handles and parameters. Most
+Transform IR ops support operand values that are mapped to multiple payload
+objects. They usually apply the respective transformation for every mapped
+object ("batched execution"). Deviations from this convention are described in
+the documentation of Transform IR ops.
+
+The transform IR values have transform IR types, which should implement exactly one of:
+
+  * [TransformHandleTypeInterface](Transform.md#transformhandletypeinterface-transformhandletypeinterface),
+
+  * [TransformValueHandleTypeInterface](Transform.md#transformvaluehandletypeinterface-transformvaluehandletypeinterface),
+
+  * [TransformParamTypeInterface](Transform.md##transformparamtypeinterface-transformparamtypeinterface).
+
+The goal of these type interfaces, beyond providing a common base for accepted
+types, is to verify the properties of the associated objects. For example, a
+handle type interface implementation may check whether all associated payload IR
+operations implement the "TileableOp" interface or have a specific "loop" kind.
+Similarly, a value handle type interface implementation may check if the
+associated payload IR values are block arguments or have a specific type, or a
+parameter type interface may check whether the associated attributes contain
+non-negative integer values. These properties are used to statically indicate
+ pre- and post-conditions of a transformation connected to a Transform dialect
+operation. The conditions are verified when payload objects operations are first
+associated with a transform handle. By convention, Transform dialect operations
+are expected to indicate narrow preconditions for their operands by enforcing
+operand type constraints in the their definitions and verifiers. On the
+contrary, operations are expected to have few constraints on their results.
+Specific instances of a transform operation can then be created with a more
+restricted result type than the constraint in the operation (e.g., the "find"
+operation only constrains the result type to be a transform IR type while its
+concrete instance can have a type with stricter constraints such as implementing
+the "tilable" interface). The verification will then happen at transform
+execution time. This approach allows one to capture payload IR operation
+properties in the transform IR without resorting to excessive use of type casts
+or coupling dialect extensions between themselves. It is a trade-off between
+verbosity/complexity and static hardening, which can be revised in the future.
 
 Overall, Transform IR ops are expected to be contained in a single top-level
 op. Such top-level ops specify how to apply the transformations described
 by the operations they contain, e.g., `transform.sequence` executes
 transformations one by one and fails if any of them fails. Such ops are
 expected to have the `PossibleTopLevelTransformOpTrait` and may be used
 without arguments.
 
 A program transformation expressed using the Transform dialect can be
 programmatically triggered by calling:
 
 ```c++
 LogicalResult transform::applyTransforms(
     Operation *payloadRoot,
-    ArrayRef<ArrayRef<PointerUnion<Operation *, Attribute>> extraMappings,
+    const RaggedArray<transform::MappedValue> &extraMappings,
     TransformOpInterface transform,
     const TransformOptions &options);
 ```
 
 that applies the transformations specified by the top-level `transform` to
 payload IR contained in `payloadRoot`. The payload root operation will be
 associated with the first argument of the entry block of the top-level transform
 op. This block may have additional arguments, handles or parameters. They will

 The presence of interface implementations is checked at runtime when the
 dialect is loaded to allow for those implementations to be supplied by
 separate dialect extensions if desired.
 
 Similarly to operations, additional types can be injected into the dialect using
 the same extension mechanism. The types must:
 
   * Implement exactly one of `TransformHandleTypeInterface`,
-    `TransformParamTypeInterface`.
+    `TransformValueHandleTypeInterface`, `TransformParamTypeInterface`.
 
 ## Side Effects
 
 The Transform dialect relies on MLIR side effect modelling to enable
 optimization of the transform IR. More specifically, it provides several
 side effect resource objects and expects operations to describe their
 effects on these resources.
 

 would fail a postcondition. Some "navigation" operations that identify
 payload IR targets for the following transformation may have a conceptual
 "failure to match" that is considered a successful execution in the
 execution model but results in handles associated with empty payload IR
 operation lists.
 
 ## Handle Invalidation
 
-The execution model of the transform dialect allows a payload IR operation
-to be associated with _multiple_ handles as well as nested payload IR
-operations to be associated with different handles. A transform IR operation
-that consumes a handle automatically _invalidates_ all the other handles
-associated with the same payload IR operations, or with any of their
-descendants, as the consumed handle. Note that the _entire_ handle is
-invalidated, even if some of the payload IR operations associated with it
-or their ancestors were not associated with the consumed handle. Any use of
-the invalidated handle results in undefined behavior since the payload IR
-operations associated with it are likely to have been mutated or erased. The
-mere fact of the handle being invalidated does _not_ trigger undefined
-behavior, only its appearance as an operand does.
+The execution model of the transform dialect allows a payload IR operation to be
+associated with _multiple_ handles as well as nested payload IR operations to be
+associated with different handles. Similarly, a payload IR value may be
+associated with multiple transform IR value handles. When a transform IR
+operation consumes a handle, it usually indicates that the corresponding payload
+IR object was destroyed and should no longer be referenced. Transform IR handles
+that _may_ be pointing to an erased payload IR object are _invalidated_. The
+mere presence of an invalidated handle in the transform IR is not a problem, but
+_using_ it results in undefined behavior. Invalidated handles can be thought of
+as dangling pointers. Note that the _entire_ handle is invalidated, even if some
+of the payload IR objects associated with it remain live.
+
+The following handle invalidation rules apply.
+
+  * When an operation handle is consumed, are invalidated:
+
+    - operation handles associated with one of the payload operations that the
+      consumed handle is associated with;
+
+    - operation handles associated with one of the operations _nested_ in the
+      payload operations described above;
+
+    - value handles associated with any result of any operation described above;
+    
+    - value handles associated with any argument of a block contained in a
+      region attached to any operation described above.
+
+  * When a value handle is consumed, are invalidated:
+
+    - operation handles associated with payload operations that produce as
+      result any value associated with the consumed handle (when the associated
+      is an operation result);
+
+    - operation handles associated with payload operations _nested_ in the
+      payload operations described above;
+
+    - operation handles associated with payload operations (recursively)
+      _contained_ in the block that defines as argument any value associated
+      with the consumed handle (when the associated value is a block argument);
+      note that the adjacent blocks are not affected;
+
+    - value handles associated with any result of any operation described above,
+      including all results of the operation defining as result the value
+      associated with the consumed handle;
+    
+    - value handles associated with any argument of a block contained in a
+      region attached to any operation described above.
+
+More intuitively, consuming a handle invalidates any handle that may be pointing
+to an object defined or contained in the payload IR subtree rooted at the
+closest operation or block.
 
 The Transform dialect infrastructure has the capability of checking whether
 the transform IR op operand is invalidated before applying the
 transformation. However, such a check is computationally expensive and
 must be enabled explicitly through `TransformOptions`. Additionally, the
 `transform-dialect-check-uses` pass emits warnings when a handle may be used
 after it has been consumed, but does so abstractly, without processing the
 payload IR.

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

 //===- TransformInterfaces.h - Transform Dialect Interfaces -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H
 
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/Dialect/Transform/Utils/DiagnosedSilenceableFailure.h"
+#include "mlir/Dialect/Transform/Utils/RaggedArray.h"
 #include "mlir/IR/OpDefinition.h"
-
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 namespace transform {
 
 class TransformOpInterface;
 
public:
   /// 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>;
+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,
-                ArrayRef<ArrayRef<MappedValue>> extraMapping = {},
+                const RaggedArray<MappedValue> &extraMapping = {},
                 const TransformOptions &options = TransformOptions());
 
 /// The state maintained across applications of various ops implementing the
 /// TransformOpInterface. The operations implementing this interface and the
 /// surrounding structure are referred to as transform IR. The operations to
 /// which transformations apply are referred to as payload IR. Transform IR
 /// operates on values that can be associated either with a list of payload IR
 /// operations (such values are referred to as handles) or with a list of
private:
   /// associated with.
   using TransformOpReverseMapping =
       DenseMap<Operation *, SmallVector<Value, 2>>;
 
   /// Mapping between a Value in the transform IR and the corresponding list of
   /// parameters.
   using ParamMapping = DenseMap<Value, SmallVector<Param>>;
 
+  /// Mapping between a Value in the transform IR and the corrsponding list of
+  /// values in the payload IR. Also works for reverse mappings.
+  using ValueMapping = DenseMap<Value, SmallVector<Value>>;
+
   /// The bidirectional mappings between transform IR values and payload IR
   /// operations, and the mapping between transform IR values and parameters.
   struct Mappings {
     TransformOpMapping direct;
     TransformOpReverseMapping reverse;
     ParamMapping params;
+    ValueMapping values;
+    ValueMapping reverseValues;
   };
 
   friend LogicalResult applyTransforms(Operation *, TransformOpInterface,
-                                       ArrayRef<ArrayRef<MappedValue>>,
+                                       const RaggedArray<MappedValue> &,
                                        const TransformOptions &);
 
 public:
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
   /// 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 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;
+
   /// Populates `handles` with all handles pointing to the given Payload IR op.
   /// Returns success if such handles exist, failure otherwise.
   LogicalResult getHandlesForPayloadOp(Operation *op,
                                        SmallVectorImpl<Value> &handles) const;
 
+  /// Populates `handles` with all handles pointing to the given payload IR
+  /// value. Returns success if such handles exist, failure otherwise.
+  LogicalResult
+  getHandlesForPayloadValue(Value payloadValue,
+                            SmallVectorImpl<Value> &handles) const;
+
   /// Applies the transformation specified by the given transform op and updates
   /// the state accordingly.
   DiagnosedSilenceableFailure applyTransform(TransformOpInterface transform);
 
   /// Records the mapping between a block argument in the transform IR and a
   /// list of operations in the payload IR. The arguments must be defined in
   /// blocks of the currently processed transform IR region, typically after a
   /// region scope is defined.
private:
   /// Identifier for storing top-level value in the `operations` mapping.
   static constexpr Value kTopLevelValue = Value();
 
   /// Creates a state for transform ops living in the given region. The second
   /// argument points to the root operation in the payload IR being transformed,
   /// which may or may not contain the region with transform ops. Additional
   /// options can be provided through the trailing configuration object.
   TransformState(Region *region, Operation *payloadRoot,
-                 ArrayRef<ArrayRef<MappedValue>> extraMappings = {},
+                 const RaggedArray<MappedValue> &extraMappings = {},
                  const TransformOptions &options = TransformOptions());
 
-  /// Returns the mappings frame for the reigon in which the value is defined.
+  /// Returns the mappings frame for the region in which the value is defined.
   const Mappings &getMapping(Value value) const {
     return const_cast<TransformState *>(this)->getMapping(value);
   }
   Mappings &getMapping(Value value) {
     auto it = mappings.find(value.getParentRegion());
     assert(it != mappings.end() &&
            "trying to find a mapping for a value from an unmapped region");
     return it->second;
   }
 
   /// Returns the mappings frame for the region in which the operation resides.
   const Mappings &getMapping(Operation *operation) const {
     return const_cast<TransformState *>(this)->getMapping(operation);
   }
   Mappings &getMapping(Operation *operation) {
     auto it = mappings.find(operation->getParentRegion());
     assert(it != mappings.end() &&
            "trying to find a mapping for an operation from an unmapped region");
     return it->second;
   }
 
-  /// Removes the mapping between the given payload IR operation and the given
-  /// transform IR value.
-  void dropReverseMapping(Mappings &mappings, Operation *op, Value value);
-
   /// 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" }
   ///
   /// which is valid by itself. However, calling a hypothetical "rewrite and
   /// rename function" transform on both handles:
   ///
   ///   transform.rewrite_and_rename %0 { new_name = "func" }
   ///   transform.rewrite_and_rename %1 { new_name = "func" }
   ///
   /// is invalid given the transformation "consumes" the handle as expressed
   /// by side effects. Practically, a transformation consuming a handle means
   /// that the associated payload operation may no longer exist.
   ///
+  /// Similarly, operation handles may be invalidate and should not be used
+  /// after a transform that consumed a value handle pointing to a payload value
+  /// defined by the operation as either block argument or op result. For
+  /// example, in the following sequence, the last transform operation rewrites
+  /// the callee to not return a specified result:
+  ///
+  ///   %0 = transform.find_call "myfunc"
+  ///   %1 = transform.find_results_of_calling "myfunc"
+  ///   transform.drop_call_result_from_signature %1[0]
+  ///
+  /// which requires the call operations to be recreated. Therefore, the handle
+  /// %0 becomes associated with a dangling pointer and should not be used.
+  ///
   /// Returns failure if the payload does not satisfy the conditions associated
   /// with the type of the handle value. The value is expected to have a type
   /// implementing TransformHandleTypeInterface.
   LogicalResult setPayloadOps(Value value, ArrayRef<Operation *> targets);
 
+  /// Sets the payload IR values association with the given transform IR value
+  /// (handle). A payload value may be associated with multiple handles as long
+  /// as at most one of them is consumed by further transformations. For
+  /// example, a hypothetical "get results of calls to function with the given
+  /// name" transform may be performed twice in a row producing handles pointing
+  /// to the same values:
+  ///
+  ///   %0 = transform.find_results_of_calling "myfunc"
+  ///   %1 = transform.find_results_of_calling "myfunc"
+  ///
+  /// which is valid by itself. However, calling a hypothetical "erase value
+  /// producer" transform on both handles:
+  ///
+  ///   transform.erase_value_produce %0
+  ///   transform.erase_value_produce %1
+  ///
+  /// is invalid provided the transformation "consumes" the handle as expressed
+  /// by side effects (which themselves reflect the semantics of the transform
+  /// erasing the producer and making the handle dangling). Practically, a
+  /// transformation consuming a handle means the associated payload value may
+  /// no longer exist.
+  ///
+  /// Similarly, value handles are invalidated and should not be used after a
+  /// transform that consumed an operation handle pointing to the payload IR
+  /// operation defining the values associated the value handle, as either block
+  /// arguments or op results, or any ancestor operation. For example,
+  ///
+  ///   %0 = transform.find_call "myfunc"
+  ///   %1 = transform.find_results_of_calling "myfunc"
+  ///   transform.rewrite_and_rename %0 { new_name = "func" }
+  ///
+  /// makes %1 unusable after the last transformation if it consumes %0. When an
+  /// operation handle is consumed, it usually indicates that the operation was
+  /// destroyed or heavily modified, meaning that the values it defines may no
+  /// longer exist.
+  ///
+  /// Returns failure if the payload values do not satisfy the conditions
+  /// associated with the type of the handle value. The value is expected to
+  /// have a type implementing TransformValueHandleTypeInterface.
+  LogicalResult setPayloadValues(Value handle, ValueRange payloadValues);
+
   /// Sets the parameters associated with the given transform IR value. Returns
   /// failure if the parameters do not satisfy the conditions associated with
   /// the type of the value. The value is expected to have a type implementing
   /// TransformParamTypeInterface.
   LogicalResult setParams(Value value, ArrayRef<Param> params);
 
-  /// Forgets the payload IR ops associated with the given transform IR value.
-  void removePayloadOps(Value value);
+  /// Forgets the payload IR ops associated with the given transform IR value,
+  /// as well as any association between value handles and the results of said
+  /// payload IR op.
+  void forgetMapping(Value opHandle, ValueRange origOpFlatResults);
+
+  void forgetValueMapping(Value valueHandle,
+                          ArrayRef<Operation *> payloadOperations);
 
   /// Updates the payload IR ops associated with the given transform IR value.
   /// The callback function is called once per associated operation and is
   /// expected to return the modified operation or nullptr. In the latter case,
   /// the corresponding operation is no longer associated with the transform IR
-  /// value.
+  /// value. Value handles associated with the results of the operation are
+  /// also updated to be associated with the results of the new operation. For
+  /// this reason, the new operation must have the same number of results.
   ///
   /// Returns failure if the payload does not satisfy the conditions associated
   /// with the type of the handle value.
-  LogicalResult
-  updatePayloadOps(Value value,
-                   function_ref<Operation *(Operation *)> callback);
+  LogicalResult replacePayloadOp(Operation *op, Operation *replacement);
 
   /// If the operand is a handle consumed by the operation, i.e. has the "free"
   /// memory effect associated with it, identifies other handles that are
   /// pointing to payload IR operations nested in the operations pointed to by
   /// the consumed handle. Marks all such handles as invalidated to trigger
-  /// errors if they are used.
-  void recordHandleInvalidation(OpOperand &handle);
-  void recordHandleInvalidationOne(OpOperand &handle, Operation *payloadOp,
-                                   Value otherHandle);
+  /// errors if they are used. If `throughValue` is passed, record the fact that
+  /// an op handle was invalidated because a value handle associated with
+  /// results of the payload op or its block arguments was invalidated.
+  void recordOpHandleInvalidation(OpOperand &consumingHandle,
+                                  ArrayRef<Operation *> potentialAncestors,
+                                  Value throughValue = nullptr);
+  void recordOpHandleInvalidationOne(OpOperand &handle,
+                                     ArrayRef<Operation *> potentialAncestors,
+                                     Operation *payloadOp, Value otherHandle,
+                                     Value throughValue = nullptr);
+
+  void recordValueHandleInvalidationByOpHandleOne(
+      OpOperand &opHandle, ArrayRef<Operation *> potentialAncestors,
+      Value payloadValue, Value valueHandle);
+
+  void recordValueHandleInvalidation(OpOperand &valueHandle);
 
   /// 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);
 
   /// 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;
 
   /// 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;
 
-  /// Storage for extra mapped values (payload operations or parameters) to be
+  /// Extra mapped values (payload operations, values or parameters) to be
   /// associated with additional entry block arguments of the top-level
-  /// transform operation. Each entry in `topLevelMappedValues` is a reference
-  /// to a contiguous block in `topLevelMappedValueStorage`.
-  // TODO: turn this into a proper named data structure, there are several more
-  // below.
-  SmallVector<ArrayRef<MappedValue>> topLevelMappedValues;
-  SmallVector<MappedValue> topLevelMappedValueStorage;
+  /// transform operation.
+  RaggedArray<MappedValue> topLevelMappedValues;
 
   /// Additional options controlling the transformation state behavior.
   TransformOptions options;
 
   /// The mapping from invalidated handles to the error-reporting functions that
   /// describe when the handles were invalidated. Calling such a function emits
   /// a user-visible diagnostic with an additional note pointing to the given
   /// location.

 /// Local mapping between values defined by a specific op implementing the
 /// TransformOpInterface and the payload IR ops they correspond to.
 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. The value must have a type
-  /// implementing TransformHandleTypeInterface.
+  /// 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);
 
   /// 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. The value must have a type implementing
-  /// TransformParamTypeInterface.
+  /// 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
+  /// succeeding. The value must have a type implementing
+  /// TransformValueHandleTypeInterface.
+  void setValues(OpResult handle, ValueRange values);
+
 private:
   /// Creates an instance of TransformResults that expects mappings for
   /// `numSegments` values, which may be associated with payload operations or
   /// parameters.
   explicit TransformResults(unsigned numSegments);
 
   /// Gets the list of operations associated with the result identified by its
   /// number in the list of operation results. The result must have been set to
   /// be associated with payload IR operations.
   ArrayRef<Operation *> get(unsigned resultNumber) const;
 
   /// Gets the list of parameters associated with the result identified by its
   /// number in the list of operation results. The result must have been set to
   /// be associated with parameters.
   ArrayRef<TransformState::Param> getParams(unsigned resultNumber) const;
 
+  /// Gets the list of payload IR values associated with the result identified
+  /// by its number in the list of operation results. The result must have been
+  /// set to be associated with payload IR values.
+  ArrayRef<Value> getValues(unsigned resultNumber) const;
+
   /// Returns `true` if the result identified by its number in the list of
-  /// operation results is associated with a list of parameters, `false` if it
-  /// is associated with the list of payload IR operations.
+  /// operation results is associated with a list of parameters, `false`
+  /// otherwise.
   bool isParam(unsigned resultNumber) const;
 
   /// Returns `true` if the result identified by its number in the list of
+  /// operation results is associated with a list of payload IR value, `false`
+  /// otherwise.
+  bool isValue(unsigned resultNumber) const;
+
+  /// Returns `true` if the result identified by its number in the list of
   /// operation results is associated with something.
   bool isSet(unsigned resultNumber) const;
 
-  /// Storage for pointers to payload IR ops that are associated with results of
-  /// a transform IR op. `segments` contains as many entries as the transform IR
-  /// op has results, even if some of them are not associated with payload IR
-  /// operations. Each entry is a reference to a contiguous segment in the
-  /// `operations` list that contains the pointers to operations. This allows
-  /// for operations to be stored contiguously without nested vectors and for
-  /// different segments to be set in any order.
-  SmallVector<ArrayRef<Operation *>, 2> segments;
-  SmallVector<Operation *> operations;
-
-  /// Storage for parameters that are associated with results of the transform
-  /// IR op. `paramSegments` contains as many entries as the transform IR op has
-  /// results, even if some of them are not associated with parameters. Each
-  /// entry is a reference to a contiguous segment in the `params` list that
-  /// contains the actual parameters. This allows for parameters to be stored
-  /// contiguously without nested vectors and for different segments to be set
-  /// in any order.
-  SmallVector<ArrayRef<TransformState::Param>, 2> paramSegments;
-  SmallVector<TransformState::Param> params;
+  /// Pointers to payload IR ops that are associated with results of a transform
+  /// IR op.
+  RaggedArray<Operation *> operations;
+
+  /// Parameters that are associated with results of the transform IR op.
+  RaggedArray<Param> params;
+
+  /// Payload IR values that are associated with results of a transform IR op.
+  RaggedArray<Value> values;
 };
 
 TransformState::RegionScope TransformState::make_region_scope(Region &region) {
   return RegionScope(*this, region);
 }
 
 namespace detail {
 /// Maps the only block argument of the op with PossibleTopLevelTransformOpTrait
public:
 
   /// Checks that the op matches the expectations of this trait.
   static LogicalResult verifyTrait(Operation *op);
 };
 
 /// Side effect resource corresponding to the mapping between Transform IR
 /// values and Payload IR operations. An Allocate effect from this resource
 /// means creating a new mapping entry, it is always accompanied by a Write
-/// effet. A Read effect from this resource means accessing the mapping. A Free
+/// effect. A Read effect from this resource means accessing the mapping. A Free
 /// effect on this resource indicates the removal of the mapping entry,
 /// typically after a transformation that modifies the Payload IR operations
 /// associated with one of the Transform IR operation's operands. It is always
 /// accompanied by a Read effect. Read-after-Free and double-Free are not
 /// allowed (they would be problematic with "regular" memory effects too) as
 /// they indicate an attempt to access Payload IR operations that have been
-/// modified, potentially erased, by the previous tranfsormations.
+/// modified, potentially erased, by the previous transformations.
 // TODO: consider custom effects if these are not enabling generic passes such
 // as CSE/DCE to work.
 struct TransformMappingResource
     : public SideEffects::Resource::Base<TransformMappingResource> {
   StringRef getName() override { return "transform.mapping"; }
 };
 
 /// Side effect resource corresponding to the Payload IR itself. Only Read and

 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h.inc"
 
 namespace mlir {
 namespace transform {
 
 /// A single result of applying a transform op with `ApplyEachOpTrait` to a
 /// single payload operation.
-using ApplyToEachResult = llvm::PointerUnion<Operation *, Attribute>;
+using ApplyToEachResult = MappedValue;
 
 /// A list of results of applying a transform op with `ApplyEachOpTrait` to a
 /// single payload operation, co-indexed with the results of the transform op.
 class ApplyToEachResultList {
 public:
   ApplyToEachResultList() = default;
   explicit ApplyToEachResultList(unsigned size) : results(size) {}
 
   /// Sets the list of results to `size` null pointers.
   void assign(unsigned size, std::nullptr_t) { results.assign(size, nullptr); }
 
   /// Sets the list of results to the given range of values.
   template <typename Range>
   void assign(Range &&range) {
     // This is roughly the implementation of SmallVectorImpl::assign.
     // Dispatching to it with map_range and template type inference would result
     // in more complex code here.
     results.clear();
     results.reserve(llvm::size(range));
     for (auto element : range) {
       if constexpr (std::is_convertible_v<decltype(*std::begin(range)),
                                           Operation *>) {
         results.push_back(static_cast<Operation *>(element));
+      } else if constexpr (std::is_convertible_v<decltype(*std::begin(range)),
+                                                 Value>) {
+        results.push_back(element.template get<Value>());
       } else {
         results.push_back(static_cast<Attribute>(element));
       }
     }
   }
 
   /// Appends an element to the list.
+  // Using ApplyToEachResult that can be implicitly constructed from a Value but
+  // not from a concrete Op that is implicitly convertible to a Value to avoid
+  // ambiguity.
   void push_back(Operation *op) { results.push_back(op); }
   void push_back(Attribute attr) { results.push_back(attr); }
+  void push_back(ApplyToEachResult r) { results.push_back(r); }
 
   /// Reserves space for `size` elements in the list.
   void reserve(unsigned size) { results.reserve(size); }
 
   /// Iterators over the list.
   auto begin() { return results.begin(); }
   auto end() { return results.end(); }
   auto begin() const { return results.begin(); }

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

let extraSharedClassDeclaration = [{
     }
   }];
 }
 
 def TransformHandleTypeInterface
     : TransformTypeInterfaceBase<"TransformHandleTypeInterface",
                                  "::mlir::Operation *"> {
   let description = [{
-    Types that can be used for the Transform dialect handle values. Such types
-    define the properties of Payload IR operations associated with the handle.
-    A user of such a handle can assume that these properties have been verified
-    for any Payload IR operation associated with it.
+    Types that can be used for the Transform dialect operation handle values.
+    Such types define the properties of Payload IR operations associated with
+    the handle. A user of such a handle can assume that these properties have
+    been verified for any Payload IR operation associated with it.
   }];
 }
 
 def TransformParamTypeInterface
     : TransformTypeInterfaceBase<"TransformParamTypeInterface",
                                  "::mlir::Attribute"> {
   let description = [{
     Types that can be used for the Transform dialect parameter values. Such types
     define the structure of the parameters associated with the value, e.g., their
     underlying type. A user of the value can assume that the parameter has been
     verified.
   }];
 }
 
+def TransformValueHandleTypeInterface
+    : TransformTypeInterfaceBase<"TransformValueHandleTypeInterface",
+                                 "::mlir::Value"> {
+  let description = [{
+    Types that can be used for the Transform dialect handle values pointing to
+    Payload IR values. Such types define the properties of Payload IR values
+    associated with the handle. Users of such a handle can assume that these
+    properties have been verified for any Payload IR value associated with it.
+  }];
+}
+
 def Transform_AnyHandleOrParamType
   : Type<Or<[TransformParamTypeInterface.predicate,
-             TransformHandleTypeInterface.predicate]>,
+             TransformHandleTypeInterface.predicate,
+             TransformValueHandleTypeInterface.predicate]>,
          "any transform handle or parameter">;
 
 def FunctionalStyleTransformOpTrait
     : NativeOpTrait<"FunctionalStyleTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def TransformEachOpTrait : NativeOpTrait<"TransformEachOpTrait"> {

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

def Transform_ParamType : TypeDef<Transform_Dialect, "Param",
   let mnemonic = "param";
   let parameters = (ins
     TypeParameter<"::mlir::Type", "Underlying type of the parameter">:$type
   );
   let assemblyFormat = "`<` $type `>`";
   let genVerifyDecl = 1;
 }
 
+def Transform_AnyValue : TypeDef<Transform_Dialect, "AnyValue",
+    [DeclareTypeInterfaceMethods<TransformValueHandleTypeInterface>]> {
+  let description = [{
+    Transform IR value that can be associated with a list of Payload IR values.
+  }];
+  let mnemonic = "any_value";
+  let assemblyFormat = "";
+}
+
 class Transform_ConcreteOpType<string opname>
   : Type<And<[Transform_OperationType.predicate,
               CPred<"$_self.cast<::mlir::transform::OperationType>()"
                     ".getOperationName() == \"" # opname # "\"">]>,
          "Transform IR handle to " # opname # " operations",
          "::mlir::transform::OperationType">;
 
 #endif  // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMTYPES

mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h

 interpreterBaseInitializeImpl(MLIRContext *context, StringRef transformFileName,
                               std::shared_ptr<OwningOpRef<ModuleOp>> &module);
 
 /// Template-free implementation of
 /// TransformInterpreterPassBase::runOnOperation.
 LogicalResult interpreterBaseRunOnOperationImpl(
     Operation *target, StringRef passName,
     const std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
-    ArrayRef<ArrayRef<MappedValue>> extraMappings,
+    const RaggedArray<MappedValue> &extraMappings,
     const TransformOptions &options,
     const Pass::Option<std::string> &transformFileName,
     const Pass::Option<std::string> &debugPayloadRootTag,
     const Pass::Option<std::string> &debugTransformRootTag,
     StringRef binaryName);
 } // namespace detail
 
 /// Base class for transform dialect interpreter passes that can consume and

mlir/include/mlir/Dialect/Transform/Utils/RaggedArray.h

+//===- RaggedArray.h - 2D array with different inner lengths ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+/// A 2D array where each row may have different length. Elements of each row
+/// are stored contiguously, but rows don't have a fixed order in the storage.
+template <typename T>
+class RaggedArray {
+public:
+  /// Returns the number of rows in the 2D array.
+  size_t size() const { return slices.size(); }
+
+  /// Returns true if the are no rows in the 2D array. Note that an array with a
+  /// non-zero number of empty rows is *NOT* empty.
+  bool empty() const { return slices.empty(); }
+
+  /// Accesses `pos`-th row.
+  ArrayRef<T> operator[](size_t pos) const { return at(pos); }
+  ArrayRef<T> at(size_t pos) const { return slices[pos]; }
+  MutableArrayRef<T> operator[](size_t pos) { return at(pos); }
+  MutableArrayRef<T> at(size_t pos) { return slices[pos]; }
+
+  /// Iterator over rows.
+  auto begin() { return slices.begin(); }
+  auto begin() const { return slices.begin(); }
+  auto end() { return slices.end(); }
+  auto end() const { return slices.end(); }
+
+  /// Reserve space to store `size` rows with `nestedSize` elements each.
+  void reserve(size_t size, size_t nestedSize = 0) {
+    slices.reserve(size);
+    storage.reserve(size * nestedSize);
+  }
+
+  /// Appends the given range of elements as a new row to the 2D array. May
+  /// invalidate the end iterator.
+  template <typename Range>
+  void push_back(Range &&elements) {
+    slices.push_back(appendToStorage(std::forward<Range>(elements)));
+  }
+
+  /// Replaces the `pos`-th row in the 2D array with the given range of
+  /// elements. Invalidates iterators and references to `pos`-th and all
+  /// succeeding rows.
+  template <typename Range>
+  void replace(size_t pos, Range &&elements) {
+    auto from = slices[pos].data();
+    if (from != nullptr) {
+      auto to = std::next(from, slices[pos].size());
+      auto newFrom = storage.erase(from, to);
+      // Update the array refs after the underlying storage was shifted.
+      for (size_t i = pos + 1, e = size(); i < e; ++i) {
+        slices[i] = MutableArrayRef<T>(newFrom, slices[i].size());
+        std::advance(newFrom, slices[i].size());
+      }
+    }
+    slices[pos] = appendToStorage(std::forward<Range>(elements));
+  }
+
+  /// Appends `num` empty rows to the array.
+  void appendEmptyRows(size_t num) { slices.resize(slices.size() + num); }
+
+private:
+  /// Appends the given elements to the storage and returns an ArrayRef pointing
+  /// to them in the storage.
+  template <typename Range>
+  MutableArrayRef<T> appendToStorage(Range &&elements) {
+    size_t start = storage.size();
+    llvm::append_range(storage, std::forward<Range>(elements));
+    return MutableArrayRef<T>(storage).drop_front(start);
+  }
+
+  /// Outer elements of the ragged array. Each entry is a reference to a
+  /// contiguous segment in the `storage` list that contains the actual
+  /// elements. This allows for elements to be stored contiguously without
+  /// nested vectors and for different segments to be set or replaced in any
+  /// order.
+  SmallVector<MutableArrayRef<T>> slices;
+
+  /// Dense storage for ragged array elements.
+  SmallVector<T> storage;
+};
+} // namespace mlir

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

void transform::detail::checkImplementsTransformOpInterface(
   assert(opName.hasInterface<MemoryEffectOpInterface>() &&
          "ops injected into the transform dialect must implement "
          "MemoryEffectsOpInterface");
 }
 
 void transform::detail::checkImplementsTransformHandleTypeInterface(
     TypeID typeID, MLIRContext *context) {
   const auto &abstractType = AbstractType::lookup(typeID, context);
-  assert(
-      (abstractType.hasInterface(
+  assert((abstractType.hasInterface(
               TransformHandleTypeInterface::getInterfaceID()) ||
           abstractType.hasInterface(
-           TransformParamTypeInterface::getInterfaceID())) &&
-      "expected Transform dialect type to implement one of the two interfaces");
+              TransformParamTypeInterface::getInterfaceID()) ||
+          abstractType.hasInterface(
+              TransformValueHandleTypeInterface::getInterfaceID())) &&
+         "expected Transform dialect type to implement one of the three "
+         "interfaces");
 }
 #endif // NDEBUG
 
 namespace {
 struct PDLOperationTypeTransformHandleTypeInterfaceImpl
     : public transform::TransformHandleTypeInterface::ExternalModel<
           PDLOperationTypeTransformHandleTypeInterfaceImpl,
           pdl::OperationType> {

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

 //===- TransformInterfaces.cpp - Transform Dialect Interfaces -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DEBUG_PRINT_AFTER_ALL "transform-dialect-print-top-level-after-all"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TransformState
 //===----------------------------------------------------------------------===//
 
 constexpr const Value transform::TransformState::kTopLevelValue;
 
 transform::TransformState::TransformState(
     Region *region, Operation *payloadRoot,
-    ArrayRef<ArrayRef<MappedValue>> extraMappings,
+    const RaggedArray<MappedValue> &extraMappings,
     const TransformOptions &options)
     : topLevel(payloadRoot), options(options) {
   topLevelMappedValues.reserve(extraMappings.size());
-  for (ArrayRef<MappedValue> mapping : extraMappings) {
-    size_t start = topLevelMappedValueStorage.size();
-    llvm::append_range(topLevelMappedValueStorage, mapping);
-    topLevelMappedValues.push_back(
-        ArrayRef<MappedValue>(topLevelMappedValueStorage)
-            .slice(start, mapping.size()));
-  }
+  for (ArrayRef<MappedValue> mapping : extraMappings)
+    topLevelMappedValues.push_back(mapping);
 
   auto result = mappings.try_emplace(region);
   assert(result.second && "the region scope is already present");
   (void)result;
 #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 {
   const TransformOpMapping &operationMapping = getMapping(value).direct;
   auto iter = operationMapping.find(value);
-  assert(iter != operationMapping.end() &&
-         "cannot find mapping for payload handle (param handle provided?)");
+  assert(
+      iter != operationMapping.end() &&
+      "cannot find mapping for payload handle (param/value handle provided?)");
   return iter->getSecond();
 }
 
 ArrayRef<Attribute> transform::TransformState::getParams(Value value) const {
   const ParamMapping &mapping = getMapping(value).params;
   auto iter = mapping.find(value);
-  assert(iter != mapping.end() &&
-         "cannot find mapping for param handle (payload handle provided?)");
+  assert(iter != mapping.end() && "cannot find mapping for param handle "
+                                  "(operation/value handle provided?)");
+  return iter->getSecond();
+}
+
+ArrayRef<Value>
+transform::TransformState::getPayloadValues(Value handleValue) const {
+  const ValueMapping &mapping = getMapping(handleValue).values;
+  auto iter = mapping.find(handleValue);
+  assert(iter != mapping.end() && "cannot find mapping for value handle "
+                                  "(param/operation handle provided?)");
   return iter->getSecond();
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadOp(
     Operation *op, SmallVectorImpl<Value> &handles) const {
   bool found = false;
   for (const Mappings &mapping : llvm::make_second_range(mappings)) {
     auto iterator = mapping.reverse.find(op);
     if (iterator != mapping.reverse.end()) {
       llvm::append_range(handles, iterator->getSecond());
       found = true;
     }
   }
 
   return success(found);
 }
 
+LogicalResult transform::TransformState::getHandlesForPayloadValue(
+    Value payloadValue, SmallVectorImpl<Value> &handles) const {
+  bool found = false;
+  for (const Mappings &mapping : llvm::make_second_range(mappings)) {
+    auto iterator = mapping.reverseValues.find(payloadValue);
+    if (iterator != mapping.reverseValues.end()) {
+      llvm::append_range(handles, iterator->getSecond());
+      found = true;
+    }
+  }
+
+  return success(found);
+}
+
 LogicalResult
 transform::TransformState::mapBlockArgument(BlockArgument argument,
                                             ArrayRef<MappedValue> values) {
   if (argument.getType().isa<TransformHandleTypeInterface>()) {
     SmallVector<Operation *> operations;
     operations.reserve(values.size());
     for (MappedValue value : values) {
       if (auto *op = value.dyn_cast<Operation *>()) {
         operations.push_back(op);
         continue;
       }
       return emitError(argument.getLoc())
              << "wrong kind of value provided for top-level operation handle";
     }
     return setPayloadOps(argument, operations);
   }
 
+  if (argument.getType().isa<TransformValueHandleTypeInterface>()) {
+    SmallVector<Value> payloadValues;
+    payloadValues.reserve(values.size());
+    for (MappedValue value : values) {
+      if (auto v = value.dyn_cast<Value>()) {
+        payloadValues.push_back(v);
+        continue;
+      }
+      return emitError(argument.getLoc())
+             << "wrong kind of value provided for the top-level value handle";
+    }
+    return setPayloadValues(argument, payloadValues);
+  }
+
   assert(argument.getType().isa<TransformParamTypeInterface>() &&
          "unsupported kind of block argument");
   SmallVector<Param> parameters;
   parameters.reserve(values.size());
   for (MappedValue value : values) {
     if (auto attr = value.dyn_cast<Attribute>()) {
       parameters.push_back(attr);
       continue;
     }
     return emitError(argument.getLoc())
            << "wrong kind of value provided for top-level parameter";
   }
   return setParams(argument, parameters);
 }
 
 LogicalResult
 transform::TransformState::setPayloadOps(Value value,
                                          ArrayRef<Operation *> targets) {
   assert(value != kTopLevelValue &&
          "attempting to reset the transformation root");
-  assert(!value.getType().isa<TransformParamTypeInterface>() &&
-         "cannot associate payload ops with a value of parameter type");
+  assert(value.getType().isa<TransformHandleTypeInterface>() &&
+         "wrong handle type");
 
   for (Operation *target : targets) {
     if (target)
       continue;
     return emitError(value.getLoc())
            << "attempting to assign a null payload op to this transform value";
   }
 
transform::TransformState::setPayloadOps(Value value,
   (void)inserted;
 
   for (Operation *op : targets)
     mappings.reverse[op].push_back(value);
 
   return success();
 }
 
+LogicalResult
+transform::TransformState::setPayloadValues(Value handle,
+                                            ValueRange payloadValues) {
+  assert(handle != nullptr && "attempting to set params for a null value");
+  assert(handle.getType().isa<TransformValueHandleTypeInterface>() &&
+         "wrong handle type");
+
+  for (Value payload : payloadValues) {
+    if (payload)
+      continue;
+    return emitError(handle.getLoc()) << "attempting to assign a null payload "
+                                         "value to this transform handle";
+  }
+
+  auto iface = handle.getType().cast<TransformValueHandleTypeInterface>();
+  SmallVector<Value> payloadValueVector = llvm::to_vector(payloadValues);
+  DiagnosedSilenceableFailure result =
+      iface.checkPayload(handle.getLoc(), payloadValueVector);
+  if (failed(result.checkAndReport()))
+    return failure();
+
+  Mappings &mappings = getMapping(handle);
+  bool inserted =
+      mappings.values.insert({handle, std::move(payloadValueVector)}).second;
+  assert(
+      inserted &&
+      "value handle is already associated with another list of payload values");
+  (void)inserted;
+
+  for (Value payload : payloadValues)
+    mappings.reverseValues[payload].push_back(handle);
+
+  return success();
+}
+
 LogicalResult transform::TransformState::setParams(Value value,
                                                    ArrayRef<Param> params) {
   assert(value != nullptr && "attempting to set params for a null value");
 
   for (Attribute attr : params) {
     if (attr)
       continue;
     return emitError(value.getLoc())
LogicalResult transform::TransformState::setParams(Value value,
   Mappings &mappings = getMapping(value);
   bool inserted =
       mappings.params.insert({value, llvm::to_vector(params)}).second;
   assert(inserted && "value is already associated with another list of params");
   (void)inserted;
   return success();
 }
 
-void transform::TransformState::dropReverseMapping(Mappings &mappings,
-                                                   Operation *op, Value value) {
-  auto it = mappings.reverse.find(op);
-  if (it == mappings.reverse.end())
+template <typename Mapping, typename Key, typename Mapped>
+void dropMappingEntry(Mapping &mapping, Key key, Mapped mapped) {
+  auto it = mapping.find(key);
+  if (it == mapping.end())
     return;
 
-  llvm::erase_value(it->getSecond(), value);
+  llvm::erase_value(it->getSecond(), mapped);
   if (it->getSecond().empty())
-    mappings.reverse.erase(it);
+    mapping.erase(it);
 }
 
-void transform::TransformState::removePayloadOps(Value value) {
-  Mappings &mappings = getMapping(value);
-  for (Operation *op : mappings.direct[value])
-    dropReverseMapping(mappings, op, value);
-  mappings.direct.erase(value);
+void transform::TransformState::forgetMapping(Value opHandle,
+                                              ValueRange origOpFlatResults) {
+  Mappings &mappings = getMapping(opHandle);
+  for (Operation *op : mappings.direct[opHandle])
+    dropMappingEntry(mappings.reverse, op, opHandle);
+  mappings.direct.erase(opHandle);
+
+  for (Value opResult : origOpFlatResults) {
+    SmallVector<Value> resultHandles;
+    (void)getHandlesForPayloadValue(opResult, resultHandles);
+    for (Value resultHandle : resultHandles) {
+      Mappings &localMappings = getMapping(resultHandle);
+      dropMappingEntry(localMappings.values, resultHandle, opResult);
+      dropMappingEntry(localMappings.reverseValues, opResult, resultHandle);
+    }
+  }
 }
 
-LogicalResult transform::TransformState::updatePayloadOps(
-    Value value, function_ref<Operation *(Operation *)> callback) {
-  Mappings &mappings = getMapping(value);
-  auto it = mappings.direct.find(value);
-  assert(it != mappings.direct.end() && "unknown handle");
-  SmallVector<Operation *, 2> &association = it->getSecond();
-  SmallVector<Operation *, 2> updated;
-  updated.reserve(association.size());
+void transform::TransformState::forgetValueMapping(
+    Value valueHandle, ArrayRef<Operation *> payloadOperations) {
+  Mappings &mappings = getMapping(valueHandle);
+  for (Value payloadValue : mappings.reverseValues[valueHandle])
+    dropMappingEntry(mappings.reverseValues, payloadValue, valueHandle);
+  mappings.values.erase(valueHandle);
+
+  for (Operation *payloadOp : payloadOperations) {
+    SmallVector<Value> opHandles;
+    (void)getHandlesForPayloadOp(payloadOp, opHandles);
+    for (Value opHandle : opHandles) {
+      Mappings &localMappings = getMapping(opHandle);
+      dropMappingEntry(localMappings.direct, opHandle, payloadOp);
+      dropMappingEntry(localMappings.reverse, payloadOp, opHandle);
+    }
+  }
+}
 
-  for (Operation *op : association) {
-    dropReverseMapping(mappings, op, value);
-    if (Operation *updatedOp = callback(op)) {
-      updated.push_back(updatedOp);
-      mappings.reverse[updatedOp].push_back(value);
+LogicalResult
+transform::TransformState::replacePayloadOp(Operation *op,
+                                            Operation *replacement) {
+  // Drop the mapping between the op and all handles that point to it. Don't
+  // care if there are on such handles.
+  SmallVector<Value> opHandles;
+  (void)getHandlesForPayloadOp(op, opHandles);
+  for (Value handle : opHandles) {
+    Mappings &mappings = getMapping(handle);
+    dropMappingEntry(mappings.reverse, op, handle);
+  }
+
+  // Drop the mapping between the op results and all value handles that point to
+  // them. Don't care if there are no such handles.
+  RaggedArray<Value> resultValueHandles;
+  for (Value opResult : op->getResults()) {
+    SmallVector<Value> valueHandles;
+    (void)getHandlesForPayloadValue(opResult, valueHandles);
+    for (Value handle : valueHandles) {
+      Mappings &localMappings = getMapping(handle);
+      dropMappingEntry(localMappings.reverseValues, opResult, handle);
+    }
+    resultValueHandles.push_back(std::move(valueHandles));
+  }
+
+  // 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);
+    }
+    for (Value opResult : op->getResults()) {
+      SmallVector<Value> valueHandles;
+      (void)getHandlesForPayloadValue(opResult, valueHandles);
+      for (Value handle : valueHandles) {
+        Mappings &localMappings = getMapping(handle);
+        dropMappingEntry(localMappings.values, handle, opResult);
       }
     }
+    return success();
+  }
 
-  auto iface = value.getType().cast<TransformHandleTypeInterface>();
-  DiagnosedSilenceableFailure result =
-      iface.checkPayload(value.getLoc(), updated);
-  if (failed(result.checkAndReport()))
-    return failure();
+  // Otherwise, replace the pointed-to object of all handles while preserving
+  // their relative order.
+  if (op->getNumResults() != replacement->getNumResults()) {
+    return emitError(op->getLoc())
+           << "cannot replace an op with another op producing a different "
+              "number of results while tracking handles";
+  }
+
+  // Replace the mapped operation if present.
+  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;
+    }
+    mappings.reverse[replacement].push_back(handle);
+  }
+
+  // Replace the mapped results of the operation.
+  for (auto [origResult, replacementResult, handleList] : llvm::zip(
+           op->getResults(), replacement->getResults(), resultValueHandles)) {
+    for (Value resultHandle : handleList) {
+      Mappings &mappings = getMapping(resultHandle);
+      auto it = mappings.values.find(resultHandle);
+      if (it == mappings.values.end())
+        continue;
+
+      SmallVector<Value> &association = it->getSecond();
+      for (Value &mapped : association) {
+        if (mapped == origResult)
+          mapped = replacementResult;
+      }
+      mappings.reverseValues[replacementResult].push_back(resultHandle);
+    }
+  }
 
-  it->second = updated;
   return success();
 }
 
-void transform::TransformState::recordHandleInvalidationOne(
-    OpOperand &handle, Operation *payloadOp, Value otherHandle) {
-  ArrayRef<Operation *> potentialAncestors = getPayloadOps(handle.get());
+void transform::TransformState::recordOpHandleInvalidationOne(
+    OpOperand &consumingHandle, ArrayRef<Operation *> potentialAncestors,
+    Operation *payloadOp, Value otherHandle, Value throughValue) {
   // If the op is associated with invalidated handle, skip the check as it
   // may be reading invalid IR.
   if (invalidatedHandles.count(otherHandle))
     return;
 
   for (Operation *ancestor : potentialAncestors) {
     if (!ancestor->isAncestor(payloadOp))
       continue;
 
     // Make sure the error-reporting lambda doesn't capture anything
     // by-reference because it will go out of scope. Additionally, extract
     // location from Payload IR ops because the ops themselves may be
     // deleted before the lambda gets called.
     Location ancestorLoc = ancestor->getLoc();
     Location opLoc = payloadOp->getLoc();
-    Operation *owner = handle.getOwner();
-    unsigned operandNo = handle.getOperandNumber();
+    Operation *owner = consumingHandle.getOwner();
+    unsigned operandNo = consumingHandle.getOperandNumber();
+    std::optional<Location> throughValueLoc =
+        throughValue ? std::make_optional(throughValue.getLoc()) : std::nullopt;
     invalidatedHandles[otherHandle] = [ancestorLoc, opLoc, owner, operandNo,
-                                       otherHandle](Location currentLoc) {
+                                       otherHandle,
+                                       throughValueLoc](Location currentLoc) {
       InFlightDiagnostic diag = emitError(currentLoc)
                                 << "op uses a handle invalidated by a "
                                    "previously executed transform op";
       diag.attachNote(otherHandle.getLoc()) << "handle to invalidated ops";
       diag.attachNote(owner->getLoc())
           << "invalidated by this transform op that consumes its operand #"
           << operandNo
-          << " and invalidates handles to payload ops nested in payload "
-             "ops associated with the consumed handle";
+          << " and invalidates all handles to payload IR entities associated "
+             "with this operand and entities nested in them";
       diag.attachNote(ancestorLoc) << "ancestor payload op";
       diag.attachNote(opLoc) << "nested payload op";
+      if (throughValueLoc) {
+        diag.attachNote(*throughValueLoc)
+            << "consumed handle points to this payload value";
+      }
+    };
+  }
+}
+
+void transform::TransformState::recordValueHandleInvalidationByOpHandleOne(
+    OpOperand &consumingHandle, ArrayRef<Operation *> potentialAncestors,
+    Value payloadValue, Value valueHandle) {
+  // If the op is associated with invalidated handle, skip the check as it
+  // may be reading invalid IR.
+  if (invalidatedHandles.count(valueHandle))
+    return;
+
+  for (Operation *ancestor : potentialAncestors) {
+    Operation *definingOp;
+    std::optional<unsigned> resultNo = std::nullopt;
+    unsigned argumentNo, blockNo, regionNo;
+    if (auto opResult = payloadValue.dyn_cast<OpResult>()) {
+      definingOp = opResult.getOwner();
+      resultNo = opResult.getResultNumber();
+    } else {
+      auto arg = payloadValue.cast<BlockArgument>();
+      definingOp = arg.getParentBlock()->getParentOp();
+      argumentNo = arg.getArgNumber();
+      blockNo = std::distance(arg.getOwner()->getParent()->begin(),
+                              arg.getOwner()->getIterator());
+      regionNo = arg.getOwner()->getParent()->getRegionNumber();
+    }
+    assert(definingOp && "expected the value to be defined by an op as result "
+                         "or block argument");
+    if (!ancestor->isAncestor(definingOp))
+      continue;
+
+    Operation *owner = consumingHandle.getOwner();
+    unsigned operandNo = consumingHandle.getOperandNumber();
+    Location ancestorLoc = ancestor->getLoc();
+    Location opLoc = definingOp->getLoc();
+    Location valueLoc = payloadValue.getLoc();
+    invalidatedHandles[valueHandle] =
+        [valueHandle, owner, operandNo, resultNo, argumentNo, blockNo, regionNo,
+         ancestorLoc, opLoc, valueLoc](Location currentLoc) {
+          InFlightDiagnostic diag = emitError(currentLoc)
+                                    << "op uses a handle invalidated by a "
+                                       "previously executed transform op";
+          diag.attachNote(valueHandle.getLoc()) << "invalidated handle";
+          diag.attachNote(owner->getLoc())
+              << "invalidated by this transform op that consumes its operand #"
+              << operandNo
+              << " and invalidates all handles to payload IR entities "
+                 "associated with this operand and entities nested in them";
+          diag.attachNote(ancestorLoc)
+              << "ancestor op associated with the consumed handle";
+          if (resultNo) {
+            diag.attachNote(opLoc)
+                << "op defining the value as result #" << *resultNo;
+          } else {
+            diag.attachNote(opLoc)
+                << "op defining the value as block argument #" << argumentNo
+                << " of block #" << blockNo << " in region #" << regionNo;
+          }
+          diag.attachNote(valueLoc) << "payload value";
         };
   }
 }
 
-void transform::TransformState::recordHandleInvalidation(OpOperand &handle) {
-  for (const Mappings &mapping : llvm::make_second_range(mappings))
-    for (const auto &[payloadOp, otherHandles] : mapping.reverse)
+void transform::TransformState::recordOpHandleInvalidation(
+    OpOperand &handle, ArrayRef<Operation *> potentialAncestors,
+    Value throughValue) {
+  // Iterate over the mapping and invalidate aliasing handles. This is quite
+  // expensive and only necessary for error reporting in case of transform
+  // dialect misuse with dangling handles. Iteration over the handles is based
+  // on the assumption that the number of handles is significantly less than the
+  // number of IR objects (operations and values). Alternatively, we could walk
+  // the IR nested in each payload op associated with the given handle and look
+  // for handles associated with each operation and value.
+  for (const Mappings &mapping : llvm::make_second_range(mappings)) {
+    // Go over all op handle mappings and mark as invalidated any handle
+    // pointing to any of the payload ops associated with the given handle or
+    // any op nested in them.
+    for (const auto &[payloadOp, otherHandles] : mapping.reverse) {
       for (Value otherHandle : otherHandles)
-        recordHandleInvalidationOne(handle, payloadOp, otherHandle);
+        recordOpHandleInvalidationOne(handle, potentialAncestors, payloadOp,
+                                      otherHandle, throughValue);
+    }
+    // Go over all value handle mappings and mark as invalidated any handle
+    // pointing to any result of the payload op associated with the given handle
+    // or any op nested in them. Similarly invalidate handles to argument of
+    // blocks belonging to any region of any payload op associated with the
+    // given handle or any op nested in them.
+    for (const auto &[payloadValue, valueHandles] : mapping.reverseValues) {
+      for (Value valueHandle : valueHandles)
+        recordValueHandleInvalidationByOpHandleOne(handle, potentialAncestors,
+                                                   payloadValue, valueHandle);
+    }
+  }
+}
+
+void transform::TransformState::recordValueHandleInvalidation(
+    OpOperand &valueHandle) {
+  // Invalidate other handles to the same value.
+  for (Value payloadValue : getPayloadValues(valueHandle.get())) {
+    SmallVector<Value> otherValueHandles;
+    (void)getHandlesForPayloadValue(payloadValue, otherValueHandles);
+    for (Value otherHandle : otherValueHandles) {
+      Operation *owner = valueHandle.getOwner();
+      unsigned operandNo = valueHandle.getOperandNumber();
+      Location valueLoc = payloadValue.getLoc();
+      invalidatedHandles[otherHandle] = [otherHandle, owner, operandNo,
+                                         valueLoc](Location currentLoc) {
+        InFlightDiagnostic diag = emitError(currentLoc)
+                                  << "op uses a handle invalidated by a "
+                                     "previously executed transform op";
+        diag.attachNote(otherHandle.getLoc()) << "invalidated handle";
+        diag.attachNote(owner->getLoc())
+            << "invalidated by this transform op that consumes its operand #"
+            << operandNo
+            << " and invalidates handles to the same values as associated with "
+               "it";
+        diag.attachNote(valueLoc) << "payload value";
+      };
+    }
+
+    if (auto opResult = payloadValue.dyn_cast<OpResult>()) {
+      Operation *payloadOp = opResult.getOwner();
+      recordOpHandleInvalidation(valueHandle, payloadOp, payloadValue);
+    } else {
+      auto arg = payloadValue.dyn_cast<BlockArgument>();
+      for (Operation &payloadOp : *arg.getOwner())
+        recordOpHandleInvalidation(valueHandle, &payloadOp, payloadValue);
+    }
+  }
 }
 
 LogicalResult transform::TransformState::checkAndRecordHandleInvalidation(
     TransformOpInterface transform) {
   auto memoryEffectsIface =
       cast<MemoryEffectOpInterface>(transform.getOperation());
   SmallVector<MemoryEffects::EffectInstance> effects;
   memoryEffectsIface.getEffectsOnResource(
       transform::TransformMappingResource::get(), effects);
 
   for (OpOperand &target : transform->getOpOperands()) {
     // If the operand uses an invalidated handle, report it.
     auto it = invalidatedHandles.find(target.get());
     if (!transform.allowsRepeatedHandleOperands() &&
         it != invalidatedHandles.end())
       return it->getSecond()(transform->getLoc()), failure();
 
     // Invalidate handles pointing to the operations nested in the operation
     // associated with the handle consumed by this operation.
     auto consumesTarget = [&](const MemoryEffects::EffectInstance &effect) {
       return isa<MemoryEffects::Free>(effect.getEffect()) &&
              effect.getValue() == target.get();
     };
-    if (llvm::any_of(effects, consumesTarget))
-      recordHandleInvalidation(target);
+    if (llvm::any_of(effects, consumesTarget)) {
+      if (target.get().getType().isa<TransformHandleTypeInterface>()) {
+        ArrayRef<Operation *> payloadOps = getPayloadOps(target.get());
+        recordOpHandleInvalidation(target, payloadOps);
+      } else if (target.get()
+                     .getType()
+                     .isa<TransformValueHandleTypeInterface>()) {
+        recordValueHandleInvalidation(target);
+      }
+    }
   }
 
   return success();
 }
 
+template <typename T>
+DiagnosedSilenceableFailure
+checkRepeatedConsumptionInOperand(ArrayRef<T> payload,
+                                  transform::TransformOpInterface transform,
+                                  unsigned operandNumber) {
+  DenseSet<T> seen;
+  for (T p : payload) {
+    if (!seen.insert(p).second) {
+      DiagnosedSilenceableFailure diag =
+          transform.emitSilenceableError()
+          << "a handle passed as operand #" << operandNumber
+          << " and consumed by this operation points to a payload "
+             "entity more than once";
+      if constexpr (std::is_pointer_v<T>)
+        diag.attachNote(p->getLoc()) << "repeated target op";
+      else
+        diag.attachNote(p.getLoc()) << "repeated target value";
+      return diag;
+    }
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG(DBGS() << "applying: " << transform << "\n");
   auto printOnFailureRAII = llvm::make_scope_exit([this] {
     (void)this;
     DEBUG_WITH_TYPE(DEBUG_PRINT_AFTER_ALL, {
       DBGS() << "Top-level payload:\n";
       getTopLevel()->print(llvm::dbgs(),
                            mlir::OpPrintingFlags().printGenericOpForm());
     });
   });
   if (options.getExpensiveChecksEnabled()) {
     if (failed(checkAndRecordHandleInvalidation(transform)))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     for (OpOperand &operand : transform->getOpOperands()) {
       if (!isHandleConsumed(operand.get(), transform))
         continue;
 
-      DenseSet<Operation *> seen;
-      for (Operation *op : getPayloadOps(operand.get())) {
-        if (!seen.insert(op).second) {
-          DiagnosedSilenceableFailure diag =
-              transform.emitSilenceableError()
-              << "a handle passed as operand #" << operand.getOperandNumber()
-              << " and consumed by this operation points to a payload "
-                 "operation more than once";
-          diag.attachNote(op->getLoc()) << "repeated target op";
-          return diag;
-        }
-      }
-    }
+      Type operandType = operand.get().getType();
+      if (operandType.isa<TransformHandleTypeInterface>()) {
+        DiagnosedSilenceableFailure check =
+            checkRepeatedConsumptionInOperand<Operation *>(
+                getPayloadOps(operand.get()), transform,
+                operand.getOperandNumber());
+        if (!check.succeeded())
+          return check;
+      } else if (operandType.isa<TransformValueHandleTypeInterface>()) {
+        DiagnosedSilenceableFailure check =
+            checkRepeatedConsumptionInOperand<Value>(
+                getPayloadValues(operand.get()), transform,
+                operand.getOperandNumber());
+        if (!check.succeeded())
+          return check;
+
+  // Find which operands are consumed.
+  DenseSet<unsigned> consumedOperands;
+  auto memEffectInterface =
+      cast<MemoryEffectOpInterface>(transform.getOperation());
+  SmallVector<MemoryEffects::EffectInstance, 2> effects;
+  for (OpOperand &target : transform->getOpOperands()) {
+    effects.clear();
+    memEffectInterface.getEffectsOnValue(target.get(), effects);
+    if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
+          return isa<transform::TransformMappingResource>(
+                     effect.getResource()) &&
+                 isa<MemoryEffects::Free>(effect.getEffect());
+        })) {
+      consumedOperands.insert(target.getOperandNumber());
+    }
+  }
+
+  // Remember the results of the payload ops associated with the consumed
+  // op handles or the ops defining the value handles so we can drop the
+  // association with them later. This must happen here because the
+  // transformation may destroy or mutate them so we cannot traverse the payload
+  // IR after that.
+  SmallVector<Value> origOpFlatResults;
+  SmallVector<Operation *> origAssociatedOps;
+  for (unsigned index : consumedOperands) {
+    Value operand = transform->getOperand(index);
+    if (operand.getType().isa<TransformHandleTypeInterface>()) {
+      for (Operation *payloadOp : getPayloadOps(operand))
+        llvm::append_range(origOpFlatResults, payloadOp->getResults());
+      continue;
+    }
+    if (operand.getType().isa<TransformValueHandleTypeInterface>()) {
+      for (Value payloadValue : getPayloadValues(operand)) {
+        if (payloadValue.isa<OpResult>()) {
+          origAssociatedOps.push_back(payloadValue.getDefiningOp());
+          continue;
+        }
+        llvm::append_range(
+            origAssociatedOps,
+            llvm::map_range(*payloadValue.cast<BlockArgument>().getOwner(),
+                            [](Operation &op) { return &op; }));
+      }
+      continue;
+    }
+    DiagnosedDefiniteFailure diag =
+        emitDefiniteFailure(transform->getLoc())
+        << "unexpectedly consumed a value that is not a handle as operand #"
+        << index;
+    diag.attachNote(operand.getLoc())
+        << "value defined here with type " << operand.getType();
+    return diag;
   }
 
-  transform::TransformResults results(transform->getNumResults());
   // 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));
   if (result.isDefiniteFailure())
     return result;
 
   // If a silenceable failure was produced, some results may be unset, set them
   // to empty lists.
   if (result.isSilenceableFailure()) {
     for (OpResult opResult : transform->getResults()) {
       if (results.isSet(opResult.getResultNumber()))
         continue;
 
       if (opResult.getType().isa<TransformParamTypeInterface>())
         results.setParams(opResult, {});
       else
         results.set(opResult, {});
     }
   }
 
   // Remove the mapping for the operand if it is consumed by the operation. This
   // allows us to catch use-after-free with assertions later on.
-  auto memEffectInterface =
-      cast<MemoryEffectOpInterface>(transform.getOperation());
-  SmallVector<MemoryEffects::EffectInstance, 2> effects;
-  for (OpOperand &target : transform->getOpOperands()) {
-    effects.clear();
-    memEffectInterface.getEffectsOnValue(target.get(), effects);
-    if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
-          return isa<transform::TransformMappingResource>(
-                     effect.getResource()) &&
-                 isa<MemoryEffects::Free>(effect.getEffect());
-        })) {
-      removePayloadOps(target.get());
+  for (unsigned index : consumedOperands) {
+    Value operand = transform->getOperand(index);
+    if (operand.getType().isa<TransformHandleTypeInterface>()) {
+      forgetMapping(operand, origOpFlatResults);
+    } else if (operand.getType().isa<TransformValueHandleTypeInterface>()) {
+      forgetValueMapping(operand, origAssociatedOps);
     }
   }
 
   for (OpResult result : transform->getResults()) {
     assert(result.getDefiningOp() == transform.getOperation() &&
            "payload IR association for a value other than the result of the "
            "current transform op");
     if (result.getType().isa<TransformParamTypeInterface>()) {
       assert(results.isParam(result.getResultNumber()) &&
              "expected parameters for the parameter-typed result");
       if (failed(
               setParams(result, results.getParams(result.getResultNumber())))) {
         return DiagnosedSilenceableFailure::definiteFailure();
       }
+    } else if (result.getType().isa<TransformValueHandleTypeInterface>()) {
+      assert(results.isValue(result.getResultNumber()) &&
+             "expected values for value-type-result");
+      if (failed(setPayloadValues(
+              result, results.getValues(result.getResultNumber())))) {
+        return DiagnosedSilenceableFailure::definiteFailure();
+      }
     } else {
       assert(!results.isParam(result.getResultNumber()) &&
              "expected payload ops for the non-parameter typed result");
       if (failed(
               setPayloadOps(result, results.get(result.getResultNumber())))) {
         return DiagnosedSilenceableFailure::definiteFailure();
       }
     }

 
 LogicalResult
 transform::TransformState::Extension::replacePayloadOp(Operation *op,
                                                        Operation *replacement) {
   SmallVector<Value> handles;
   if (failed(state.getHandlesForPayloadOp(op, handles)))
     return failure();
 
-  for (Value handle : handles) {
-    LogicalResult result =
-        state.updatePayloadOps(handle, [&](Operation *current) {
-          return current == op ? replacement : current;
-        });
-    if (failed(result))
-      return failure();
-  }
-  return success();
+  // TODO: we may need to invalidate handles to operations and values nested in
+  // the operation being replaced.
+  return state.replacePayloadOp(op, replacement);
 }
 
 //===----------------------------------------------------------------------===//
 // TransformResults
 //===----------------------------------------------------------------------===//
 
 transform::TransformResults::TransformResults(unsigned numSegments) {
-  segments.resize(numSegments,
-                  ArrayRef<Operation *>(nullptr, static_cast<size_t>(0)));
-  paramSegments.resize(numSegments, ArrayRef<TransformState::Param>(
-                                        nullptr, static_cast<size_t>(0)));
+  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>(segments.size()) &&
+  assert(position < static_cast<int64_t>(operations.size()) &&
          "setting results for a non-existent handle");
-  assert(segments[position].data() == nullptr && "results already set");
-  int64_t start = operations.size();
-  llvm::append_range(operations, ops);
-  segments[position] = ArrayRef(operations).drop_front(start);
+  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>(paramSegments.size()) &&
+  assert(position < static_cast<int64_t>(this->params.size()) &&
          "setting params for a non-existent handle");
-  assert(paramSegments[position].data() == nullptr && "params already set");
-  size_t start = this->params.size();
-  llvm::append_range(this->params, params);
-  paramSegments[position] = ArrayRef(this->params).drop_front(start);
+  assert(this->params[position].data() == nullptr && "params already set");
+  assert(operations[position].data() == nullptr &&
+         "another kind of results already set");
+  assert(values[position].data() == nullptr &&
+         "another kind of results already set");
+  this->params.replace(position, params);
+}
+
+void transform::TransformResults::setValues(OpResult handle,
+                                            ValueRange values) {
+  int64_t position = handle.getResultNumber();
+  assert(position < static_cast<int64_t>(values.size()) &&
+         "setting values for a non-existent handle");
+  assert(this->values[position].data() == nullptr && "values already set");
+  assert(operations[position].data() == nullptr &&
+         "another kind of results already set");
+  assert(params[position].data() == nullptr &&
+         "another kind of results already set");
+  this->values.replace(position, values);
 }
 
 ArrayRef<Operation *>
 transform::TransformResults::get(unsigned resultNumber) const {
-  assert(resultNumber < segments.size() &&
+  assert(resultNumber < operations.size() &&
          "querying results for a non-existent handle");
-  assert(segments[resultNumber].data() != nullptr &&
-         "querying unset results (param expected?)");
-  return segments[resultNumber];
+  assert(operations[resultNumber].data() != nullptr &&
+         "querying unset results (values or params expected?)");
+  return operations[resultNumber];
 }
 
 ArrayRef<transform::TransformState::Param>
 transform::TransformResults::getParams(unsigned resultNumber) const {
-  assert(resultNumber < paramSegments.size() &&
+  assert(resultNumber < params.size() &&
+         "querying params for a non-existent handle");
+  assert(params[resultNumber].data() != nullptr &&
+         "querying unset params (ops or values expected?)");
+  return params[resultNumber];
+}
+
+ArrayRef<Value>
+transform::TransformResults::getValues(unsigned resultNumber) const {
+  assert(resultNumber < params.size() &&
          "querying params for a non-existent handle");
-  assert(paramSegments[resultNumber].data() != nullptr &&
-         "querying unset params (payload ops expected?)");
-  return paramSegments[resultNumber];
+  assert(values[resultNumber].data() != nullptr &&
+         "querying unset values (ops or params expected?)");
+  return values[resultNumber];
 }
 
 bool transform::TransformResults::isParam(unsigned resultNumber) const {
-  assert(resultNumber < paramSegments.size() &&
+  assert(resultNumber < params.size() &&
          "querying association for a non-existent handle");
-  return paramSegments[resultNumber].data() != nullptr;
+  return params[resultNumber].data() != nullptr;
+}
+
+bool transform::TransformResults::isValue(unsigned resultNumber) const {
+  assert(resultNumber < values.size() &&
+         "querying association for a non-existent handle");
+  return values[resultNumber].data() != nullptr;
 }
 
 bool transform::TransformResults::isSet(unsigned resultNumber) const {
-  assert(resultNumber < paramSegments.size() &&
+  assert(resultNumber < params.size() &&
          "querying association for a non-existent handle");
-  return paramSegments[resultNumber].data() != nullptr ||
-         segments[resultNumber].data() != nullptr;
+  return params[resultNumber].data() != nullptr ||
+         operations[resultNumber].data() != nullptr ||
+         values[resultNumber].data() != nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for TransformEachOpTrait.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 transform::detail::checkApplyToOne(Operation *transformOp,
void transform::detail::setApplyToOneResults(
     ArrayRef<ApplyToEachResultList> results) {
   for (OpResult r : transformOp->getResults()) {
     if (r.getType().isa<TransformParamTypeInterface>()) {
       auto params = llvm::to_vector(
           llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
             return oneResult[r.getResultNumber()].get<Attribute>();
           }));
       transformResults.setParams(r, params);
+    } else if (r.getType().isa<TransformValueHandleTypeInterface>()) {
+      auto values = llvm::to_vector(
+          llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
+            return oneResult[r.getResultNumber()].get<Value>();
+          }));
+      transformResults.setValues(r, values);
     } else {
       auto payloads = llvm::to_vector(
           llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
             return oneResult[r.getResultNumber()].get<Operation *>();
           }));
       transformResults.set(r, payloads);
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 
 LogicalResult transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
     TransformState &state, Operation *op, Region &region) {
   SmallVector<Operation *> targets;
   SmallVector<SmallVector<MappedValue>> extraMappings;
   if (op->getNumOperands() != 0) {
     llvm::append_range(targets, state.getPayloadOps(op->getOperand(0)));
     for (Value operand : op->getOperands().drop_front()) {
       SmallVector<MappedValue> &mapped = extraMappings.emplace_back();
       if (operand.getType().isa<TransformHandleTypeInterface>()) {
         llvm::append_range(mapped, state.getPayloadOps(operand));
+      } else if (operand.getType().isa<TransformValueHandleTypeInterface>()) {
+        llvm::append_range(mapped, state.getPayloadValues(operand));
       } else {
         assert(operand.getType().isa<TransformParamTypeInterface>() &&
                "unsupported kind of transform dialect value");
         llvm::append_range(mapped, state.getParams(operand));
       }
     }
   } else {
     if (state.getNumTopLevelMappings() !=
if (op->getNumOperands() != 0) {
     if (arg.getType() != op->getOperand(0).getType()) {
       return op->emitOpError()
              << "expects the type of the block argument to match "
                 "the type of the operand";
     }
   }
   for (BlockArgument arg : body->getArguments().drop_front()) {
     if (arg.getType()
-            .isa<TransformHandleTypeInterface, TransformParamTypeInterface>())
+            .isa<TransformHandleTypeInterface, TransformParamTypeInterface,
+                 TransformValueHandleTypeInterface>())
       continue;
 
     InFlightDiagnostic diag =
         op->emitOpError()
         << "expects trailing entry block arguments to be of type implementing "
-           "TransformHandleTypeInterface or TransformParamTypeInterface";
+           "TransformHandleTypeInterface, TransformValueHandleTypeInterface or "
+           "TransformParamTypeInterface";
     diag.attachNote() << "argument #" << arg.getArgNumber() << " does not";
     return diag;
   }
 
   if (auto *parent =
           op->getParentWithTrait<PossibleTopLevelTransformOpTrait>()) {
     if (op->getNumOperands() != body->getNumArguments()) {
       InFlightDiagnostic diag =

 //===----------------------------------------------------------------------===//
 
 void transform::detail::getParamProducerTransformOpTraitEffects(
     Operation *op, SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   producesHandle(op->getResults(), effects);
   bool hasPayloadOperands = false;
   for (Value operand : op->getOperands()) {
     onlyReadsHandle(operand, effects);
-    if (operand.getType().isa<TransformHandleTypeInterface>())
+    if (operand.getType()
+            .isa<TransformHandleTypeInterface,
+                 TransformValueHandleTypeInterface>())
       hasPayloadOperands = true;
   }
   if (hasPayloadOperands)
     onlyReadsPayload(effects);
 }
 
 LogicalResult
 transform::detail::verifyParamProducerTransformOpTrait(Operation *op) {

 
 //===----------------------------------------------------------------------===//
 // Entry point.
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 transform::applyTransforms(Operation *payloadRoot,
                            TransformOpInterface transform,
-                           ArrayRef<ArrayRef<MappedValue>> extraMapping,
+                           const RaggedArray<MappedValue> &extraMapping,
                            const TransformOptions &options) {
 #ifndef NDEBUG
   if (!transform->hasTrait<PossibleTopLevelTransformOpTrait>() ||
       transform->getNumOperands() != 0) {
     transform->emitError()
         << "expected transform to start at the top-level transform op";
     llvm::report_fatal_error("could not run transforms",
                              /*gen_crash_diag=*/false);

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

if (integerAttr.getType() != getType()) {
       return emitSilenceableError(loc)
              << "expected the type of the parameter attribute ("
              << integerAttr.getType() << ") to match the parameter type ("
              << getType() << ")";
     }
   }
   return DiagnosedSilenceableFailure::success();
 }
+
+//===----------------------------------------------------------------------===//
+// transform::AnyValueType
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::AnyValueType::checkPayload(Location loc,
+                                      ArrayRef<Value> payload) const {
+  return DiagnosedSilenceableFailure::success();
+}

mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp

DEBUG_WITH_TYPE(DEBUG_TYPE_DUMP_FILE, {
     saveReproToTempFile(llvm::dbgs(), target, transform, passName,
                         debugPayloadRootTag, debugTransformRootTag, binaryName);
   });
 }
 
 LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
     Operation *target, StringRef passName,
     const std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
-    ArrayRef<ArrayRef<MappedValue>> extraMappings,
+    const RaggedArray<MappedValue> &extraMappings,
     const TransformOptions &options,
     const Pass::Option<std::string> &transformFileName,
     const Pass::Option<std::string> &debugPayloadRootTag,
     const Pass::Option<std::string> &debugTransformRootTag,
     StringRef binaryName) {
 
   // Step 1
   // ------

mlir/test/Dialect/Linalg/transform-op-match.mlir

 // RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
 
 func.func @bar() {
   // expected-remark @below {{matched op name}}
   // expected-remark @below {{matched attr name}}
   %0 = arith.constant {my_attr} 0: i32
   // expected-remark @below {{matched op name}}
   %1 = arith.constant 1 : i32
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %match_name = transform.structured.match ops{["arith.constant"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.test_print_remark_at_operand %match_name, "matched op name" : !pdl.operation
-  transform.test_consume_operand %match_name
+  transform.test_consume_operand %match_name : !pdl.operation
 
   %match_attr = transform.structured.match ops{["arith.constant"]} attributes{my_attr} in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.test_print_remark_at_operand %match_attr, "matched attr name" : !pdl.operation
-  transform.test_consume_operand %match_attr
+  transform.test_consume_operand %match_attr : !pdl.operation
 }
 
 // -----
 
 func.func @by_type() {
   %0 = arith.constant 0: i32
   // expected-remark @below {{matched op name}}
   %1 = arith.constant 1.0 : f32
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %match_name = transform.structured.match
     ops{["arith.constant"]} filter_result_type = f32 in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.test_print_remark_at_operand %match_name, "matched op name" : !pdl.operation
-  transform.test_consume_operand %match_name
+  transform.test_consume_operand %match_name : !pdl.operation
 }
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 func.func @match_complex_attribute(%arg0: tensor<12x128x32xf32>)
     -> tensor<128x12x32xf32> {
^bb1(%arg1: !pdl.operation):
   %match_attr = transform.structured.match
       ops{["linalg.generic"]}
       attributes{iterator_types = [
         #linalg.iterator_type<parallel>,
         #linalg.iterator_type<parallel>,
         #linalg.iterator_type<parallel>]}
       in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.test_print_remark_at_operand %match_attr, "matched complex attr" : !pdl.operation
-  transform.test_consume_operand %match_attr
+  transform.test_consume_operand %match_attr : !pdl.operation
 
   %no_match = transform.structured.match
       attributes{iterator_types = [
         #linalg.iterator_type<parallel>,
         #linalg.iterator_type<parallel>,
         #linalg.iterator_type<reduction>]}
       in %arg1 : (!pdl.operation) -> !pdl.operation
 // expected-remark @below {{0}}
   transform.test_print_number_of_associated_payload_ir_ops %no_match
 }

mlir/test/Dialect/Transform/check-use-after-free.mlir

 // RUN: mlir-opt %s --transform-dialect-check-uses --split-input-file --verify-diagnostics
 
 func.func @use_after_free_branching_control_flow() {
   // expected-note @below {{allocated here}}
-  %0 = transform.test_produce_param_or_forward_operand 42
+  %0 = transform.test_produce_self_handle_or_forward_operand
   transform.test_transform_op_with_regions {
     "transform.test_branching_transform_op_terminator"() : () -> ()
   },
   {
   ^bb0:
     "transform.test_branching_transform_op_terminator"()[^bb1, ^bb2] : () -> ()
   ^bb1:
     // expected-note @below {{freed here}}
-    transform.test_consume_operand_if_matches_param_or_fail %0[42]
+    transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
     "transform.test_branching_transform_op_terminator"()[^bb3] : () -> ()
   ^bb2:
     "transform.test_branching_transform_op_terminator"()[^bb3] : () -> ()
   ^bb3:
     // expected-warning @below {{operand #0 may be used after free}}
     transform.sequence %0 : !pdl.operation failures(propagate) {
     ^bb0(%arg0: !pdl.operation):
     }
     "transform.test_branching_transform_op_terminator"() : () -> ()
   }
   return
 }
 
 // -----
 
 func.func @use_after_free_in_nested_op() {
   // expected-note @below {{allocated here}}
-  %0 = transform.test_produce_param_or_forward_operand 42
+  %0 = transform.test_produce_self_handle_or_forward_operand
   // expected-note @below {{freed here}}
   transform.test_transform_op_with_regions {
     "transform.test_branching_transform_op_terminator"() : () -> ()
   },
   {
   ^bb0:
     "transform.test_branching_transform_op_terminator"()[^bb1, ^bb2] : () -> ()
   ^bb1:
-    transform.test_consume_operand_if_matches_param_or_fail %0[42]
+    transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
     "transform.test_branching_transform_op_terminator"()[^bb3] : () -> ()
   ^bb2:
     "transform.test_branching_transform_op_terminator"()[^bb3] : () -> ()
   ^bb3:
     "transform.test_branching_transform_op_terminator"() : () -> ()
   }
   // expected-warning @below {{operand #0 may be used after free}}
   transform.sequence %0 : !pdl.operation failures(propagate) {
^bb0(%arg0: !pdl.operation):
     ^bb3(%arg3: !pdl.operation):
     }
 
     // `transform.sequence` has recursive side effects so it has the same "free"
     // as the child op it contains.
     // expected-note @below {{freed here}}
     transform.sequence %0 : !pdl.operation failures(propagate) attributes { ord = 4 } {
     ^bb4(%arg4: !pdl.operation):
-      test_consume_operand_if_matches_param_or_fail %0[42]
+      test_consume_operand_of_op_kind_or_fail %0, "transform.sequence"
     }
     // expected-warning @below {{operand #0 may be used after free}}
     transform.sequence %0 : !pdl.operation failures(propagate) attributes { ord = 5 } {
     ^bb3(%arg3: !pdl.operation):
     }
   }
   return
 }
^bb0(%arg0: !pdl.operation):
     transform.sequence %0 : !pdl.operation failures(propagate) attributes { ord = 2 } {
     ^bb2(%arg2: !pdl.operation):
     }
     transform.sequence %0 : !pdl.operation failures(propagate) attributes { ord = 3 } {
     ^bb3(%arg3: !pdl.operation):
     }
 
     // expected-note @below {{freed here}}
-    test_consume_operand_if_matches_param_or_fail %0[42]
+    test_consume_operand_of_op_kind_or_fail %0, "transform.sequence"
     // expected-warning @below {{operand #0 may be used after free}}
     transform.sequence %0 : !pdl.operation failures(propagate) attributes { ord = 5 } {
     ^bb3(%arg3: !pdl.operation):
     }
   }
   return
 }
 
 // -----
 
 // In the case of a control flow cycle, the operation that uses the value may
 // precede the one that frees it in the same block. Both operations should
 // be reported as use-after-free.
 func.func @use_after_free_self_cycle() {
   // expected-note @below {{allocated here}}
-  %0 = transform.test_produce_param_or_forward_operand 42
+  %0 = transform.test_produce_self_handle_or_forward_operand
   transform.test_transform_op_with_regions {
     "transform.test_branching_transform_op_terminator"() : () -> ()
   },
   {
   ^bb0:
     "transform.test_branching_transform_op_terminator"()[^bb1] : () -> ()
   ^bb1:
     // expected-warning @below {{operand #0 may be used after free}}
     transform.sequence %0 : !pdl.operation failures(propagate) {
     ^bb0(%arg0: !pdl.operation):
     }
     // expected-warning @below {{operand #0 may be used after free}}
     // expected-note @below {{freed here}}
-    transform.test_consume_operand_if_matches_param_or_fail %0[42]
+    transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
     "transform.test_branching_transform_op_terminator"()[^bb1, ^bb2] : () -> ()
   ^bb2:
     "transform.test_branching_transform_op_terminator"() : () -> ()
   }
   return
 }
 
 
 // -----
 
 // Check that the "free" that happens in a cycle is also reported as potential
 // use-after-free.
 func.func @use_after_free_cycle() {
   // expected-note @below {{allocated here}}
-  %0 = transform.test_produce_param_or_forward_operand 42
+  %0 = transform.test_produce_self_handle_or_forward_operand
   transform.test_transform_op_with_regions {
     "transform.test_branching_transform_op_terminator"() : () -> ()
   },
   {
   ^bb0:
     "transform.test_branching_transform_op_terminator"()[^bb1, ^bb2] : () -> ()
   ^bb1:
     // expected-warning @below {{operand #0 may be used after free}}
     // expected-note @below {{freed here}}
-    transform.test_consume_operand_if_matches_param_or_fail %0[42]
+    transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
     "transform.test_branching_transform_op_terminator"()[^bb2, ^bb3] : () -> ()
   ^bb2:
     "transform.test_branching_transform_op_terminator"()[^bb1] : () -> ()
   ^bb3:
     "transform.test_branching_transform_op_terminator"() : () -> ()
   }
   return
 }

mlir/test/Dialect/Transform/expensive-checks.mlir

^bb0(%arg0: !pdl.operation):
   }
 
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     // expected-note @below {{handle to invalidated ops}}
     %0 = pdl_match @return in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
     // expected-note @below {{invalidated by this transform op that consumes its operand #0}}
-    test_consume_operand %1
+    test_consume_operand %1 : !pdl.operation
     // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
     test_print_remark_at_operand %0, "remark" : !pdl.operation
   }
 }
 
 // -----
 
 func.func @func1() {
pdl.pattern @return : benefit(1) {
     rewrite %2 with "transform.dialect"
   }
 
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @func in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = pdl_match @return in %arg1 : (!pdl.operation) -> !pdl.operation
     %2 = replicate num(%0) %1 : !pdl.operation, !pdl.operation
-    // expected-error @below {{a handle passed as operand #0 and consumed by this operation points to a payload operation more than once}}
-    test_consume_operand %2
+    // expected-error @below {{a handle passed as operand #0 and consumed by this operation points to a payload entity more than once}}
+    test_consume_operand %2 : !pdl.operation
     test_print_remark_at_operand %0, "remark" : !pdl.operation
   }
 }
 
 
 // -----
 
 // expected-note @below {{ancestor payload op}}
 // expected-note @below {{nested payload op}}
 module {
 
   transform.sequence failures(propagate) {
   ^bb0(%0: !pdl.operation):
     %1 = transform.test_copy_payload %0
     // expected-note @below {{handle to invalidated ops}}
     %2 = transform.test_copy_payload %0
     // expected-note @below {{invalidated by this transform op that consumes its operand #0}}
-    transform.test_consume_operand %1
+    transform.test_consume_operand %1 : !pdl.operation
     // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
-    transform.test_consume_operand %2
+    transform.test_consume_operand %2 : !pdl.operation
   }
 }
 
 // -----
 
 // expected-note @below {{ancestor payload op}}
 // expected-note @below {{nested payload op}}
 module {
 
   transform.sequence failures(propagate) {
   ^bb0(%0: !pdl.operation):
     %1 = transform.test_copy_payload %0
     // expected-note @below {{handle to invalidated ops}}
     %2 = transform.test_copy_payload %0
     // Consuming two handles in the same operation is invalid if they point
     // to overlapping sets of payload IR ops.
     //
     // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
-    // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates handles}}
-    transform.test_consume_operand %1, %2
+    // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities}}
+    transform.test_consume_operand %1, %2 : !pdl.operation
   }
 }
 
 // -----
 
 // Deduplication attribute allows "merge_handles" to take repeated operands.
 
 module {
 
   transform.sequence failures(propagate) {
   ^bb0(%0: !pdl.operation):
     %1 = transform.test_copy_payload %0
     %2 = transform.test_copy_payload %0
     transform.merge_handles %1, %2 { deduplicate } : !pdl.operation
   }
 }
+// -----
+
+// expected-note @below {{payload value}}
+%0 = "test.match_anchor"() : () -> (i32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %2 = transform.structured.match ops{["test.match_anchor"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_result %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated handle}}
+  %4 = test_produce_value_handle_to_result %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates handles to the same values as associated with it}}
+  test_consume_operand %3 : !transform.any_value
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %4 : !transform.any_value
+}
+
+// -----
+
+// expected-note @below {{ancestor op associated with the consumed handle}}
+// expected-note @below {{payload value}}
+// expected-note @below {{op defining the value as result #0}}
+%0 = "test.match_anchor"() : () -> (i32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %2 = transform.structured.match ops{["test.match_anchor"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{invalidated handle}}
+  %3 = test_produce_value_handle_to_result %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %2 : !transform.any_op
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %3 : !transform.any_value
+}
+
+// -----
+
+// expected-note @below {{ancestor op associated with the consumed handle}}
+"test.match_anchor_1"() ({
+^bb0:
+  // expected-note @below {{op defining the value as result #0}}
+  // expected-note @below {{payload value}}
+  %0 = "test.match_anchor_2"() : () -> (i32)
+  "test.region_terminator"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{invalidated handle}}
+  %3 = test_produce_value_handle_to_result %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %1 : !transform.any_op
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %3 : !transform.any_value
+}
+
+// -----
+
+// expected-note @below {{ancestor op associated with the consumed handle}}
+// expected-note @below {{op defining the value as block argument #0 of block #0 in region #0}}
+"test.match_anchor_1"() ({
+// expected-note @below {{payload value}}
+^bb0(%arg0: i32):
+  %0 = "test.match_anchor_2"() : () -> (i32)
+  "test.region_terminator"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{invalidated handle}}
+  %3 = test_produce_value_handle_to_argument_of_parent_block %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %1 : !transform.any_op
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %3 : !transform.any_value
+}
+
+// -----
+
+// expected-note @below {{ancestor op associated with the consumed handle}}
+"test.match_anchor_1"() ({
+^bb:
+  // expected-note @below {{op defining the value as block argument #0 of block #0 in region #0}}
+  "test.op_with_regions"() ({
+  // expected-note @below {{payload value}}
+  ^bb0(%arg0: i32):
+    %0 = "test.match_anchor_2"() : () -> (i32)
+    "test.region_terminator"() : () -> ()
+  }): () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{invalidated handle}}
+  %3 = test_produce_value_handle_to_argument_of_parent_block %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %1 : !transform.any_op
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %3 : !transform.any_value
+}
+
+// -----
+
+// expected-note @below {{ancestor payload op}}
+// expected-note @below {{nested payload op}}
+// expected-note @below {{consumed handle points to this payload value}}
+%0 = "test.match_anchor"() : () -> (i32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  // expected-note @below {{handle to invalidated ops}}
+  %2 = transform.structured.match ops{["test.match_anchor"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_result %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %3 : !transform.any_value
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %2 : !transform.any_op 
+}
+
+// -----
+
+// expected-note @below {{ancestor payload op}}
+// expected-note @below {{consumed handle points to this payload value}}
+%0 = "test.match_anchor_1"() ({
+^bb0:
+  // expected-note @below {{nested payload op}}
+  "test.match_anchor_2"() : () -> ()
+  "test.region_terminator"() : () -> ()
+}) : () -> (i32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{handle to invalidated ops}}
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_result %1, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %3 : !transform.any_value
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %2 : !transform.any_op
+}
+
+
+// -----
+
+"test.match_anchor_1"() ({
+// expected-note @below {{consumed handle points to this payload value}}
+^bb0(%arg0: f32):
+  // expected-note @below {{ancestor payload op}}
+  // expected-note @below {{nested payload op}}
+  "test.match_anchor_2"() : () -> ()
+  "test.region_terminator"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  // expected-note @below {{handle to invalidated ops}}
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_argument_of_parent_block %2, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %3 : !transform.any_value
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %2 : !transform.any_op
+}
+
+// -----
+
+"test.op_with_regions"() ({
+// expected-note @below {{consumed handle points to this payload value}}
+^bb(%arg0: i32):
+  // expected-note @below {{ancestor payload op}}
+  "test.op_with_regions"() ({
+  ^bb0:
+    // expected-note @below {{nested payload op}}
+    "test.match_anchor_2"() : () -> ()
+    "test.region_terminator"() : () -> ()
+  }): () -> ()
+  "test.match_anchor_1"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-note @below {{handle to invalidated ops}}
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_argument_of_parent_block %1, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
+  test_consume_operand %3 : !transform.any_value
+  // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
+  test_consume_operand %2 : !transform.any_op
+}
+
+// -----
+
+// Removing a block argument does not invalidate handles to operations in another block.
+// Not expecting an error here.
+
+"test.op_with_regions"() ({
+^bb1(%arg0: i32):
+  "test.match_anchor_1"() : () -> ()
+^bb2:
+  "test.match_anchor_2"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_argument_of_parent_block %1, 0 : (!transform.any_op) -> !transform.any_value
+  test_consume_operand %3 : !transform.any_value
+  test_consume_operand %2 : !transform.any_op
+}

mlir/test/Dialect/Transform/multi-arg-top-level-ops.mlir

 }
 
 func.func @foo() {
   return
 }
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_value):
+  // expected-error @above {{wrong kind of value provided for the top-level value handle}}
+}
+
+func.func @foo() {
+  return
+}
+
+// -----
+
 // expected-error @below {{operation expects 1 extra value bindings, but 2 were provided to the interpreter}}
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op):
 }
 
 // -----
 
 transform.sequence failures(propagate) {

mlir/test/Dialect/Transform/multi-arg-top-level-values.mlir

+// RUN: mlir-opt %s --pass-pipeline='builtin.module(test-transform-dialect-interpreter{bind-first-extra-to-results-of-ops=test.some_returning_op bind-second-extra-to-results-of-ops=test.some_other_returning_op})' \
+// RUN:             --split-input-file --verify-diagnostics
+
+// Note that diagnostic checker will merge two diagnostics with the same message
+// at the same location, so only check the remark once.
+// 
+// expected-remark @below {{first extra}}
+// expected-note @below {{value handle points to an op result #0}}
+// expected-note @below {{value handle points to an op result #1}}
+%0:2 = "test.some_returning_op"() : () -> (i32, i64)
+
+// expected-remark @below {{first extra}}
+// expected-note @below {{value handle points to an op result #0}}
+%1 = "test.some_returning_op"() : () -> index
+
+// Note that diagnostic checker will merge two diagnostics with the same message
+// at the same location, so only check the remark once.
+// 
+// expected-remark @below {{second extra}}
+// expected-note @below {{value handle points to an op result #0}}
+// expected-note @below {{value handle points to an op result #1}}
+%2:2 = "test.some_other_returning_op"() : () -> (f32, f64)
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_value, %arg2: !transform.any_value):
+  test_print_remark_at_operand_value %arg1, "first extra" : !transform.any_value
+  test_print_remark_at_operand_value %arg2, "second extra" : !transform.any_value
+}
+
+// -----
+
+%0:2 = "test.some_returning_op"() : () -> (i32, i64)
+%1 = "test.some_returning_op"() : () -> index
+
+transform.sequence failures(propagate) {
+// expected-error @below {{wrong kind of value provided for top-level operation handle}}
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_op, %arg2: !transform.any_value):
+}
+
+// -----
+
+// expected-error @below {{operation expects 1 extra value bindings, but 2 were provided to the interpreter}}
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op, %arg1: !transform.any_value):
+}

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

^bb0(%arg0: !pdl.operation):
   // expected-error @below {{expects operands to be provided for a nested op}}
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
   }
 }
 
 // -----
 
-// expected-error @below {{'transform.sequence' op expects trailing entry block arguments to be of type implementing TransformHandleTypeInterface or TransformParamTypeInterface}}
+// expected-error @below {{'transform.sequence' op expects trailing entry block arguments to be of type implementing TransformHandleTypeInterface, TransformValueHandleTypeInterface or TransformParamTypeInterface}}
 // expected-note @below {{argument #1 does not}}
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op, %arg1: i64):
 }
 
 // -----
 
 // expected-error @below {{expected children ops to implement TransformOpInterface}}
^bb1:
   "test.potential_terminator"() : () -> ()
 }) : () -> ()
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{result #0 has more than one potential consumer}}
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   // expected-note @below {{used here as operand #0}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
   // expected-note @below {{used here as operand #0}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{result #0 has more than one potential consumer}}
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   // expected-note @below {{used here as operand #0}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
   // expected-note @below {{used here as operand #0}}
   transform.sequence %0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
-    test_consume_operand_if_matches_param_or_fail %arg1[42]
+    test_consume_operand_of_op_kind_or_fail %arg1, "transform.test_produce_self_handle_or_forward_operand"
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{result #0 has more than one potential consumer}}
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   // expected-note @below {{used here as operand #0}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
   transform.sequence %0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     // expected-note @below {{used here as operand #0}}
-    test_consume_operand_if_matches_param_or_fail %0[42]
+    test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{result #0 has more than one potential consumer}}
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   // expected-note @below {{used here as operand #0}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
   // expected-note @below {{used here as operand #0}}
   transform.sequence %0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     transform.sequence %arg1 : !pdl.operation failures(propagate) {
     ^bb2(%arg2: !pdl.operation):
-      test_consume_operand_if_matches_param_or_fail %arg2[42]
+      test_consume_operand_of_op_kind_or_fail %arg2, "transform.test_produce_self_handle_or_forward_operand"
     }
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
^bb0:
   transform.yield
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{result #0 has more than one potential consumer}}
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   // expected-note @below {{used here as operand #0}}
   transform.foreach %0 : !pdl.operation {
   ^bb1(%arg1: !pdl.operation):
-    transform.test_consume_operand %arg1
+    transform.test_consume_operand %arg1 : !pdl.operation
   }
   // expected-note @below {{used here as operand #0}}
-  transform.test_consume_operand %0
+  transform.test_consume_operand %0 : !pdl.operation
 }
 
 // -----
 
 transform.sequence failures(suppress) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{TransformOpInterface requires memory effects on operands to be specified}}
   // expected-note @below {{no effects specified for operand #0}}

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

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

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

 // RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-remark @below {{applying transformation}}
   transform.test_transform_op
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
-  %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
+  %0 = transform.test_produce_self_handle_or_forward_operand { foo = "bar" }
   // expected-remark @below {{succeeded}}
-  transform.test_consume_operand_if_matches_param_or_fail %0[42]
+  transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
-  %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
-  // expected-error @below {{expected the operand to be associated with 21 got 42}}
-  transform.test_consume_operand_if_matches_param_or_fail %0[21]
+  %0 = transform.test_produce_self_handle_or_forward_operand { foo = "bar" }
+  // expected-error @below {{expected the operand to be associated a payload op of kind transform.sequence got transform.test_produce_self_handle_or_forward_operand}}
+  transform.test_consume_operand_of_op_kind_or_fail %0, "transform.sequence"
 }
 
 // -----
 
 // It is okay to have multiple handles to the same payload op as long
 // as only one of them is consumed. The expensive checks mode is necessary
 // to detect double-consumption.
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
-  %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
+  %0 = transform.test_produce_self_handle_or_forward_operand { foo = "bar" }
   %1 = transform.test_copy_payload %0
   // expected-remark @below {{succeeded}}
-  transform.test_consume_operand_if_matches_param_or_fail %0[42]
+  transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
^bb0(%arg0: !pdl.operation):
   // expected-remark @below {{applying transformation "e"}}
   test_transform_op "e"
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
-  %0 = test_produce_param_or_forward_operand 42
+  %0 = test_produce_self_handle_or_forward_operand
   sequence %0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     // expected-remark @below {{succeeded}}
-    test_consume_operand_if_matches_param_or_fail %arg1[42]
+    test_consume_operand_of_op_kind_or_fail %arg1, "transform.test_produce_self_handle_or_forward_operand"
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   %0 = sequence %arg0 : !pdl.operation -> !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
-    %1 = test_produce_param_or_forward_operand 42
+    %1 = test_produce_self_handle_or_forward_operand
     yield %1 : !pdl.operation
   }
   // expected-remark @below {{succeeded}}
-  test_consume_operand_if_matches_param_or_fail %0[42]
+  test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
 }
 
 // -----
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
^bb0(%arg0: !pdl.operation):
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     // This is necessary to run the transformation on something other than the
     // top-level module, "alternatives" cannot be run on that.
     %0 = pdl_match @match_func in %arg1 : (!pdl.operation) -> !pdl.operation
     transform.alternatives %0 : !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
-      %1 = transform.test_produce_param_or_forward_operand 42
+      %1 = transform.test_produce_self_handle_or_forward_operand
       // This operation fails, which triggers the next alternative without
       // reporting the error.
-      transform.test_consume_operand_if_matches_param_or_fail %1[43]
+      transform.test_consume_operand_of_op_kind_or_fail %1, "transform.sequence"
     }, {
     ^bb2(%arg2: !pdl.operation):
-      %1 = transform.test_produce_param_or_forward_operand 42
+      %1 = transform.test_produce_self_handle_or_forward_operand
       // expected-remark @below {{succeeded}}
-      transform.test_consume_operand_if_matches_param_or_fail %1[42]
+      transform.test_consume_operand_of_op_kind_or_fail %1, "transform.test_produce_self_handle_or_forward_operand"
     }
   }
 }
 
 // -----
 
 func.func private @bar()
 
^bb0(%arg0: !pdl.operation):
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @match_call in %arg1 : (!pdl.operation) -> !pdl.operation
     %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
     %2 = transform.alternatives %1 : !pdl.operation -> !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
       %3 = transform.pdl_match @match_call in %arg2 : (!pdl.operation) -> !pdl.operation
       // expected-remark @below {{applying}}
       transform.test_emit_remark_and_erase_operand %3, "applying" {fail_after_erase}
-      %4 = transform.test_produce_param_or_forward_operand 43
+      %4 = transform.test_produce_self_handle_or_forward_operand %3
       transform.yield %4 : !pdl.operation
     }, {
     ^bb2(%arg2: !pdl.operation):
-      %4 = transform.test_produce_param_or_forward_operand 42
+      %4 = transform.test_produce_self_handle_or_forward_operand
       transform.yield %4 : !pdl.operation
     }
     // The first alternative failed, so the returned value is taken from the
-    // second alternative.
+    // second alternative, associated test_produce_self_handle_or_forward_operand rather
+    // than pdl_match.
     // expected-remark @below {{succeeded}}
-    transform.test_consume_operand_if_matches_param_or_fail %2[42]
+    transform.test_consume_operand_of_op_kind_or_fail %2, "transform.test_produce_self_handle_or_forward_operand"
   }
 }
 
 // -----
 
 // expected-note @below {{scope}}
 module {
   func.func @foo() {
     %0 = arith.constant 0 : i32
     return
   }
 
   func.func @bar() {
     %0 = arith.constant 0 : i32
     %1 = arith.constant 1 : i32
     return
   }
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     // expected-error @below {{scope must not contain the transforms being applied}}
     transform.alternatives %arg1 : !pdl.operation {
     ^bb2(%arg2: !pdl.operation):
-      %0 = transform.test_produce_param_or_forward_operand 42
-      transform.test_consume_operand_if_matches_param_or_fail %0[43]
+      %0 = transform.test_produce_self_handle_or_forward_operand
+      transform.test_consume_operand_of_op_kind_or_fail %0, "transform.sequence"
     }, {
     ^bb2(%arg2: !pdl.operation):
-      %0 = transform.test_produce_param_or_forward_operand 42
-      transform.test_consume_operand_if_matches_param_or_fail %0[42]
+      %0 = transform.test_produce_self_handle_or_forward_operand
+      transform.test_consume_operand_of_op_kind_or_fail %0, "transform.test_produce_self_handle_or_forward_operand"
     }
   }
 }
 
 // -----
 
 func.func @foo(%arg0: index, %arg1: index, %arg2: index) {
   // expected-note @below {{scope}}

 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{attempting to assign a null parameter to this transform value}}
   %0 = transform.test_produce_null_param : !transform.param<i64>
 }
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{attempting to assign a null payload value to this transform handle}}
+  %0 = transform.test_produce_null_value : !transform.any_value
+}
+
+// -----
+
 // expected-error @below {{could not find a nested top-level transform op}}
 // expected-note @below {{use the 'transform-file-name' option to provide transform as external file}}
 module {
 }
 
 // -----
 
 // expected-note @below {{previous top-level transform op}}
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
 }
 
-// expected-error @below {{ore than one top-level transform op}}
+// expected-error @below {{more than one top-level transform op}}
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+// expected-remark @below {{value handle}}
+// expected-note @below {{value handle points to a block argument #0 in block #0 in region #0}}
+^bb1(%arg0: !transform.any_op):
+  %0 = test_produce_value_handle_to_self_operand %arg0 : (!transform.any_op) -> !transform.any_value
+  test_print_remark_at_operand_value %0, "value handle" : !transform.any_value
+}
+
+// -----
+
+// expected-remark @below {{result handle}}
+// expected-note @below {{value handle points to an op result #1}}
+%0:2 = "test.get_two_results"() : () -> (i32, i32)
+// expected-remark @below {{result handle}}
+// expected-note @below {{value handle points to an op result #1}}
+%1:3 = "test.get_three_results"() : () -> (i32, i32, f32)
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %2 = transform.structured.match ops{["test.get_two_results", "test.get_three_results"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_result %2, 1 : (!transform.any_op) -> !transform.any_value
+  test_print_remark_at_operand_value %3, "result handle" : !transform.any_value
+}
+
+// -----
+
+"test.op_with_regions"() ({
+^bb0:
+  "test.regon_terminator"() : () -> ()
+}, {
+^bb1:
+  "test.regon_terminator"() : () -> ()
+// expected-remark @below {{block argument handle}}
+// expected-note @below {{value handle points to a block argument #2 in block #1 in region #1}}
+^bb2(%arg0: i32, %arg1: f64, %arg3: index):
+  "test.match_anchor"() : () -> ()
+  "test.regon_terminator"() : () -> ()
+}) : () -> ()
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %2 = transform.structured.match ops{["test.match_anchor"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %3 = test_produce_value_handle_to_argument_of_parent_block %2, 2 : (!transform.any_op) -> !transform.any_value
+  test_print_remark_at_operand_value %3, "block argument handle" : !transform.any_value
+}
+
+// -----
+
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
+  // expected-note @below {{value defined here with type '!transform.test_dialect_param'}}
+  %0 = test_produce_param_with_number_of_test_ops %arg0 : !transform.any_op
+  // expected-error @below {{unexpectedly consumed a value that is not a handle as operand #0}}
+  test_consume_operand %0 : !transform.test_dialect_param
 }

mlir/test/Dialect/Transform/transform-state-extension.mlir

^bb0(%arg0: !pdl.operation):
     test_remap_operand_to_self %arg0
     // expected-remark @below {{extension present, A}}
     test_check_if_test_extension_present %arg0
   }
 }
 
 // -----
 
+// expected-error @below {{cannot replace an op with another op producing a different number of results while tracking handles}}
+module {
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !pdl.operation):
+    test_add_test_extension "A"
+    %dummy = test_remap_operand_to_self %arg0 : !transform.any_op
+  }
+}
+
+
+// -----
+
 module {
   transform.sequence failures(suppress) {
   ^bb0(%arg0: !pdl.operation):
     // expected-error @below {{TestTransformStateExtension missing}}
     test_remap_operand_to_self %arg0
   }
 }

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

public:
   }
 
   // No side effects.
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 };
 } // namespace
 
 DiagnosedSilenceableFailure
-mlir::test::TestProduceParamOrForwardOperandOp::apply(
+mlir::test::TestProduceSelfHandleOrForwardOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   if (getOperation()->getNumOperands() != 0) {
     results.set(getResult().cast<OpResult>(),
                 getOperation()->getOperand(0).getDefiningOp());
   } else {
-    results.set(getResult().cast<OpResult>(),
-                reinterpret_cast<Operation *>(*getParameter()));
+    results.set(getResult().cast<OpResult>(), getOperation());
   }
   return DiagnosedSilenceableFailure::success();
 }
 
-void mlir::test::TestProduceParamOrForwardOperandOp::getEffects(
+void mlir::test::TestProduceSelfHandleOrForwardOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   if (getOperand())
     transform::onlyReadsHandle(getOperand(), effects);
   transform::producesHandle(getRes(), effects);
 }
 
-LogicalResult mlir::test::TestProduceParamOrForwardOperandOp::verify() {
-  if (getParameter().has_value() ^ (getNumOperands() != 1))
-    return emitOpError() << "expects either a parameter or an operand";
-  return success();
+DiagnosedSilenceableFailure
+mlir::test::TestProduceValueHandleToSelfOperand::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  results.setValues(getOut().cast<OpResult>(), getIn());
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestProduceValueHandleToSelfOperand::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::producesHandle(getOut(), effects);
+  transform::onlyReadsPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestProduceValueHandleToResult::applyToOne(
+    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  if (target->getNumResults() <= getNumber())
+    return emitSilenceableError() << "payload has no result #" << getNumber();
+  results.push_back(target->getResult(getNumber()));
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestProduceValueHandleToResult::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::producesHandle(getOut(), effects);
+  transform::onlyReadsPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestProduceValueHandleToArgumentOfParentBlock::applyToOne(
+    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  if (!target->getBlock())
+    return emitSilenceableError() << "payload has no parent block";
+  if (target->getBlock()->getNumArguments() <= getNumber())
+    return emitSilenceableError()
+           << "parent of the payload has no argument #" << getNumber();
+  results.push_back(target->getBlock()->getArgument(getNumber()));
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestProduceValueHandleToArgumentOfParentBlock::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::producesHandle(getOut(), effects);
+  transform::onlyReadsPayload(effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestConsumeOperand::apply(transform::TransformResults &results,
                                       transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestConsumeOperand::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getOperand(), effects);
   if (getSecondOperand())
     transform::consumesHandle(getSecondOperand(), effects);
   transform::modifiesPayload(effects);
 }
 
-DiagnosedSilenceableFailure
-mlir::test::TestConsumeOperandIfMatchesParamOrFail::apply(
+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");
-  auto value = reinterpret_cast<intptr_t>(payload[0]);
-  if (static_cast<uint64_t>(value) != getParameter()) {
+  if (payload[0]->getName().getStringRef() != getOpKind()) {
     return emitSilenceableError()
-           << "op expected the operand to be associated with " << getParameter()
-           << " got " << value;
+           << "op expected the operand to be associated a payload op of kind "
+           << getOpKind() << " got " << payload[0]->getName().getStringRef();
   }
 
   emitRemark() << "succeeded";
   return DiagnosedSilenceableFailure::success();
 }
 
-void mlir::test::TestConsumeOperandIfMatchesParamOrFail::getEffects(
+void mlir::test::TestConsumeOperandOfOpKindOrFail::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::consumesHandle(getOperand(), effects);
   transform::modifiesPayload(effects);
 }
 
 DiagnosedSilenceableFailure mlir::test::TestPrintRemarkAtOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   ArrayRef<Operation *> payload = state.getPayloadOps(getOperand());
   for (Operation *op : payload)
     op->emitRemark() << getMessage();
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintRemarkAtOperandOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::onlyReadsPayload(effects);
 }
 
+DiagnosedSilenceableFailure mlir::test::TestPrintRemarkAtOperandValue::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  ArrayRef<Value> values = state.getPayloadValues(getIn());
+  for (Value value : values) {
+    std::string note;
+    llvm::raw_string_ostream os(note);
+    if (auto arg = value.dyn_cast<BlockArgument>()) {
+      os << "a block argument #" << arg.getArgNumber() << " in block #"
+         << std::distance(arg.getOwner()->getParent()->begin(),
+                          arg.getOwner()->getIterator())
+         << " in region #" << arg.getOwner()->getParent()->getRegionNumber();
+    } else {
+      os << "an op result #" << value.cast<OpResult>().getResultNumber();
+    }
+    InFlightDiagnostic diag = ::emitRemark(value.getLoc()) << getMessage();
+    diag.attachNote() << "value handle points to " << os.str();
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestPrintRemarkAtOperandValue::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::onlyReadsPayload(effects);
+}
+
 DiagnosedSilenceableFailure
 mlir::test::TestAddTestExtensionOp::apply(transform::TransformResults &results,
                                           transform::TransformState &state) {
   state.addExtension<TestTransformStateExtension>(getMessageAttr());
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
if (failed(extension->updateMapping(state.getPayloadOps(getOperand()).front(),
                                       getOperation())))
     return DiagnosedSilenceableFailure::definiteFailure();
   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::TestProduceNullParamOp::apply(transform::TransformResults &results,
                                           transform::TransformState &state) {
   results.setParams(getOut().cast<OpResult>(), Attribute());
   return DiagnosedSilenceableFailure::success();
 }
 
+void mlir::test::TestProduceNullValueOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::producesHandle(getOut(), effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestProduceNullValueOp::apply(transform::TransformResults &results,
+                                          transform::TransformState &state) {
+  results.setValues(getOut().cast<OpResult>(), Value());
+  return DiagnosedSilenceableFailure::success();
+}
+
 void mlir::test::TestRequiredMemoryEffectsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   if (getHasOperandEffect())
     transform::consumesHandle(getIn(), effects);
 
   if (getHasResultEffect())
     transform::producesHandle(getOut(), effects);
   else

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

: TypeDef<Transform_Dialect, "TestDialectParam",
       [DeclareTypeInterfaceMethods<TransformParamTypeInterface>]> {
   let description = [{
     Parameter associated with an i32 attribute for testing purposes.
   }];
   let mnemonic = "test_dialect_param";
   let assemblyFormat = "";
 }
 
-def TestProduceParamOrForwardOperandOp
-  : Op<Transform_Dialect, "test_produce_param_or_forward_operand",
+def TestProduceSelfHandleOrForwardOperandOp
+  : Op<Transform_Dialect, "test_produce_self_handle_or_forward_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let arguments = (ins
-     Optional<PDL_Operation>:$operand,
-     OptionalAttr<I64Attr>:$parameter);
+  let arguments = (ins Optional<PDL_Operation>:$operand);
   let results = (outs PDL_Operation:$res);
-  let assemblyFormat = "(`from` $operand^)? ($parameter^)? attr-dict";
+  let assemblyFormat = "($operand^)? attr-dict";
   let cppNamespace = "::mlir::test";
-  let hasVerifier = 1;
+}
+
+def TestProduceValueHandleToSelfOperand
+  : Op<Transform_Dialect, "test_produce_value_handle_to_self_operand",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$in);
+  let results = (outs TransformValueHandleTypeInterface:$out);
+  let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+  
+}
+
+def TestProduceValueHandleToResult
+  : Op<Transform_Dialect, "test_produce_value_handle_to_result",
+       [TransformEachOpTrait, TransformOpInterface,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$in,
+                       I64Attr:$number);
+  let results = (outs TransformValueHandleTypeInterface:$out);
+  let assemblyFormat = "$in `,` $number attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+    
+def TestProduceValueHandleToArgumentOfParentBlock
+  : Op<Transform_Dialect, "test_produce_value_handle_to_argument_of_parent_block",
+       [TransformEachOpTrait, TransformOpInterface,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$in,
+                       I64Attr:$number);
+  let results = (outs TransformValueHandleTypeInterface:$out);
+  let assemblyFormat = "$in `,` $number attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
 }
 
 def TestConsumeOperand : Op<Transform_Dialect, "test_consume_operand",
      [DeclareOpInterfaceMethods<TransformOpInterface>,
       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins
-    PDL_Operation:$operand,
+    Transform_AnyHandleOrParamType:$operand,
     Optional<PDL_Operation>:$second_operand);
-  let assemblyFormat = "$operand (`,` $second_operand^)? attr-dict";
+  let assemblyFormat = "$operand (`,` $second_operand^)? attr-dict `:` type($operand)";
   let cppNamespace = "::mlir::test";
 }
 
-def TestConsumeOperandIfMatchesParamOrFail
-  : Op<Transform_Dialect, "test_consume_operand_if_matches_param_or_fail",
+def TestConsumeOperandOfOpKindOrFail
+  : Op<Transform_Dialect, "test_consume_operand_of_op_kind_or_fail",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins
     PDL_Operation:$operand,
-    I64Attr:$parameter);
-  let assemblyFormat = "$operand `[` $parameter `]` attr-dict";
+    StrAttr:$op_kind);
+  let assemblyFormat = "$operand `,` $op_kind attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestPrintRemarkAtOperandOp
   : Op<Transform_Dialect, "test_print_remark_at_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins
     TransformHandleTypeInterface:$operand,
     StrAttr:$message);
   let assemblyFormat =
     "$operand `,` $message attr-dict `:` type($operand)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestPrintRemarkAtOperandValue
+  : Op<Transform_Dialect, "test_print_remark_at_operand_value",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins TransformValueHandleTypeInterface:$in,
+                       StrAttr:$message);
+  let assemblyFormat = "$in `,` $message attr-dict `:` type($in)";
+  let cppNamespace = "::mlir::test";
+}
+
 def TestAddTestExtensionOp
   : Op<Transform_Dialect, "test_add_test_extension",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         NoMemoryEffect]> {
   let arguments = (ins StrAttr:$message);
   let assemblyFormat = "$message attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestCheckIfTestExtensionPresentOp
   : Op<Transform_Dialect, "test_check_if_test_extension_present",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins PDL_Operation:$operand);
   let assemblyFormat = "$operand attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestRemapOperandPayloadToSelfOp
   : Op<Transform_Dialect, "test_remap_operand_to_self",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let arguments = (ins  PDL_Operation:$operand);
-  let assemblyFormat = "$operand attr-dict";
+  let arguments = (ins PDL_Operation:$operand);
+  let results = (outs Optional<TransformHandleTypeInterface>:$out);        
+  let assemblyFormat = "$operand attr-dict (`:` type($out)^)?";
   let cppNamespace = "::mlir::test";
 }
 
 def TestRemoveTestExtensionOp
   : Op<Transform_Dialect, "test_remove_test_extension",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         NoMemoryEffect]> {
   let assemblyFormat = "attr-dict";
def TestProduceNullParamOp
   : Op<Transform_Dialect, "test_produce_null_param",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let results = (outs TransformParamTypeInterface:$out);
   let assemblyFormat = "attr-dict `:` type($out)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestProduceNullValueOp
+  : Op<Transform_Dialect, "test_produce_null_value",
+       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let results = (outs TransformValueHandleTypeInterface:$out);
+  let assemblyFormat = "attr-dict `:` type($out)";
+  let cppNamespace = "::mlir::test";
+}
+
 def TestRequiredMemoryEffectsOp
   : Op<Transform_Dialect, "test_required_memory_effects",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   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";
 }
 
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD

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

public:
   StringRef getArgument() const override {
     return "test-transform-dialect-interpreter";
   }
 
   StringRef getDescription() const override {
     return "apply transform dialect operations one by one";
   }
 
-  ArrayRef<transform::MappedValue>
-  findOperationsByName(Operation *root, StringRef name,
-                       SmallVectorImpl<transform::MappedValue> &storage) {
-    size_t start = storage.size();
+  void findOperationsByName(Operation *root, StringRef name,
+                            SmallVectorImpl<Operation *> &operations) {
     root->walk([&](Operation *op) {
       if (op->getName().getStringRef() == name) {
-        storage.push_back(op);
+        operations.push_back(op);
       }
     });
-    return ArrayRef(storage).drop_front(start);
   }
 
-  ArrayRef<transform::MappedValue>
-  createParameterMapping(MLIRContext &context, ArrayRef<int> values,
-                         SmallVectorImpl<transform::MappedValue> &storage) {
-    size_t start = storage.size();
-    llvm::append_range(storage, llvm::map_range(values, [&](int v) {
-                         Builder b(&context);
-                         return transform::MappedValue(b.getI64IntegerAttr(v));
-                       }));
-    return ArrayRef(storage).drop_front(start);
+  void createParameterMapping(MLIRContext &context, ArrayRef<int> values,
+                              RaggedArray<transform::MappedValue> &result) {
+    SmallVector<transform::MappedValue> storage =
+        llvm::to_vector(llvm::map_range(values, [&](int v) {
+    result.push_back(std::move(storage));
+  }
+
+  void
+  createOpResultMapping(Operation *root, StringRef name,
+                        RaggedArray<transform::MappedValue> &extraMapping) {
+    SmallVector<Operation *> operations;
+    findOperationsByName(root, name, operations);
+    SmallVector<Value> results;
+    for (Operation *op : operations)
+      llvm::append_range(results, op->getResults());
+    extraMapping.push_back(results);
+  }
+
+  unsigned numberOfSetOptions(const Option<std::string> &ops,
+                              const ListOption<int> &params,
+                              const Option<std::string> &values) {
+    unsigned numSetValues = 0;
+    numSetValues += !ops.empty();
+    numSetValues += !params.empty();
+    numSetValues += !values.empty();
+    return numSetValues;
   }
 
   void runOnOperation() override {
-    if (!bindFirstExtraToOps.empty() && !bindFirstExtraToParams.empty()) {
-      emitError(UnknownLoc::get(&getContext()))
-          << "cannot bind the first extra top-level argument to both "
-             "operations and parameters";
+    unsigned firstSetOptions =
+        numberOfSetOptions(bindFirstExtraToOps, bindFirstExtraToParams,
+                           bindFirstExtraToResultsOfOps);
+    unsigned secondSetOptions =
+        numberOfSetOptions(bindSecondExtraToOps, bindSecondExtraToParams,
+                           bindSecondExtraToResultsOfOps);
+    auto loc = UnknownLoc::get(&getContext());
+    if (firstSetOptions > 1) {
+      emitError(loc) << "cannot bind the first extra top-level argument to "
+                        "multiple entities";
       return signalPassFailure();
     }
-    if (!bindSecondExtraToOps.empty() && !bindSecondExtraToParams.empty()) {
-      emitError(UnknownLoc::get(&getContext()))
-          << "cannot bind the second extra top-level argument to both "
-             "operations and parameters";
+    if (secondSetOptions > 1) {
+      emitError(loc) << "cannot bind the second extra top-level argument to "
+                        "multiple entities";
       return signalPassFailure();
     }
-    if ((!bindSecondExtraToOps.empty() || !bindSecondExtraToParams.empty()) &&
-        bindFirstExtraToOps.empty() && bindFirstExtraToParams.empty()) {
-      emitError(UnknownLoc::get(&getContext()))
-          << "cannot bind the second extra top-level argument without binding "
-             "the first";
-      return signalPassFailure();
+    if (firstSetOptions == 0 && secondSetOptions != 0) {
+      emitError(loc) << "cannot bind the second extra top-level argument "
+                        "without bindings the first";
     }
 
-    SmallVector<transform::MappedValue> extraMappingStorage;
-    SmallVector<ArrayRef<transform::MappedValue>> extraMapping;
+    RaggedArray<transform::MappedValue> extraMapping;
     if (!bindFirstExtraToOps.empty()) {
-      extraMapping.push_back(findOperationsByName(
-          getOperation(), bindFirstExtraToOps.getValue(), extraMappingStorage));
+      SmallVector<Operation *> operations;
+      findOperationsByName(getOperation(), bindFirstExtraToOps.getValue(),
+                           operations);
+      extraMapping.push_back(operations);
     } else if (!bindFirstExtraToParams.empty()) {
-      extraMapping.push_back(createParameterMapping(
-          getContext(), bindFirstExtraToParams, extraMappingStorage));
+      createParameterMapping(getContext(), bindFirstExtraToParams,
+                             extraMapping);
+    } else if (!bindFirstExtraToResultsOfOps.empty()) {
+      createOpResultMapping(getOperation(), bindFirstExtraToResultsOfOps,
+                            extraMapping);
     }
+
     if (!bindSecondExtraToOps.empty()) {
-      extraMapping.push_back(findOperationsByName(
-          getOperation(), bindSecondExtraToOps, extraMappingStorage));
+      SmallVector<Operation *> operations;
+      findOperationsByName(getOperation(), bindSecondExtraToOps, operations);
+      extraMapping.push_back(operations);
     } else if (!bindSecondExtraToParams.empty()) {
-      extraMapping.push_back(createParameterMapping(
-          getContext(), bindSecondExtraToParams, extraMappingStorage));
+      createParameterMapping(getContext(), bindSecondExtraToParams,
+                             extraMapping);
+    } else if (!bindSecondExtraToResultsOfOps.empty()) {
+      createOpResultMapping(getOperation(), bindSecondExtraToResultsOfOps,
+                            extraMapping);
     }
 
     options = options.enableExpensiveChecks(enableExpensiveChecks);
     if (failed(transform::detail::interpreterBaseRunOnOperationImpl(
             getOperation(), getArgument(), getSharedTransformModule(),
             extraMapping, options, transformFileName, debugPayloadRootTag,
             debugTransformRootTag, getBinaryName())))
       return signalPassFailure();
   }
 
   Option<bool> enableExpensiveChecks{
       *this, "enable-expensive-checks", llvm::cl::init(false),
       llvm::cl::desc("perform expensive checks to better report errors in the "
                      "transform IR")};
 
   Option<std::string> bindFirstExtraToOps{
       *this, "bind-first-extra-to-ops",
       llvm::cl::desc("bind the first extra argument of the top-level op to "
                      "payload operations of the given kind")};
   ListOption<int> bindFirstExtraToParams{
       *this, "bind-first-extra-to-params",
       llvm::cl::desc("bind the first extra argument of the top-level op to "
                      "the given integer parameters")};
+  Option<std::string> bindFirstExtraToResultsOfOps{
+      *this, "bind-first-extra-to-results-of-ops",
+      llvm::cl::desc("bind the first extra argument of the top-level op to "
+                     "results of payload operations of the given kind")};
 
   Option<std::string> bindSecondExtraToOps{
       *this, "bind-second-extra-to-ops",
       llvm::cl::desc("bind the second extra argument of the top-level op to "
                      "payload operations of the given kind")};
   ListOption<int> bindSecondExtraToParams{
       *this, "bind-second-extra-to-params",
       llvm::cl::desc("bind the second extra argument of the top-level op to "
                      "the given integer parameters")};
+  Option<std::string> bindSecondExtraToResultsOfOps{
+      *this, "bind-second-extra-to-results-of-ops",
+      llvm::cl::desc("bind the second extra argument of the top-level op to "
+                     "results of payload operations of the given kind")};
+
   Option<std::string> transformFileName{
       *this, "transform-file-name", llvm::cl::init(""),
       llvm::cl::desc(
           "Optional filename containing a transform dialect specification to "
           "apply. If left empty, the IR is assumed to contain one top-level "
           "transform dialect operation somewhere in the module.")};
   Option<std::string> debugPayloadRootTag{
       *this, "debug-payload-root-tag", llvm::cl::init(""),