[mlir][linalg][NFC] Return tensor::PadOp handle from transform op

[springerm]

Summary

"transform.structured.pad" now returns all tensor::PadOp in addition to the padded ops.

Also add a test case that shows how to force an allocation for "tensor.pad" ops with a custom memory space.

Depends On: D153554

[springerm]

update

[springerm]

rebase

[springerm]

clang-format

[springerm]

rebase

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

 }
 
 //===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//
 
 def PadOp : Op<Transform_Dialect, "structured.pad",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-     TransformOpInterface, TransformEachOpTrait,
+     DeclareOpInterfaceMethods<TransformOpInterface>,
      ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Pads the operations pointed to by the target handle using the options
-    provides as operation attributes.
+    provides as operation attributes. The operation returns a handle to the
+    padded operation and to the padding operation ("tensor.pad").
 
     #### Return modes
 
     This operation ignores non-Linalg ops and drops them in the return.
     This operation may produce a definiteFailure if the padding fails for any
     reason.
+
     If all the operations referred to by the `target` handle pad
     properly, the transform succeeds. Otherwise the transform silently fails.
     The return handle points to only the subset of successfully produced
     padded operations, which can be empty.
   }];
 
   let arguments =
     (ins TransformHandleTypeInterface:$target,
          DefaultValuedAttr<ArrayAttr, "{}">:$padding_values,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$padding_dimensions,
          OptionalAttr<I64ArrayAttr>:$pad_to_multiple_of,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$pack_paddings,
          DefaultValuedAttr<
           TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
           "{}">:$transpose_paddings,
          DefaultValuedAttr<BoolAttr, "true">:$copy_back);
-  let results = (outs TransformHandleTypeInterface:$transformed);
+  let results = (outs TransformHandleTypeInterface:$padded,
+                      TransformHandleTypeInterface:$pad);
 
   let assemblyFormat =
-    "$target attr-dict `:` "
-    "custom<SemiFunctionType>(type($target), type($transformed))";
+    "$target attr-dict `:` functional-type(operands, results)";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 // TODO: Add support for scf.parallel and affine.for loops.
 SmallVector<Value> peelLoop(RewriterBase &rewriter, Operation *op);
 
 /// Peel 'loops' and applies affine_min/max bounds simplification on the fly
 /// where relevant.
 void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
 
 /// Pad the iterator dimensions `paddingDimensions` of all `opToPad` operands
-/// to a static bounding box. `padToMultipleOf` indicates that each padding
-/// dimension should be padded to the specified multiple. If the derived padding
-/// sizes should not be rounded up to any multiple, use "1". Use `paddingValues`
-/// and `packPaddings` to set padding value and nofold attribute of the created
-/// tensor::PadOps, respectively. Update `paddedOp` to the cloned operation with
-/// statically shaped `paddingDimensions` and return the extracted dynamically
-/// shaped results. If padding fails, return failure. If `copyBack` is set, the
-/// unpadded result is copied back into the original destination tensor.
-FailureOr<SmallVector<Value>> rewriteAsPaddedOp(
-    RewriterBase &rewriter, LinalgOp opToPad,
-    ArrayRef<int64_t> paddingDimensions, ArrayRef<int64_t> padToMultipleOf,
-    ArrayRef<Attribute> paddingValues, ArrayRef<bool> packPaddings,
-    LinalgOp &paddedOp, bool copyBack);
+/// to a static bounding box. The original `opToPad` is cloned and operates on
+/// the padded tensors.
+///
+/// * `padToMultipleOf` indicates that each padding dimension should be padded
+///   to the specified multiple. If the derived padding sizes should not be
+///   rounded up to any multiple, use "1".
+/// * Use `paddingValues` and `packPaddings` to set padding value and nofold
+///   attribute of the created tensor::PadOps, respectively.
+/// * The unpadded results (extracted slice of the cloned operation) are
+///   returned via `replacements`.
+/// * The tensor::PadOps are returned via `padOps`.
+/// * If `copyBack` is set to "true", the unpadded result is copied back to the
+///   original destination tensor.
+LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
+                                ArrayRef<int64_t> paddingDimensions,
+                                ArrayRef<int64_t> padToMultipleOf,
+                                ArrayRef<Attribute> paddingValues,
+                                ArrayRef<bool> packPaddings, LinalgOp &paddedOp,
+                                SmallVector<Value> &replacements,
+                                SmallVector<tensor::PadOp> &padOps,
+                                bool copyBack);
 
 namespace detail {
 
 /// Helper struct to hold the results of building a packing loop nest.
 struct PackingResult {
   SmallVector<OpFoldResult> offsets, sizes, strides;
   SmallVector<Value> clonedLoopIvs, leadingPackedTensorIndexings;
   GenericOp maybeTransposeOp;

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

transform::PackTransposeOp::apply(transform::TransformRewriter &rewriter,
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::PadOp::applyToOne(transform::TransformRewriter &rewriter,
-                             LinalgOp target,
-                             transform::ApplyToEachResultList &results,
-                             transform::TransformState &state) {
+transform::PadOp::apply(transform::TransformRewriter &rewriter,
+                        transform::TransformResults &results,
+  SmallVector<Operation *> paddedOps, padOps;
+
+  for (Operation *target : state.getPayloadOps(getTarget())) {
+    auto linalgTarget = dyn_cast<LinalgOp>(target);
+    if (!linalgTarget) {
+      auto diag = emitSilenceableError() << "expected LinalgOp target";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+
     // 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;
     for (auto const &it :
-       llvm::zip(getPaddingValues(), target->getOperandTypes())) {
+         llvm::zip(getPaddingValues(), linalgTarget->getOperandTypes())) {
       auto attr = dyn_cast<TypedAttr>(std::get<0>(it));
       if (!attr) {
         emitOpError("expects padding values to be typed attributes");
         return DiagnosedSilenceableFailure::definiteFailure();
       }
       Type elementType = getElementTypeOrSelf(std::get<1>(it));
       // Try to parse string attributes to obtain an attribute of element type.
       if (auto stringAttr = dyn_cast<StringAttr>(attr)) {
-      auto parsedAttr = dyn_cast_if_present<TypedAttr>(
-          parseAttribute(stringAttr, getContext(), elementType,
-                         /*numRead=*/nullptr, /*isKnownNullTerminated=*/true));
+        auto parsedAttr = dyn_cast_if_present<TypedAttr>(parseAttribute(
+            stringAttr, getContext(), elementType,
         if (!parsedAttr || parsedAttr.getType() != elementType) {
           auto diag = this->emitOpError("expects a padding that parses to ")
                       << elementType << ", got " << std::get<0>(it);
-        diag.attachNote(target.getLoc()) << "when applied to this op";
+          diag.attachNote(linalgTarget.getLoc()) << "when applied to this op";
           return DiagnosedSilenceableFailure::definiteFailure();
         }
         paddingValues.push_back(parsedAttr);
         continue;
       }
       // Otherwise, add the attribute directly.
       if (attr.getType() != elementType) {
         auto diag = this->emitOpError("expects a padding value of type ")
                     << elementType << ", got " << attr;
-      diag.attachNote(target.getLoc()) << "when applied to this op";
+        diag.attachNote(linalgTarget.getLoc()) << "when applied to this op";
         return DiagnosedSilenceableFailure::definiteFailure();
       }
       paddingValues.push_back(attr);
     }
 
     // Extract the transpose vectors.
     SmallVector<SmallVector<int64_t>> transposePaddings;
     for (Attribute transposeVector : cast<ArrayAttr>(getTransposePaddings()))
       transposePaddings.push_back(
           extractFromI64ArrayAttr(cast<ArrayAttr>(transposeVector)));
 
     LinalgOp paddedOp;
     SmallVector<int64_t> paddingDimensions =
         extractFromI64ArrayAttr(getPaddingDimensions());
     SmallVector<int64_t> padToMultipleOf(paddingDimensions.size(), 1);
     if (getPadToMultipleOf().has_value())
       padToMultipleOf = extractFromI64ArrayAttr(*getPadToMultipleOf());
-  FailureOr<SmallVector<Value>> result =
-      rewriteAsPaddedOp(rewriter, target, paddingDimensions, padToMultipleOf,
-                        paddingValues, packPaddings, paddedOp, getCopyBack());
-  if (succeeded(result)) {
+    SmallVector<Value> replacements;
+    SmallVector<tensor::PadOp> newPadOps;
+    if (failed(rewriteAsPaddedOp(rewriter, linalgTarget, paddingDimensions,
+                                 padToMultipleOf, paddingValues, packPaddings,
+                                 paddedOp, replacements, newPadOps,
+                                 getCopyBack()))) {
+      auto diag = emitSilenceableError() << "failed to pad op";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
     // TODO: clean this up and stop "pad and hoist" behavior more globally now
     // that we have more composable abstractions.
-    rewriter.replaceOp(target, *result);
-    results.push_back(paddedOp);
-    return DiagnosedSilenceableFailure::success();
+    rewriter.replaceOp(linalgTarget, replacements);
+    paddedOps.push_back(paddedOp);
+    padOps.append(newPadOps.begin(), newPadOps.end());
   }
 
-  return emitDefaultSilenceableFailure(target);
+  results.set(cast<OpResult>(getPadded()), paddedOps);
+  results.set(cast<OpResult>(getPad()), padOps);
+  return DiagnosedSilenceableFailure::success();
 }
 
 LogicalResult transform::PadOp::verify() {
   SmallVector<int64_t> packPaddings =
       extractFromI64ArrayAttr(getPackPaddings());
   if (any_of(packPaddings, [](int64_t packPadding) {
         return packPadding != 0 && packPadding != 1;
       })) {

mlir/lib/Dialect/Linalg/Transforms/Padding.cpp

static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
   auto paddedTensorType = RankedTensorType::get(
       paddedShape, getElementTypeOrSelf(opOperand->get()));
   LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
                     << paddedTensorType);
   return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
                                opOperand->get(), paddingValue, nofold);
 }
 
-FailureOr<SmallVector<Value>> linalg::rewriteAsPaddedOp(
+LogicalResult linalg::rewriteAsPaddedOp(
     RewriterBase &rewriter, LinalgOp opToPad,
     ArrayRef<int64_t> paddingDimensions, ArrayRef<int64_t> padToMultipleOf,
     ArrayRef<Attribute> paddingValues, ArrayRef<bool> packPaddings,
-    LinalgOp &paddedOp, bool copyBack) {
+    LinalgOp &paddedOp, SmallVector<Value> &replacements,
+    SmallVector<tensor::PadOp> &padOps, bool copyBack) {
   LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << opToPad << "\n");
   Location loc = opToPad->getLoc();
 
   // TODO: there are cases where we may still want to pad to larger sizes.
   if (!opToPad.hasTensorSemantics())
     return rewriter.notifyMatchFailure(opToPad,
                                        "expected operation on tensors");
 
for (OpOperand &opOperand : opToPad->getOpOperands()) {
     // Exit if `paddingDimensions` cannot be bounded statically.
     if (failed(paddedOperand)) {
       LLVM_DEBUG(DBGS() << "--operand cannot be bound statically : "
                         << opOperand.get() << " -> FAIL\n");
       return rewriter.notifyMatchFailure(opToPad,
                                          "operand cannot be bound statically");
     }
     newOperands.push_back(*paddedOperand);
+    if (auto padOp = paddedOperand->getDefiningOp<tensor::PadOp>())
+      padOps.push_back(padOp);
   }
 
   ReifiedRankedShapedTypeDims reifiedResultShapes;
   if (failed(reifyResultShapes(rewriter, opToPad, reifiedResultShapes))) {
     LLVM_DEBUG(DBGS() << "--failed to reify result shapes -> FAIL\n");
     return rewriter.notifyMatchFailure(opToPad,
                                        "failed to reify result shapes");
   }
for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
     int64_t rank = cast<RankedTensorType>(paddedResult.getType()).getRank();
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
     paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
         strides));
   }
 
-  if (!copyBack)
-    return paddedSubtensorResults;
+  if (!copyBack) {
+    replacements = std::move(paddedSubtensorResults);
+    return success();
+  }
 
   // Copy back unpadded results to the original destination (i.e., inits of the
   // linalg op), so that the destination buffer of the computation does not
   // change. If the padding folds away, this will materizalize as a memcpy
   // between two identical buffers, which will then also fold away.
-  SmallVector<Value> copiedBack;
+  assert(paddedSubtensorResults.size() == opToPad.getNumDpsInits() &&
+         "expected matching number of results");
   for (auto it :
        llvm::zip(paddedSubtensorResults, opToPad.getDpsInitOperands())) {
-    copiedBack.push_back(rewriter.create<bufferization::CopyTensorOp>(
+    replacements.push_back(rewriter.create<bufferization::CopyTensorOp>(
         loc, std::get<0>(it), std::get<1>(it)->get()));
   }
-  return copiedBack;
+  return success();
 }
 
 FailureOr<LinalgOp>
 mlir::linalg::padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
                                   LinalgPaddingOptions options) {
   if (!linalgOp.hasTensorSemantics())
     return rewriter.notifyMatchFailure(
         linalgOp, "only applies to Linalg ops with tensor semantics");
 
   // Pad the operation.
   LinalgOp paddedOp;
   SmallVector<int64_t> padToMultipleOf(options.paddingDimensions.size(), 1);
   if (options.padToMultipleOf.has_value())
     padToMultipleOf.assign(options.padToMultipleOf->begin(),
                            options.padToMultipleOf->end());
-  FailureOr<SmallVector<Value>> newResults =
-      rewriteAsPaddedOp(rewriter, linalgOp, options.paddingDimensions,
+  SmallVector<Value> newResults;
+  SmallVector<tensor::PadOp> padOps;
+  if (failed(rewriteAsPaddedOp(rewriter, linalgOp, options.paddingDimensions,
                                padToMultipleOf, options.paddingValues,
-                        options.packPaddings, paddedOp, /*copyBack=*/false);
-  if (failed(newResults))
+                               options.packPaddings, paddedOp, newResults,
+                               padOps, /*copyBack=*/false)))
     return rewriter.notifyMatchFailure(linalgOp,
                                        "failed to rewrite as a padded op");
 
   // Hoist the padding.
   for (const auto &en : enumerate(options.hoistPaddings)) {
     if (static_cast<int64_t>(en.index()) >= paddedOp->getNumOperands())
       break;
     OpOperand &opOperand = paddedOp->getOpOperand(en.index());
if (failed(newResult)) {
       (void)rewriter.notifyMatchFailure(linalgOp,
                                         "failed to apply hoistPadding");
       continue;
     }
     rewriter.replaceOp(padOp, *newResult);
   }
 
   // Replace the original operation to pad.
-  rewriter.replaceOp(linalgOp, *newResults);
+  rewriter.replaceOp(linalgOp, newResults);
 
   return paddedOp;
 }

mlir/test/Dialect/Linalg/pad-to-specific-memory-space.mlir

+
+// RUN: mlir-opt --test-transform-dialect-interpreter -cse -canonicalize -split-input-file -verify-diagnostics %s | FileCheck %s
+
+#map = affine_map<()[s0] -> (-s0 + 12, 7)>
+
+// CHECK-LABEL: func @pad_to_memory_space(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<24x12xf32, strided<[?, ?], offset: ?>>,
+//  CHECK-SAME:     %[[arg1:.*]]: memref<12x25xf32, strided<[?, ?], offset: ?>>,
+//  CHECK-SAME:     %[[arg2:.*]]: memref<24x25xf32, strided<[?, ?], offset: ?>>,
+func.func @pad_to_memory_space(%arg0: tensor<24x12xf32>,
+                               %arg1: tensor<12x25xf32>,
+                               %arg2: tensor<24x25xf32>,
+                               %iv0 : index, %iv1 : index,
+                               %iv2 : index) -> tensor<24x25xf32> {
+  %0 = affine.min #map()[%iv2]
+
+  // CHECK: %[[s0:.*]] = memref.subview %[[arg0]]
+  %1 = tensor.extract_slice %arg0[%iv0, %iv2] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
+  // CHECK: %[[s1:.*]] = memref.subview %[[arg1]]
+  %2 = tensor.extract_slice %arg1[%iv2, %iv1] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
+  // CHECK: %[[s2:.*]] = memref.subview %[[arg2]]
+  %3 = tensor.extract_slice %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
+
+  // CHECK: %[[alloc0:.*]] = memref.alloc() : memref<4x7xf32, 3>
+  // CHECK: linalg.fill {{.*}} outs(%[[alloc0]]
+  // CHECK: %[[alloc0_view:.*]] = memref.subview %[[alloc0]][0, 0] [4, %{{.*}}] [1, 1]
+  // CHECK: memref.copy %[[s0]], %[[alloc0_view]]
+
+  // CHECK: %[[alloc1:.*]] = memref.alloc() : memref<7x5xf32, 3>
+  // CHECK: linalg.fill {{.*}} outs(%[[alloc1]]
+  // CHECK: %[[alloc1_view:.*]] = memref.subview %[[alloc1]][0, 0] [%{{.*}}, 5] [1, 1]
+  // CHECK: memref.copy %[[s1]], %[[alloc1_view]]
+
+  // CHECK: %[[alloc2:.*]] = memref.alloc() : memref<4x5xf32, 3>
+  // CHECK: linalg.fill {{.*}} outs(%[[alloc2]]
+  // No subview because there is 0 padding
+  // CHECK: memref.copy %[[s2]], %[[alloc2]]
+
+  // CHECK: linalg.matmul ins(%[[alloc0]], %[[alloc1]] : {{.*}}) outs(%[[alloc2]] : {{.*}})
+  // Copy back result.
+  // CHECK: memref.copy %[[alloc2]], %[[s2]]
+  %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
+
+  // insert_slice bufferizes to a no-op.
+  %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
+  func.return %5 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %padded, %pad = transform.structured.pad %0 {
+    padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 1]
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  %pad_result = transform.get_result %pad[0] : (!transform.any_op) -> !transform.any_value
+  %buffer, %replacement = transform.structured.bufferize_to_allocation %pad_result {memory_space = 3}
+  %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
+
+}

mlir/test/Dialect/Linalg/transform-op-hoist-pad.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
-  %matmul_padded = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0 = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back = false
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   // In this case, the pad op is actually empty: we only tile the first dimension
   // and it does not have an impact on the RHS operand.
   // expected-error @below {{incompatible payload operation name}}
   %pad = transform.get_producer_of_operand %matmul_padded[1]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
   // We do not even reach this transform op.

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
-  %matmul_padded = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0 = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back = false
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %pad = transform.get_producer_of_operand %matmul_padded[2]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
   // We do not know yet how to hoist the init.
   // expected-error @below {{transform.structured.hoist_pad failed to apply}}
   transform.structured.hoist_pad %pad by 1 loops
      : (!transform.op<"tensor.pad">) -> !transform.any_op

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
-  %matmul_padded = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0 = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back = false
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %pad = transform.get_producer_of_operand %matmul_padded[0]
     : (!transform.any_op) -> !transform.any_op
 
   transform.structured.hoist_pad %pad by 1 loops
      : (!transform.any_op) -> !transform.any_op
 }
 

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
-  %matmul_padded = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0 = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back = false
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %pad = transform.get_producer_of_operand %matmul_padded[0]
     : (!transform.any_op) -> !transform.any_op
 
   transform.structured.hoist_pad %pad by 1 loops, transpose by [1, 0]
      : (!transform.any_op) -> !transform.any_op
 }
 

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1:2 = transform.structured.tile_to_scf_for %matmul [5, 0, 7] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 
-  %matmul_padded = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0 = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back = false
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %pad = transform.get_producer_of_operand %matmul_padded[2]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
   transform.structured.hoist_pad %pad by 1 loops
      : (!transform.op<"tensor.pad">) -> !transform.any_op
 }

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

func.func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
 
 // CHECK-LABEL: @pad_to_multiple
 func.func @pad_to_multiple(%arg0: tensor<24x12xf32>,
func.func @pad_to_multiple(%arg0: tensor<24x12xf32>,
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pad_to_multiple_of=[2, 2, 1],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
 
 // CHECK-LABEL: @static_sizes_output_divisible_on_empty_op
 func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
   %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{op expects a padding value of type 'f32', got 0 : i32}}
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0: i32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{expects a padding that parses to 'f32', got "{foo}"}}
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=["{foo}", 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 // CHECK-LABEL: @pad(
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // This is attached to an error that is silenceable and is not reported by this transform
   //   {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(suppress) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // This error is silenceable and is not reported by this transform
   //   {{transform.structured.pad failed to apply}}
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 // Check that the padding can be applied even when the output argument of the
 // linalg op is not produced by an empty op or an extract_slice op.
 
 // CHECK-DAG: #[[$MAP_MIN:.*]] = affine_map<(d0) -> (-d0 + 2044, 16)>
%10 = scf.for %arg3 = %c0 to %c2044 step %c16 iter_args(%arg4 = %9) -> (tensor<128x128xf32>) {
     scf.yield %res : tensor<128x128xf32>
   }
   return %10 : tensor<128x128xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
 
 // CHECK-LABEL: @pack_everything
 func.func @pack_everything(%arg0: tensor<24x12xf32>,
func.func @pack_everything(%arg0: tensor<24x12xf32>,
   //      CHECK: %[[T8:.*]] = tensor.insert_slice %[[T7]] into %{{.*}}
   %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
   func.return %5 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = transform.structured.pad %0 {
+  %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
-  } : (!transform.any_op) -> !transform.any_op
+  } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }