[mlir][linalg][transform] Return copy_back op from PadOp.

[ingomueller-net]

Summary

This patch makes the transform.structured.pad op return also a handle
to the copy op that it inserts. This allows to continue transformation
on that op, such as mapping it to a GPU thread.

The patch was mainly authored by @springerm as part of the WIP patch
https://reviews.llvm.org/D156371, which also has an example usage of
this change.

[ingomueller-net]

• Fix mix-in by adding default type for new result value.

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

(ins TransformHandleTypeInterface:$target,
          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<StrAttr, "::mlir::bufferization::CopyTensorOp::getOperationName()">:$copy_back_op);
   let results = (outs TransformHandleTypeInterface:$padded,
-                      TransformHandleTypeInterface:$pad);
+                      TransformHandleTypeInterface:$pad,
+                      TransformHandleTypeInterface:$copy);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
   let hasVerifier = 1;
 
   let builders = [
     // Builder for a transform::PadOp with automatic inference of padding
     // value. Warning: this will set the value 0 for the inferred elemental

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

return build(/*builder=*/b,
                /*transposePaddings=*/b.getArrayAttr(transposePaddings),
                /*copyBackOp=*/b.getStringAttr(copyBackOp));
 }
 
 DiagnosedSilenceableFailure
 transform::PadOp::apply(transform::TransformRewriter &rewriter,
                         transform::TransformResults &results,
                         transform::TransformState &state) {
-  SmallVector<Operation *> paddedOps, padOps;
+  SmallVector<Operation *> paddedOps, padOps, copyBackOps;
 
   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;
     }
for (Operation *target : state.getPayloadOps(getTarget())) {
     // 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(linalgTarget, replacements);
     paddedOps.push_back(paddedOp);
     padOps.append(newPadOps.begin(), newPadOps.end());
+    if (options.copyBackOp != LinalgPaddingOptions::CopyBackOp::None) {
+      for (Value v : replacements) {
+        Operation *copyBackOp = v.getDefiningOp();
+        if (llvm::find(copyBackOps, copyBackOp) == copyBackOps.end())
+          copyBackOps.push_back(copyBackOp);
+      }
+    }
   }
 
   results.set(cast<OpResult>(getPadded()), paddedOps);
   results.set(cast<OpResult>(getPad()), padOps);
+  results.set(cast<OpResult>(getCopy()), copyBackOps);
   return DiagnosedSilenceableFailure::success();
 }
 
 LogicalResult transform::PadOp::verify() {
   SmallVector<int64_t> packPaddings =
       extractFromIntegerArrayAttr<int64_t>(getPackPaddings());
   if (any_of(packPaddings, [](int64_t packPadding) {
         return packPadding != 0 && packPadding != 1;

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

):
         pad_to_multiple_of = _get_int_array_attr(pad_to_multiple_of)
         pack_paddings = _get_int_array_attr(pack_paddings)
         transpose_paddings = _get_int_array_array_attr(transpose_paddings)
 
         pdl_operation_type = pdl.OperationType.get()
         super().__init__(
             pdl_operation_type,
             pdl_operation_type,
+            pdl_operation_type,
             target,
             padding_values=padding_values,
             padding_dimensions=padding_dimensions,
             pad_to_multiple_of=pad_to_multiple_of,
             pack_paddings=pack_paddings,
             transpose_paddings=transpose_paddings,
             copy_back_op=copy_back_op,
             loc=loc,

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

func.func @pad_to_memory_space(%arg0: tensor<24x12xf32>,
   // 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 {
+  %padded, %pad, %copy_back = 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, !transform.any_op, !transform.any_op)
   %buffer, %new_ops = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op
   %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
func.func @vectorize_and_bufferize_pad(%arg0: tensor<24x12xf32>,
   // 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 {
+  %padded, %pad, %copy_back = 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, !transform.any_op, !transform.any_op)
   transform.structured.masked_vectorize %pad vector_sizes [10, 12] : !transform.any_op
   %vector_write = transform.structured.match ops{["vector.transfer_write"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %mask_op = transform.get_parent_op %vector_write {op_name = "vector.mask"} : (!transform.any_op) -> !transform.any_op
   %buffer, %new_ops = transform.structured.bufferize_to_allocation %mask_op {memory_space = 3} : !transform.any_op
   %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-build-packing-loop-nest.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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2]
-  } : (!transform.any_op) -> (!transform.any_op, !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.
   %pad = transform.get_producer_of_operand %matmul_padded[1]
     : (!transform.any_op) -> !transform.any_op
 
   // expected-error @below {{requires exactly 2 non-null handles}}
   transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1

 
 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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2]
-  } : (!transform.any_op) -> (!transform.any_op, !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.any_op
 
   // We do not know yet how to hoist the init.
   // expected-error @below {{could not build packing loop nest}}
   transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1
      : (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2]
-  } : (!transform.any_op) -> (!transform.any_op, !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.build_packing_loop_nest %pad above %loops_l1
      : (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2]
-  } : (!transform.any_op) -> (!transform.any_op, !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.build_packing_loop_nest %pad above %loops_l1, transpose by [1, 0]
      : (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2]
-  } : (!transform.any_op) -> (!transform.any_op, !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.any_op
 
   transform.structured.hoist_pad.build_packing_loop_nest %pad above %loops_l1#1
      : (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back_op = "none"
-  } : (!transform.any_op) -> (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back_op = "none"
-  } : (!transform.any_op) -> (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back_op = "none"
-  } : (!transform.any_op) -> (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back_op = "none"
-  } : (!transform.any_op) -> (!transform.any_op, !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, %0 = transform.structured.pad %matmul_l1 {
+  %matmul_padded, %0, %copy_back = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     copy_back_op = "none"
-  } : (!transform.any_op) -> (!transform.any_op, !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>,
 
   //      CHECK: %[[T3:.*]] = tensor.pad %[[T0]] nofold
   //      CHECK: tensor.yield %[[CST]]
   //      CHECK: %[[T4:.*]] = tensor.pad %[[T1]] nofold
 
   //      CHECK: %[[T5:.*]] = linalg.matmul
   // CHECK-SAME:              ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>)
   // CHECK-SAME:              outs(%[[T2]] : tensor<4x5xf32>)
+
+  //      CHECK: %[[T6:.*]] = tensor.extract_slice %[[T5]]
+  //      CHECK: %[[T7:.*]] = bufferization.copy_tensor %[[T6]], %[[T2]]
   %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
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !transform.any_op, !transform.op<"bufferization.copy_tensor">)
+  // expected-remark @below {{1}}
+  test_print_number_of_associated_payload_ir_ops %copy_back : !transform.op<"bufferization.copy_tensor">
 }
 
 // -----
 
 #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
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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}}
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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}"}}
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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}}
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !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
-  %padded, %pad = transform.structured.pad %0 {
+  %padded, %pad, %copy_back = 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, !transform.any_op, !transform.any_op)
 }

mlir/test/Dialect/Linalg/transform-ops-invalid.mlir

 // RUN: mlir-opt %s --split-input-file --verify-diagnostics
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error@below {{'transform.structured.interchange' op expects iterator_interchange to be a permutation, found 1, 1}}
   transform.structured.interchange %arg0 iterator_interchange = [1, 1] : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error@below {{expects padding_dimensions to contain positive integers, found [1, -7]}}
-  transform.structured.pad %arg0 {padding_dimensions=[1, -7]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  transform.structured.pad %arg0 {padding_dimensions=[1, -7]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error@below {{expects pack_paddings to contain booleans (0/1), found [1, 7]}}
-  transform.structured.pad %arg0 {pack_paddings=[1, 7]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  transform.structured.pad %arg0 {pack_paddings=[1, 7]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error@below {{expects transpose_paddings to be a permutation, found [1, 1]}}
-  transform.structured.pad %arg0 {transpose_paddings=[[1, 1]]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  transform.structured.pad %arg0 {transpose_paddings=[[1, 1]]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error@below {{'transform.structured.interchange' op attribute 'iterator_interchange' failed to satisfy constraint: i64 dense array attribute whose value is non-negative}}
   transform.structured.interchange %arg0 iterator_interchange = [-3, 1] : (!transform.any_op) -> !transform.any_op

mlir/test/Dialect/Linalg/transform-ops.mlir

 // Check that operations are registered correctly through the extension
 // mechanism. Their syntax is generated and requires no additional testing since
 // we test the generator.
 //===----------------------------------------------------------------------===//
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !transform.any_op):
   // CHECK: transform.structured.pad
-  %0, %1 = transform.structured.pad %arg0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+  %0, %1, %2 = transform.structured.pad %arg0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !transform.any_op):
   // CHECK: transform.structured.interchange
   %0 = transform.structured.interchange %arg0 : (!transform.any_op) -> !transform.any_op
 }