[mlir] adapt TransformEachOpTrait to parameter values

[ftynse]

Summary

Adapt the implementation of TransformEachOpTrait to the existence of
parameter values recently introduced into the transform dialect. In
particular, allow applyToOne hooks to return a list containing a mix
of Operation * that will be associated with handles and Attribute
that will be associated with parameter values by the trait
implementation of the transform interface's apply method.

Disentangle the "transposition" of the list of per-payload op partial
results to decrease its overall complexity and detemplatize the code
that doesn't really need templates. This removes the poorly documented
special handling for single-result ops with TransformEachOpTrait that
could have assigned null pointer values to handles.

[springerm]

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

544–545

can have arbitrary number of input params?

545–546

operation handles and params?

738

How does this work? Attribute is not a pointer.

769

Why is this useful? Also, is this adding an nullptr operation* or a null attribute?

793

Does this class own the attributes or just pointers to them? Could there be an object lifetime issue?

798

When is null useful as a result?

835

new ops and params.

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

485

get asserts. Do you want dyn_cast here? Alternative: params.push_back(...get<Attribute>())

493

get asserts. Do you want dyn_cast here? Alternative: payloads.push_back(...get<Operation *>())

[ftynse]

Address review.

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

738

Fundamentally, it is a pointer. That's why we are passing it by-value everywhere. https://github.com/llvm/llvm-project/blob/eebfee8f9ea7dfc9e7cf4326decadc779969994e/mlir/include/mlir/IR/Attributes.h#L290-L300

769

Good point, it looks useless now that I finished the refactoring.

793

Attributes are owned by the MLIRContext, as usual. We usually consider them immortal.

798

I don't think they are, but this was an allowed behavior before so I prefer to change it separately.

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

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

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

let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$blockDim,
                    DefaultValuedAttr<BoolAttr, "true">:$syncAfterDistribute);
   let results = (outs PDL_Operation:$result);
 
   let assemblyFormat = "$target attr-dict";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 
 def MapForeachToBlocks :
   Op<Transform_Dialect, "gpu.map_foreach_to_blocks",
     [FunctionalStyleTransformOpTrait,
      MemoryEffectsOpInterface,
      TransformOpInterface,
      TransformEachOpTrait]> {
   let description = [{
     Target the gpu_launch op and rewrite the top level `scf.foreach_thread`
     to distributed gpu.block_id attribute. If `generate_gpu_launch` attribute
     is set, then first generates `gpu_launch` and moves the top level
     `scf.foreach_thread` inside.
 
     The operation searches top level `scf.foreach_thread` ops under
     `gpu_launch` and maps each such op to GPU blocks. Mapping is
     one-to-one and the induction variables of `scf.foreach_thread` are
-    rewritten to gpu.block_id according to the `thread_dim_apping` attribute.
+    rewritten to gpu.block_id according to the `thread_dim_mapping` attribute.
 
     Dynamic, `scf.foreach_thread` trip counts are currently not supported.
     Dynamic block dim sizes are currently not supported.
 
     Only **bufferized** scf.foreach_thread are currently supported.
     Only scf.foreach_thread distributed to **at most 3 dimensions** are
     currently supported.
 
let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$gridDim,
                    UnitAttr:$generate_gpu_launch);
   let results = (outs PDL_Operation:$result);
 
   let assemblyFormat = "$target attr-dict";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // GPU_TRANSFORM_OPS

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

 class TilingInterface;
 class RewriterBase;
 namespace linalg {
 class GenericOp;
 class LinalgOp;
 } // namespace linalg
 
 namespace transform {
+class TransformHandleTypeInterface;
 // Types needed for builders.
 struct TileSizesSpec {};
 struct NumThreadsSpec {};
 } // namespace transform
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Linalg Transform Operations

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

def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // FuseOp
 //===----------------------------------------------------------------------===//
 
def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // InterchangeOp
 //===----------------------------------------------------------------------===//
 
let assemblyFormat = [{
     $target 
     (`iterator_interchange` `=` $iterator_interchange^)? attr-dict
   }];
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::GenericOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MatchOp
 //===----------------------------------------------------------------------===//
 
def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
   let results = (outs PDL_Operation:$low_size,
                       PDL_Operation:$high_size,
                       PDL_Operation:$split_point);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//
 
def PadOp : Op<Transform_Dialect, "structured.pad",
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
let arguments = (ins PDL_Operation:$target,
                        OptionalAttr<I64Attr>:$alignment);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$result);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
OpBuilder<(ins "Value":$target,
                    CArg<"bool", "false">:$innerParallel,
                    CArg<"bool", "false">:$useScalingAlgorithm,
                    CArg<"bool", "false">:$useAlloc)>
   ];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TileReductionUsingScfOp
 //===----------------------------------------------------------------------===//
 
let assemblyFormat = [{
     $target
     `by` `tile_sizes` `=` $tile_sizes
     attr-dict
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TileReductionUsingForeachThreadOp
 //===----------------------------------------------------------------------===//
 
let assemblyFormat = [{
     (`,` `tile_sizes` `=` $tile_sizes^)?
     (`,` `mapping` `=` $mapping^)?
     attr-dict
   }];
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 
 }
 
 //===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
   let builders = [
     OpBuilder<(ins "Value":$target,
                CArg<"bool", "false">:$vectorizePadding,
                CArg<"bool", "false">:$vectorizeNDExtract)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformOpInterface]> {
   let description = [{

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

def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
 
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         memref::AllocOp target,
-        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // MEMREF_TRANSFORM_OPS

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

def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::scf::ForOp target,
-        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Applies software pipelining to the loop";
def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::scf::ForOp target,
-        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Unrolls the given loop with the given unroll factor";
def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
   let arguments = (ins TransformHandleTypeInterface:$target,
                        ConfinedAttr<I64Attr, [IntPositive]>:$factor);
 
   let assemblyFormat = "$target attr-dict `:` type($target)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
-        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // SCF_TRANSFORM_OPS

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

public:
   LogicalResult mapBlockArguments(TransformState &state) {
     assert(
         this->getOperation()->getNumRegions() == 1 &&
         "must indicate the region to map if the operation has more than one");
     return mapBlockArguments(state, this->getOperation()->getRegion(0));
   }
 };
 
 /// Trait implementing the TransformOpInterface for operations applying a
-/// transformation to a single operation handle and producing zero, one or
-/// multiple operation handles.
+/// transformation to a single operation handle and producing an arbitrary
+/// number of handles and parameter values.
 /// The op must implement a method with the following signature:
 ///   - DiagnosedSilenceableFailure applyToOne(OpTy,
 ///       SmallVector<Operation*> &results, state)
 /// to perform a transformation that is applied in turn to all payload IR
 /// operations that correspond to the handle of the transform IR operation.
 /// In `applyToOne`, OpTy is either Operation* or a concrete payload IR Op class
 /// that the transformation is applied to (and NOT the class of the transform IR
 /// op).

 
 } // namespace transform
 } // namespace mlir
 
 #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>;
+
+/// 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 {
+        results.push_back(static_cast<Attribute>(element));
+      }
+    }
+  }
+
+  /// Appends an element to the list.
+  void push_back(Operation *op) { results.push_back(op); }
+  void push_back(Attribute attr) { results.push_back(attr); }
+
+  /// 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(); }
+  auto end() const { return results.end(); }
+
+  /// Returns the number of elements in the list.
+  size_t size() const { return results.size(); }
+
+  /// Element access. Expects the index to be in bounds.
+  ApplyToEachResult &operator[](size_t index) { return results[index]; }
+  const ApplyToEachResult &operator[](size_t index) const {
+    return results[index];
+  }
+
+private:
+  /// Underlying storage.
+  SmallVector<ApplyToEachResult> results;
+};
+
 namespace detail {
+
+/// Check that the contents of `partialResult` matches the number, kind (payload
+/// op or parameter) and nullity (either all or none) requirements of
+/// `transformOp`. Report errors and return failure otherwise.
+LogicalResult checkApplyToOne(Operation *transformOp, Location payloadOpLoc,
+                              const ApplyToEachResultList &partialResult);
+
+/// "Transpose" the results produced by individual applications, arranging them
+/// per result value of the transform op, and populate `transformResults` with
+/// that. The number, kind and nullity of per-application results are assumed to
+/// have been verified.
+void setApplyToOneResults(Operation *transformOp,
+                          TransformResults &transformResults,
+                          ArrayRef<ApplyToEachResultList> results);
+
 /// Applies a one-to-one or a one-to-many transform to each of the given
 /// targets. Puts the results of transforms, if any, in `results` in the same
 /// order. Fails if any of the application fails. Individual transforms must be
 /// callable with the following signature:
 ///   - DiagnosedSilenceableFailure(OpTy,
 ///       SmallVector<Operation*> &results, state)
 /// where OpTy is either
 ///   - Operation *, in which case the transform is always applied;
 ///   - a concrete Op class, in which case a check is performed whether
 ///   `targets` contains operations of the same class and a silenceable failure
 ///   is reported if it does not.
-template <typename FnTy>
-DiagnosedSilenceableFailure applyTransformToEach(
-    Location loc, int expectedNumResults, ArrayRef<Operation *> targets,
-    SmallVectorImpl<SmallVector<Operation *>> &results, FnTy transform) {
-  SmallVector<Diagnostic> silenceableStack;
-  using OpTy = typename llvm::function_traits<FnTy>::template arg_t<0>;
+template <typename TransformOpTy>
+DiagnosedSilenceableFailure
+applyTransformToEach(TransformOpTy transformOp, ArrayRef<Operation *> targets,
+                     SmallVectorImpl<ApplyToEachResultList> &results,
+                     TransformState &state) {
+  using OpTy = typename llvm::function_traits<
+      decltype(&TransformOpTy::applyToOne)>::template arg_t<0>;
   static_assert(std::is_convertible<OpTy, Operation *>::value,
                 "expected transform function to take an operation");
+
+  SmallVector<Diagnostic> silenceableStack;
+  unsigned expectedNumResults = transformOp->getNumResults();
   for (Operation *target : targets) {
-    // Emplace back a placeholder for the returned new ops.
+    // Emplace back a placeholder for the returned new ops and params.
     // This is filled with `expectedNumResults` if the op fails to apply.
-    results.push_back(SmallVector<Operation *>());
+    ApplyToEachResultList placeholder;
+    placeholder.reserve(expectedNumResults);
+    results.push_back(std::move(placeholder));
 
     auto specificOp = dyn_cast<OpTy>(target);
     if (!specificOp) {
-      Diagnostic diag(loc, DiagnosticSeverity::Error);
+      Diagnostic diag(transformOp->getLoc(), DiagnosticSeverity::Error);
       diag << "transform applied to the wrong op kind";
       diag.attachNote(target->getLoc()) << "when applied to this op";
       // Producing `expectedNumResults` nullptr is a silenceableFailure mode.
       // TODO: encode this implicit `expectedNumResults` nullptr ==
       // silenceableFailure with a proper trait.
       results.back().assign(expectedNumResults, nullptr);
       silenceableStack.push_back(std::move(diag));
       continue;
     }
 
-    DiagnosedSilenceableFailure result = transform(specificOp, results.back());
-    if (result.isDefiniteFailure())
-      return result;
-    if (result.isSilenceableFailure())
-      result.takeDiagnostics(silenceableStack);
+    DiagnosedSilenceableFailure res =
+        transformOp.applyToOne(specificOp, results.back(), state);
+    if (res.isDefiniteFailure() ||
+        failed(detail::checkApplyToOne(transformOp, specificOp->getLoc(),
+                                       results.back()))) {
+      return DiagnosedSilenceableFailure::definiteFailure();
+    }
+
+    if (res.isSilenceableFailure())
+      res.takeDiagnostics(silenceableStack);
   }
   if (!silenceableStack.empty()) {
     return DiagnosedSilenceableFailure::silenceableFailure(
         std::move(silenceableStack));
   }
   return DiagnosedSilenceableFailure::success();
 }
 
-/// Helper function: transpose MxN into NxM; assumes that the input is a valid.
-static inline SmallVector<SmallVector<Operation *, 1>>
-transposeResults(const SmallVector<SmallVector<Operation *>, 1> &m) {
-  SmallVector<SmallVector<Operation *, 1>> res;
-  if (m.empty())
-    return res;
-  int64_t rows = m.size(), cols = m[0].size();
-  for (int64_t j = 0; j < cols; ++j)
-    res.push_back(SmallVector<Operation *, 1>(rows, nullptr));
-  for (int64_t i = 0; i < rows; ++i) {
-    assert(static_cast<int64_t>(m[i].size()) == cols);
-    for (int64_t j = 0; j < cols; ++j) {
-      res[j][i] = m[i][j];
-    }
-  }
-  return res;
-}
 } // namespace detail
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
     TransformResults &transformResults, TransformState &state) {
-  using TransformOpType = typename llvm::function_traits<
-      decltype(&OpTy::applyToOne)>::template arg_t<0>;
   ArrayRef<Operation *> targets =
       state.getPayloadOps(this->getOperation()->getOperand(0));
 
   // Step 1. Handle the corner case where no target is specified.
   // This is typically the case when the matcher fails to apply and we need to
   // propagate gracefully.
   // In this case, we fill all results with an empty vector.
   if (targets.empty()) {
-    SmallVector<Operation *> empty;
-    for (auto r : this->getOperation()->getResults())
-      transformResults.set(r.template cast<OpResult>(), empty);
+    SmallVector<Operation *> emptyPayload;
+    SmallVector<Attribute> emptyParams;
+    for (OpResult r : this->getOperation()->getResults()) {
+      if (r.getType().isa<TransformParamTypeInterface>())
+        transformResults.setParams(r, emptyParams);
+      else
+        transformResults.set(r, emptyPayload);
+    }
     return DiagnosedSilenceableFailure::success();
   }
 
   // Step 2. Call applyToOne on each target and record newly produced ops in its
   // corresponding results entry.
-  int expectedNumResults = this->getOperation()->getNumResults();
-  SmallVector<SmallVector<Operation *>, 1> results;
+  SmallVector<ApplyToEachResultList, 1> results;
+  results.reserve(targets.size());
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
-      this->getOperation()->getLoc(), expectedNumResults, targets, results,
-      [&](TransformOpType specificOp, SmallVector<Operation *> &partialResult) {
-        auto res = static_cast<OpTy *>(this)->applyToOne(specificOp,
-                                                         partialResult, state);
-        if (res.isDefiniteFailure())
-          return res;
-
-        // TODO: encode this implicit must always produce `expectedNumResults`
-        // and nullptr is fine with a proper trait.
-        if (static_cast<int>(partialResult.size()) != expectedNumResults) {
-          auto loc = this->getOperation()->getLoc();
-          auto diag = mlir::emitError(loc, "applications of ")
-                      << OpTy::getOperationName() << " expected to produce "
-                      << expectedNumResults << " results (actually produced "
-                      << partialResult.size() << ").";
-          diag.attachNote(loc)
-              << "If you need variadic results, consider a generic `apply` "
-              << "instead of the specialized `applyToOne`.";
-          diag.attachNote(loc)
-              << "Producing " << expectedNumResults << " null results is "
-              << "allowed if the use case warrants it.";
-          diag.attachNote(specificOp->getLoc()) << "when applied to this op";
-          return DiagnosedSilenceableFailure::definiteFailure();
-        }
-        // Check that all is null or none is null
-        // TODO: relax this behavior and encode with a proper trait.
-        if (llvm::any_of(partialResult, [](Operation *op) { return op; }) &&
-            llvm::any_of(partialResult, [](Operation *op) { return !op; })) {
-          auto loc = this->getOperation()->getLoc();
-          auto diag = mlir::emitError(loc, "unexpected application of ")
-                      << OpTy::getOperationName()
-                      << " produces both null and non null results.";
-          diag.attachNote(specificOp->getLoc()) << "when applied to this op";
-          return DiagnosedSilenceableFailure::definiteFailure();
-        }
-        return res;
-      });
+      cast<OpTy>(this->getOperation()), targets, results, state);
 
   // Step 3. Propagate the definite failure if any and bail out.
   if (result.isDefiniteFailure())
     return result;
 
-  // Step 4. If there are no results, return early.
-  if (OpTy::template hasTrait<OpTrait::ZeroResults>())
-    return result;
-
-  // Step 5. Perform transposition of M applications producing N results each
-  // into N results for each of the M applications.
-  SmallVector<SmallVector<Operation *, 1>> transposedResults =
-      detail::transposeResults(results);
-
-  // Step 6. Single result applies to M ops produces one single M-result.
-  if (OpTy::template hasTrait<OpTrait::OneResult>()) {
-    assert(transposedResults.size() == 1 && "Expected single result");
-    transformResults.set(
-        this->getOperation()->getResult(0).template cast<OpResult>(),
-        transposedResults[0]);
-    // ApplyToOne may have returned silenceableFailure, propagate it.
-    return result;
-  }
-
-  // Step 7. Filter out empty results and set the transformResults.
-  for (const auto &it :
-       llvm::zip(this->getOperation()->getResults(), transposedResults)) {
-    SmallVector<Operation *, 1> filtered;
-    llvm::copy_if(std::get<1>(it), std::back_inserter(filtered),
-                  [](Operation *op) { return op; });
-    transformResults.set(std::get<0>(it).template cast<OpResult>(), filtered);
-  }
+  // Step 4. "Transpose" the results produced by individual applications,
+  // arranging them per result value of the transform op. The number, kind and
+  // nullity of per-application results have been verified by the callback
+  // above.
+  detail::setApplyToOneResults(this->getOperation(), transformResults, results);
 
-  // Step 8. ApplyToOne may have returned silenceableFailure, propagate it.
+  // Step 5. ApplyToOne may have returned silenceableFailure, propagate it.
   return result;
 }
 
 template <typename OpTy>
 mlir::LogicalResult
 mlir::transform::TransformEachOpTrait<OpTy>::verifyTrait(Operation *op) {
   static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
                 "expected single-operand op");

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

def CastOp : TransformDialectOp<"cast",
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$input);
   let results = (outs TransformHandleTypeInterface:$output);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
       ::mlir::Operation *target,
-      ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+      ::mlir::transform::ApplyToEachResultList &results,
       ::mlir::transform::TransformState &state);
   }];
 }
 
 def ForeachOp : TransformDialectOp<"foreach",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface, [

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

 }
 
 //===----------------------------------------------------------------------===//
 // EmptyTensorToAllocTensorOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
-                                       SmallVector<Operation *> &results,
+                                       ApplyToEachResultList &results,
                                        transform::TransformState &state) {
   IRRewriter rewriter(target->getContext());
   rewriter.setInsertionPoint(target);
   auto alloc = rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
       target, target.getType(), target.getDynamicSizes());
   results.push_back(alloc);
   return DiagnosedSilenceableFailure::success();
 }

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

static void generateGpuBlockIds(RewriterBase &rewriter,
   blockOps = SmallVector<Value>{
       rewriter.create<BlockIdOp>(loc, indexType, Dimension::x),
       rewriter.create<BlockIdOp>(loc, indexType, Dimension::y),
       rewriter.create<BlockIdOp>(loc, indexType, Dimension::z)};
 }
 
 DiagnosedSilenceableFailure
 transform::MapForeachToBlocks::applyToOne(Operation *target,
-                                          SmallVectorImpl<Operation *> &results,
+                                          ApplyToEachResultList &results,
                                           transform::TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
   TrivialPatternRewriter rewriter(getContext());
   auto transformOp = cast<TransformOpInterface>(getOperation());
 
   if (!getGenerateGpuLaunch() && !gpuLaunch) {
-    results.assign({target});
+    results.push_back(target);
     DiagnosedSilenceableFailure diag =
         emitSilenceableError()
         << "Given target is not gpu.launch, set `generate_gpu_launch` "
            "attribute";
     diag.attachNote(target->getLoc()) << "when applied to this payload op";
     return diag;
   }
 
   scf::ForeachThreadOp topLevelForeachThreadOp;
   DiagnosedSilenceableFailure diag =
       mlir::transform::gpu::findTopLevelForeachThreadOp(
           target, topLevelForeachThreadOp, transformOp);
   if (!diag.succeeded()) {
-    results.assign({target});
+    results.push_back(target);
     diag.attachNote(target->getLoc()) << "when applied to this payload op";
     return diag;
   }
 
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(topLevelForeachThreadOp);
 
   // Generate gpu launch here and move the foreach_thread inside
   if (getGenerateGpuLaunch()) {
     DiagnosedSilenceableFailure diag =
         createGpuLaunch(rewriter, target->getLoc(), transformOp, gpuLaunch);
     if (!diag.succeeded()) {
-      results.assign({target});
+      results.push_back(target);
       return diag;
     }
     rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());
     Operation *newForeachThreadOp = rewriter.clone(*topLevelForeachThreadOp);
     rewriter.eraseOp(topLevelForeachThreadOp);
     topLevelForeachThreadOp = cast<scf::ForeachThreadOp>(newForeachThreadOp);
   }
 
diag = mlir::transform::gpu::mapForeachToBlocksImpl(
         rewriter, topLevelForeachThreadOp, generateGpuBlockIds, gridDim,
         transformOp, blockMappingAttributes);
   if (diag.succeeded()) {
     diag = alterGpuLaunch(rewriter, gpuLaunch,
                           cast<TransformOpInterface>(getOperation()),
                           gridDim[0], gridDim[1], gridDim[2]);
   }
 
-  results.assign({gpuLaunch});
+  results.push_back(gpuLaunch);
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // MapNestedForeachToThreads
 //===----------------------------------------------------------------------===//
 
 /// Searches `scf.foreach_thread` ops nested under `target` and maps each such
if (diag.succeeded()) {
           threadMappingAttributes);
     }
     return diag.succeeded() ? WalkResult::advance() : WalkResult::interrupt();
   });
   return diag;
 }
 
 DiagnosedSilenceableFailure transform::MapNestedForeachToThreads::applyToOne(
-    ::mlir::Operation *target,
-    ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
-    ::mlir::transform::TransformState &state) {
+    Operation *target, ApplyToEachResultList &results, TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
   auto transformOp = cast<TransformOpInterface>(getOperation());
 
   if (!gpuLaunch) {
-    results.assign({target});
+    results.push_back(target);
     return emitSilenceableError() << "Given target is not gpu.launch";
   }
 
   SmallVector<int64_t> blockDim = extractFromI64ArrayAttr(getBlockDim());
   blockDim.resize(/*size=*/3, /*value=*/1);
 
   DiagnosedSilenceableFailure diag =
       checkGpuLimits(transformOp, std::nullopt, std::nullopt, std::nullopt,
                      blockDim[0], blockDim[1], blockDim[2]);
   if (diag.isSilenceableFailure()) {
-    results.assign({target});
+    results.push_back(target);
     diag.attachNote(getLoc()) << getBlockDimAttrName() << " is very large";
     return diag;
   }
 
   MLIRContext *ctx = getContext();
   TrivialPatternRewriter rewriter(ctx);
   rewriter.setInsertionPoint(target);
 
   SmallVector<DeviceMappingAttrInterface> threadMappingAttributes = {
       GPUThreadMappingAttr::get(ctx, Threads::DimX),
       GPUThreadMappingAttr::get(ctx, Threads::DimY),
       GPUThreadMappingAttr::get(ctx, Threads::DimZ)};
 
   diag = mlir::transform::gpu::mapNestedForeachToThreadsImpl(
       rewriter, target, blockDim, getSyncAfterDistribute(), transformOp,
       threadMappingAttributes);
 
   if (diag.succeeded()) {
     diag = alterGpuLaunch(rewriter, gpuLaunch, transformOp, std::nullopt,
                           std::nullopt, std::nullopt, blockDim[0], blockDim[1],
                           blockDim[2]);
   }
 
-  results.assign({gpuLaunch});
+  results.push_back(gpuLaunch.getOperation());
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {

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

 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::DecomposeOp::applyToOne(linalg::LinalgOp target,
-                                   SmallVectorImpl<Operation *> &results,
+                                   transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
 #define DOWNSCALE(trans)                                                       \
   {                                                                            \
     FailureOr<LinalgOp> res = tryApply<trans>(target);                         \
     if (succeeded(res)) {                                                      \
       results.push_back(*res);                                                 \
       return DiagnosedSilenceableFailure::success();                           \
     }                                                                          \

 }
 
 //===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::GeneralizeOp::applyToOne(linalg::LinalgOp target,
-                                    SmallVectorImpl<Operation *> &results,
+                                    transform::ApplyToEachResultList &results,
                                     transform::TransformState &state) {
   // Exit early if no transformation is needed.
   if (isa<GenericOp>(target)) {
     results.push_back(target);
     return DiagnosedSilenceableFailure::success();
   }
   FailureOr<LinalgOp> generic = tryApply<LinalgGeneralizationPattern>(target);
   if (succeeded(generic)) {
     results.push_back(generic->getOperation());
     return DiagnosedSilenceableFailure::success();
   }
   results.assign(1, nullptr);
   return emitDefaultSilenceableFailure(target);
 }
 
 //===----------------------------------------------------------------------===//
 // InterchangeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::InterchangeOp::applyToOne(linalg::GenericOp target,
-                                     SmallVectorImpl<Operation *> &results,
+                                     transform::ApplyToEachResultList &results,
                                      transform::TransformState &state) {
   ArrayRef<int64_t> interchangeVector = getIteratorInterchange();
   // Exit early if no transformation is needed.
   if (interchangeVector.empty()) {
     results.push_back(target);
     return DiagnosedSilenceableFailure::success();
   }
   TrivialPatternRewriter rewriter(target->getContext());
transform::MatchOp::apply(transform::TransformResults &results,
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // MultiTileSizesOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
-    LinalgOp target, SmallVector<Operation *> &results, TransformState &state) {
+    LinalgOp target, transform::ApplyToEachResultList &results,
+    TransformState &state) {
   OpBuilder builder(target.getContext());
   builder.setInsertionPoint(target);
   OpFoldResult targetSize = builder.getIndexAttr(getTargetSize());
   OpFoldResult divisor = builder.getIndexAttr(getDivisor());
   FailureOr<MultiSizeSpecification> spec = computeMultiTileSizes(
       builder, target, getDimension(), targetSize, divisor);
   if (failed(spec)) {
     return emitSilenceableError() << "could not generate tile size computation";

 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PadOp::applyToOne(linalg::LinalgOp target,
-                             SmallVectorImpl<Operation *> &results,
+                             transform::ApplyToEachResultList &results,
                              transform::TransformState &state) {
   // Convert the integer packing flags to booleans.
   SmallVector<bool> packPaddings;
   for (int64_t packPadding : extractFromI64ArrayAttr(getPackPaddings()))
     packPaddings.push_back(static_cast<bool>(packPadding));
 
   // Convert the padding values to attributes.
   SmallVector<Attribute> paddingValues;

 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PromoteOp::applyToOne(linalg::LinalgOp target,
-                                 SmallVectorImpl<Operation *> &results,
+                                 transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {
   LinalgPromotionOptions promotionOptions;
   if (!getOperandsToPromote().empty())
     promotionOptions = promotionOptions.setOperandsToPromote(
         extractFromI64ArrayAttr(getOperandsToPromote()));
   if (getUseFullTilesByDefault())
     promotionOptions = promotionOptions.setUseFullTileBuffersByDefault(
         getUseFullTilesByDefault());

 }
 
 //===----------------------------------------------------------------------===//
 // ScalarizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
-                                   SmallVectorImpl<Operation *> &results,
+                                   transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.setTileSizeComputationFunction([&](OpBuilder &b, Operation *) {
     SmallVector<Value, 4> tileSizes;
     Location loc = target.getLoc();
     SmallVector<OpFoldResult> allShapeSizes =
         target.createFlatListOfOperandDims(b, loc);
     AffineMap map = target.getShapesToLoopsMap();
FailureOr<scf::SCFTilingResult> maybeTilingResult = tileUsingSCFForOp(
       rewriter, cast<TilingInterface>(target.getOperation()), tilingOptions);
   if (failed(maybeTilingResult))
     return emitDefaultDefiniteFailure(target);
 
   if (target->getNumResults())
     rewriter.replaceOp(target, maybeTilingResult->replacements);
   else
     rewriter.eraseOp(target);
-  results.append(maybeTilingResult->tiledOps);
+
+  results.reserve(maybeTilingResult->tiledOps.size());
+  for (Operation *tiled : maybeTilingResult->tiledOps)
+    results.push_back(tiled);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
void transform::SplitReductionOp::build(
   if (useAlloc) {
     result.addAttribute(SplitReductionOp::getUseAllocAttrName(result.name),
                         builder.getUnitAttr());
   }
   auto resultType = pdl::OperationType::get(ctx);
   result.addTypes({resultType, resultType, resultType, resultType});
 }
 
-DiagnosedSilenceableFailure
-transform::SplitReductionOp::applyToOne(linalg::LinalgOp target,
-                                        SmallVectorImpl<Operation *> &results,
-                                        transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::SplitReductionOp::applyToOne(
+    linalg::LinalgOp target, transform::ApplyToEachResultList &results,
   ControlSplitReductionFn splitFn = [&](LinalgOp) {
     return linalg::SplitReductionOptions{int64_t(getSplitFactor()),
                                          unsigned(getInsertSplitDimension()),
                                          bool(getInnerParallel())};
   };
   TrivialPatternRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<SplitReductionResult> splitResult =
void transform::TileReductionUsingScfOp::build(
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
   build(builder, result,
         /*resultTypes=*/TypeRange{opTy, opTy, opTy, opTy},
         /*target=*/target,
         /*tile_sizes=*/staticTileSizesAttr);
 }
 
 DiagnosedSilenceableFailure transform::TileReductionUsingScfOp::applyToOne(
-    linalg::LinalgOp target, SmallVectorImpl<Operation *> &results,
+    linalg::LinalgOp target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   TrivialPatternRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<scf::SCFReductionTilingResult> result = scf::tileReductionUsingScf(
       rewriter, cast<PartialReductionOpInterface>(target.getOperation()),
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getTileSizes())));
 
   if (failed(result))
build(builder, result,
         /*target=*/target,
         /*num_threads=*/staticNumThreadsAttr,
         /*tile_sizes=*/staticTileSizesAttr,
         /*mapping=*/mapping);
 }
 
 DiagnosedSilenceableFailure
 transform::TileReductionUsingForeachThreadOp::applyToOne(
-    linalg::LinalgOp target, SmallVectorImpl<Operation *> &results,
+    linalg::LinalgOp target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   TrivialPatternRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   SmallVector<OpFoldResult> numThreads =
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getNumThreads()));
   SmallVector<OpFoldResult> tileSizes =
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getTileSizes()));
   FailureOr<linalg::ForeachThreadReductionTilingResult> result =
private:
   /// Controls whether to vectorize `tensor.extract` when the input tensor is
   /// rank >= 2.
   bool vectorizeNDExtract = false;
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
-                                   SmallVectorImpl<Operation *> &results,
+                                   transform::ApplyToEachResultList &results,
                                    transform::TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     auto diag = this->emitOpError("requires isolated-from-above targets");
     diag.attachNote(target->getLoc()) << "non-isolated target";
     return DiagnosedSilenceableFailure::definiteFailure();
   }
 
   MLIRContext *ctx = getContext();

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

 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // MemRefMultiBufferOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure
-transform::MemRefMultiBufferOp::applyToOne(memref::AllocOp target,
-                                           SmallVector<Operation *> &results,
-                                           transform::TransformState &state) {
+DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::applyToOne(
+    memref::AllocOp target, transform::ApplyToEachResultList &results,
   auto newBuffer = memref::multiBuffer(target, getFactor());
   if (failed(newBuffer)) {
     Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);
     diag << "op failed to multibuffer";
     return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
   }
 
   results.push_back(*newBuffer);

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

 }
 
 //===----------------------------------------------------------------------===//
 // LoopPeelOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::LoopPeelOp::applyToOne(scf::ForOp target,
-                                  SmallVector<Operation *> &results,
+                                  transform::ApplyToEachResultList &results,
                                   transform::TransformState &state) {
   scf::ForOp result;
   IRRewriter rewriter(target->getContext());
   // This helper returns failure when peeling does not occur (i.e. when the IR
   // is not modified). This is not a failure for the op as the postcondition:
   //    "the loop trip count is divisible by the step"
   // is valid.
   LogicalResult status =
for (Operation *op : it.second) {
       unsigned cycle = opCycles[op];
       schedule.emplace_back(op, cycle / iterationInterval);
     }
   }
 }
 
 DiagnosedSilenceableFailure
 transform::LoopPipelineOp::applyToOne(scf::ForOp target,
-                                      SmallVector<Operation *> &results,
+                                      transform::ApplyToEachResultList &results,
                                       transform::TransformState &state) {
   scf::PipeliningOption options;
   options.getScheduleFn =
       [this](scf::ForOp forOp,
              std::vector<std::pair<Operation *, unsigned>> &schedule) mutable {
         loopScheduling(forOp, schedule, getIterationInterval(),
                        getReadLatency());
       };

 }
 
 //===----------------------------------------------------------------------===//
 // LoopUnrollOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::LoopUnrollOp::applyToOne(Operation *op,
-                                    SmallVector<Operation *> &results,
+                                    transform::ApplyToEachResultList &results,
                                     transform::TransformState &state) {
   LogicalResult result(failure());
   if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))
     result = loopUnrollByFactor(scfFor, getFactor());
   else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
     result = loopUnrollByFactor(affineFor, getFactor());
 
   if (failed(result)) {

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

 
 bool transform::TransformResults::isParam(unsigned resultNumber) const {
   assert(resultNumber < paramSegments.size() &&
          "querying association for a non-existent handle");
   return paramSegments[resultNumber].data() != nullptr;
 }
 
 //===----------------------------------------------------------------------===//
+// Utilities for TransformEachOpTrait.
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+transform::detail::checkApplyToOne(Operation *transformOp,
+                                   Location payloadOpLoc,
+                                   const ApplyToEachResultList &partialResult) {
+  Location transformOpLoc = transformOp->getLoc();
+  StringRef transformOpName = transformOp->getName().getStringRef();
+  unsigned expectedNumResults = transformOp->getNumResults();
+  // TODO: encode this implicit must always produce `expectedNumResults`
+  // and nullptr is fine with a proper trait.
+  if (partialResult.size() != expectedNumResults) {
+    auto diag = mlir::emitError(transformOpLoc, "applications of ")
+                << transformOpName << " expected to produce "
+                << expectedNumResults << " results (actually produced "
+                << partialResult.size() << ").";
+    diag.attachNote(transformOpLoc)
+        << "If you need variadic results, consider a generic `apply` "
+        << "instead of the specialized `applyToOne`.";
+    diag.attachNote(transformOpLoc)
+        << "Producing " << expectedNumResults << " null results is "
+        << "allowed if the use case warrants it.";
+    diag.attachNote(payloadOpLoc) << "when applied to this op";
+    return failure();
+  }
+
+  // Check that all is null or none is null
+  // TODO: relax this behavior and encode with a proper trait.
+  if (llvm::any_of(
+          partialResult,
+          [](llvm::PointerUnion<Operation *, Attribute> ptr) { return ptr; }) &&
+      llvm::any_of(partialResult,
+                   [](llvm::PointerUnion<Operation *, Attribute> ptr) {
+                     return !ptr;
+                   })) {
+    auto diag = mlir::emitError(transformOpLoc, "unexpected application of ")
+                << transformOpName
+                << " produces both null and non null results.";
+    diag.attachNote(payloadOpLoc) << "when applied to this op";
+    return failure();
+  }
+
+  // Check that the right kind of value was produced.
+  for (const auto &[ptr, res] :
+       llvm::zip(partialResult, transformOp->getResults())) {
+    if (ptr.is<Operation *>() &&
+        !res.getType().template isa<TransformHandleTypeInterface>()) {
+      mlir::emitError(transformOpLoc)
+          << "applications of " << transformOpName
+          << " expected to produce an Attribute for result #"
+          << res.getResultNumber();
+      return failure();
+    }
+    if (ptr.is<Attribute>() &&
+        !res.getType().template isa<TransformParamTypeInterface>()) {
+      mlir::emitError(transformOpLoc)
+          << "applications of " << transformOpName
+          << " expected to produce an Operation * for result #"
+          << res.getResultNumber();
+      return failure();
+    }
+  }
+  return success();
+}
+
+void transform::detail::setApplyToOneResults(
+    Operation *transformOp, TransformResults &transformResults,
+    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 {
+      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;
   if (op->getNumOperands() != 0)
     llvm::append_range(targets, state.getPayloadOps(op->getOperand(0)));

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

LogicalResult transform::AlternativesOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::CastOp::applyToOne(Operation *target,
-                              SmallVectorImpl<Operation *> &results,
+transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::CastOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsPayload(effects);

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

 }
 
 func.func private @three_test_ops(%arg0: i32) {
   "test.op_a"(%arg0) { attr = 0 : i32} : (i32) -> i32
   "test.op_a"(%arg0) { attr = 0 : i32} : (i32) -> i32
   "test.op_a"(%arg0) { attr = 0 : i32} : (i32) -> i32
   return
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{expected to produce an Operation * for result #0}}
+  transform.test_produce_transform_param_or_forward_operand %arg0
+    { first_result_is_param }
+    : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
+}
+
+// -----
+
+// expected-note @below {{when applied to this op}}
+module {
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-error @below {{produces both null and non null results}}
+    transform.test_produce_transform_param_or_forward_operand %arg0
+      { first_result_is_null }
+      : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
+  }
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{expected to produce an Attribute for result #1}}
+  transform.test_produce_transform_param_or_forward_operand %arg0
+    { second_result_is_handle }
+    : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
+}

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

for (Operation *op : state.getPayloadOps(getTarget()))
     op->erase();
 
   if (getFailAfterErase())
     return emitSilenceableError() << "silenceable error";
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestWrongNumberOfResultsOp::applyToOne(
-    Operation *target, SmallVectorImpl<Operation *> &results,
+    Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestWrongNumberOfMultiResultsOp::applyToOne(
-    Operation *target, SmallVectorImpl<Operation *> &results,
+    Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   static int count = 0;
   if (count++ == 0) {
     OperationState opState(target->getLoc(), "foo");
     results.push_back(OpBuilder(target).create(opState));
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestCorrectNumberOfMultiResultsOp::applyToOne(
-    Operation *target, SmallVectorImpl<Operation *> &results,
+    Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(OpBuilder(target).create(opState));
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestMixedNullAndNonNullResultsOp::applyToOne(
-    Operation *target, SmallVectorImpl<Operation *> &results,
+    Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(nullptr);
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
-    Operation *target, SmallVectorImpl<Operation *> &results,
+    Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   if (target->hasAttr("target_me"))
     return DiagnosedSilenceableFailure::success();
   return emitDefaultSilenceableFailure(target);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestPrintNumberOfAssociatedPayloadIROps::apply(

 
 LogicalResult mlir::test::TestProduceIntegerParamWithTypeOp::verify() {
   if (!getType().isa<IntegerType>()) {
     return emitOpError() << "expects an integer type";
   }
   return success();
 }
 
+void mlir::test::TestProduceTransformParamOrForwardOperandOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getIn(), effects);
+  transform::producesHandle(getOut(), effects);
+  transform::producesHandle(getParam(), effects);
+}
+
+DiagnosedSilenceableFailure
+mlir::test::TestProduceTransformParamOrForwardOperandOp::applyToOne(
+    Operation *target, ::transform::ApplyToEachResultList &results,
+    ::transform::TransformState &state) {
+  Builder builder(getContext());
+  if (getFirstResultIsParam()) {
+    results.push_back(builder.getI64IntegerAttr(0));
+  } else if (getFirstResultIsNull()) {
+    results.push_back(nullptr);
+  } else {
+    results.push_back(state.getPayloadOps(getIn()).front());
+  }
+
+  if (getSecondResultIsHandle()) {
+    results.push_back(state.getPayloadOps(getIn()).front());
+  } else {
+    results.push_back(builder.getI64IntegerAttr(42));
+  }
+
+  return DiagnosedSilenceableFailure::success();
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

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

def TestWrongNumberOfResultsOp
   let results = (outs PDL_Operation:$a,
                       PDL_Operation:$b,
                       PDL_Operation:$c);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestWrongNumberOfMultiResultsOp
   : Op<Transform_Dialect, "test_wrong_number_of_multi_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$result);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestCorrectNumberOfMultiResultsOp
   : Op<Transform_Dialect, "test_correct_number_of_multi_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$result1,
                       PDL_Operation:$result2);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestMixedNullAndNonNullResultsOp
   : Op<Transform_Dialect, "test_mixed_null_and_non_null_results",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$null,
                       PDL_Operation:$non_null);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestMixedSuccessAndSilenceableOp
   : Op<Transform_Dialect, "test_mixed_sucess_and_silenceable",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins PDL_Operation:$target);
   let results = (outs);
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target,
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def TestPrintNumberOfAssociatedPayloadIROps
   : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_ops",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
: Op<Transform_Dialect, "test_produce_integer_param_with_type",
         DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins TypeAttr:$type);
   let results = (outs TransformParamTypeInterface:$result);
   let assemblyFormat = "$type attr-dict `:` type($result)";
   let cppNamespace = "::mlir::test";
   let hasVerifier = 1;
 }
 
+def TestProduceTransformParamOrForwardOperandOp
+  : Op<Transform_Dialect, "test_produce_transform_param_or_forward_operand",
+       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        TransformEachOpTrait, TransformOpInterface]> {
+  let arguments = (ins TransformHandleTypeInterface:$in,
+                       UnitAttr:$first_result_is_param,
+                       UnitAttr:$first_result_is_null,
+                       UnitAttr:$second_result_is_handle);
+  let results = (outs TransformHandleTypeInterface:$out,
+                      TransformParamTypeInterface:$param);
+  let assemblyFormat = "$in 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);
+  }];
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD