[mlir][linalg] MapCopyToThreadsOp: Support tensor.pad

[springerm]

Summary

Depends On: D155946

[nicolasvasilache]

nice!

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

 }
 
 //===----------------------------------------------------------------------===//
 // MapCopyToThreadsOp
 //===----------------------------------------------------------------------===//
 
 def MapCopyToThreadsOp :
   Op<Transform_Dialect, "structured.gpu.map_copy_to_threads",
-    [FunctionalStyleTransformOpTrait, 
+    [FunctionalStyleTransformOpTrait,
      MemoryEffectsOpInterface,
-     TransformEachOpTrait, 
+     TransformEachOpTrait,
      TransformOpInterface]> {
   let description = [{
-    Targeted mapping of a copy operation on tensors to a GPU thread mapping.
+    Targeted mapping of a linalg.copy / tensor.pad operation on tensors to a GPU
+    thread mapping.
 
-    This operation implements a greedy heuristic that determines a good 
-    distribution of threads to break down the copy operation into.
-    The heuristic is driven by considerations related to the underlying 
+    This operation implements a greedy heuristic that determines a good
+    distribution of threads to break down the copy/pad operation into.
+    The heuristic is driven by considerations related to the underlying
     architecture for which good high-level decisions are needed assuming certain
     hardware features. Relevant features are exposed via first-class attributes
     to control the behavior of the transformation at a high level.
 
     For now, a single heuristic is implemented and can be extended on a per-need
     basis.
 
-    #### Return modes:
+    #### Return modes
 
-    The operation always succeeds and returns a handle to the relevant tiled
-    linalg.copy op.
+    This operation fails definitely if there is an unsupported op (i.e., not
+    linalg.copy / tensor.pad) among the targeted op. Otherwise, the operation
+    always succeeds and returns a handle to the relevant tiled linalg.copy /
+    tensor.pad op and the enclosing scf.forall op.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        I64Attr:$total_num_threads,
                        I64Attr:$desired_bit_alignment);
-  let results = (outs TransformHandleTypeInterface:$transformed);
+  let results = (outs TransformHandleTypeInterface:$forall_op,
+                      TransformHandleTypeInterface:$tiled_op);
 
   let assemblyFormat = [{
     $target
     `total_num_threads` `=` $total_num_threads
     `desired_bit_alignment` `=` $desired_bit_alignment
-    attr-dict 
+    attr-dict
     `:` functional-type(operands, results)
   }];
 
   let builders = [
     OpBuilder<(ins "Value":$target)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformRewriter &rewriter,
-        ::mlir::linalg::CopyOp copyOp,
+        ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
 
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedNumThreads();
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

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

DiagnosedSilenceableFailure diag =
       << "only InsertSliceOp and ParallelInsertSliceOp ops are supported";
   diag.attachNote(targetOp->getLoc()) << "target op";
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // MapCopyToThreadsOp
 //===----------------------------------------------------------------------===//
+
 DiagnosedSilenceableFailure transform::MapCopyToThreadsOp::applyToOne(
-    transform::TransformRewriter &rewriter, linalg::CopyOp copyOp,
+    transform::TransformRewriter &rewriter, Operation *target,
     transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
-  auto transformOp = cast<TransformOpInterface>(getOperation());
-  ShapedType resultShapedType;
-  if (copyOp) {
-    resultShapedType =
-        cast<ShapedType>(copyOp.getDpsInitOperand(0)->get().getType());
+  // Check if the op is supported.
+  if (!isa<linalg::CopyOp, tensor::PadOp>(target)) {
+    DiagnosedSilenceableFailure diag =
+        emitSilenceableError()
+        << "only linalg.copy and tensor.pad target ops are supported";
+    diag.attachNote(target->getLoc()) << "target op";
+    return diag;
   }
-  if (!copyOp || !resultShapedType.hasStaticShape()) {
+  assert(target->getNumResults() == 1 && "expected single result");
+  auto resultShapedType = cast<ShapedType>(target->getResult(0).getType());
+  if (!resultShapedType.hasStaticShape()) {
     DiagnosedSilenceableFailure diag =
-        transformOp.emitSilenceableError()
-        << "only statically sized linalg.copy ops of rank <= 3 are supported";
-    diag.attachNote(copyOp->getLoc()) << "target op";
+        emitSilenceableError()
+        << "only statically sized ops of rank <= 3 are supported";
+    diag.attachNote(target->getLoc()) << "target op";
     return diag;
   }
 
   // Conservatively set the minimum viable desired bitwidth alignment.
   int64_t desiredBitAlignment = getDesiredBitAlignment();
   int64_t eltBitwidth =
       resultShapedType.getElementType().getIntOrFloatBitWidth();
   if (desiredBitAlignment % eltBitwidth != 0) {
     desiredBitAlignment = eltBitwidth;
   }
 
   gpu::CopyMappingInfo mapping(
       /*ctx=*/getContext(),
       /*totalNumThreads=*/getTotalNumThreads(),
       /*alignment=*/desiredBitAlignment,
       /*sizes=*/resultShapedType.getShape(),
       /*favorPredication=*/false,
       /*elementalBitwidth=*/
       resultShapedType.getElementType().getIntOrFloatBitWidth());
   if (mapping.status == gpu::CopyMappingInfo::Status::Invalid) {
     DiagnosedSilenceableFailure diag =
-        transformOp.emitSilenceableError()
+        emitSilenceableError()
         << "too few threads to map copy op to threads on the most minor "
            "dimension, given alignment and vector size constraints, try "
            "smaller tile size of mapping to more threads";
-    diag.attachNote(copyOp->getLoc()) << "target op";
+    diag.attachNote(target->getLoc()) << "target op";
     return diag;
   }
 
   // OpBuilder only used to compute attributes.
   OpBuilder b(getContext());
   linalg::ForallTilingResult tilingResult;
   DiagnosedSilenceableFailure diag = tileToForallOpImpl(
       /*rewriter=*/rewriter,
       /*state=*/state,
-      /*transformOp=*/transformOp,
-      /*target=*/copyOp,
+      /*transformOp=*/*this,
+      /*target=*/target,
       /*mixedNumThreads=*/getMixedValues(mapping.numThreads, {}, b),
       /*mixedTileSizes=*/ArrayRef<OpFoldResult>{},
       /*mapping=*/b.getArrayAttr(mapping.threadMapping),
       /*tilingResult=*/tilingResult);
   if (!diag.succeeded())
     return diag;
 
+  results.push_back(tilingResult.tileOp);
   results.push_back(tilingResult.tiledOp);
   return DiagnosedSilenceableFailure::success();
 }
 
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOpsEnums.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"

mlir/test/Dialect/Linalg/transform-op-gpu-map-copy-to-threads.mlir

 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
+}
+
+// -----
+
+!tt = tensor<8xf16>
+!tin = tensor<?xf16>
+
+// CHECK-LABEL: func @pad_1d_8xf16
+func.func @pad_1d_8xf16(%t0: !tin, %sz: index) -> !tt {
+  %cst = arith.constant 0.0 : f16
+  /// Too little data for all threads, needs predication, while keeping most
+  /// minor transfer size -> 1 thread.
+  // CHECK: scf.forall {{.*}} in (1) {{.*}}
+  // CHECK:   %[[padded:.*]] = tensor.pad {{.*}}
+  // CHECK:   tensor.cast %[[padded]] : tensor<?xf16> to tensor<8xf16>
+  // CHECK: {mapping = [#gpu.linear<x>]}
+  %0 = tensor.pad %t0 low[0] high[%sz] {
+  ^bb0(%arg0: index):
+    tensor.yield %cst : f16
+  } : !tin to !tt
+  return %0 : !tt
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
+    : (!transform.any_op) -> !transform.any_op
+  transform.structured.gpu.map_copy_to_threads %0
+    total_num_threads = 32 desired_bit_alignment = 128
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"tensor.pad">)
 }
 
 // -----
 
 !tt = tensor<16xf16>
 
 // CHECK-LABEL: func @copy_1d_16xf16
 func.func @copy_1d_16xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Too little data for all threads, needs predication, while keeping most
   /// minor transfer size -> 2 threads.
   // CHECK: scf.forall {{.*}} in (2) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<8xf16>
   // CHECK: {mapping = [#gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<20xf16>
 
 // CHECK-LABEL: func @copy_1d_20xf16
 func.func @copy_1d_20xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Too little data for all threads, needs predication, while keeping most
   /// minor transfer size -> 5 threads.
   // CHECK: scf.forall {{.*}} in (5) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<4xf16>
   // CHECK: {mapping = [#gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 
 // -----
 
 !tt = tensor<20xf16>
 
 // CHECK-LABEL: func @copy_1d_20xf16
 func.func @copy_1d_20xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Too little data for all threads, needs predication, while keeping most
   /// minor transfer size -> 5 threads.
   // CHECK: scf.forall {{.*}} in (5) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<4xf16>
   // CHECK: {mapping = [#gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<128xf16>
 
 // CHECK-LABEL: func @copy_1d_128xf16
 func.func @copy_1d_128xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Enough data for all threads and no need for predication but we must reduce
   /// the transfer size to 4xf16.
   // CHECK: scf.forall {{.*}} in (32) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<4xf16>
   // CHECK: {mapping = [#gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<256xf16>
 
 // CHECK-LABEL: func @copy_1d_256xf16
 func.func @copy_1d_256xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Enough data for all threads and no need for predication.
   // CHECK: scf.forall {{.*}} in (32) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<8xf16>
   // CHECK: {mapping = [#gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<16x32x64xi8>
 
 // CHECK-LABEL: func @copy_3d_16x32x64xi8
 func.func @copy_3d_16x32x64xi8(%t0: !tt, %out: !tt) -> !tt {
   // CHECK: scf.forall {{.*}} in (1, 8, 4) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<16x4x16xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<16x32x64xi8>
 
 // CHECK-LABEL: func @copy_3d_16x32x64xi8
 func.func @copy_3d_16x32x64xi8(%t0: !tt, %out: !tt) -> !tt {
   // CHECK: scf.forall {{.*}} in (1, 4, 8) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<16x8x8xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 64
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<4x8x16xi8>
 
 // CHECK-LABEL: func @copy_3d_4x8x16xi8
 func.func @copy_3d_4x8x16xi8(%t0: !tt, %out: !tt) -> !tt {
   // CHECK: scf.forall {{.*}} in (4, 8, 1) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<1x1x16xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<4x8x16xi8>
 
 // CHECK-LABEL: func @copy_3d_4x8x16xi8
 func.func @copy_3d_4x8x16xi8(%t0: !tt, %out: !tt) -> !tt {
   // CHECK: scf.forall {{.*}} in (1, 2, 16) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<4x4x1xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 8
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<3x5x7xi8>
 
 // CHECK-LABEL: func @copy_3d_3x5x7xi8
 func.func @copy_3d_3x5x7xi8(%t0: !tt, %out: !tt) -> !tt {

 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 8
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<16x15x5xi8>
 
 // CHECK-LABEL: func @copy_3d_16x15x5xi8
 func.func @copy_3d_16x15x5xi8(%t0: !tt, %out: !tt) -> !tt {
   // DP mapping: 5 mandated most minor, then 3 to allow 8 on the outermost.
   // CHECK: scf.forall {{.*}} in (8, 3, 5) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<2x5x1xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 128 desired_bit_alignment = 8
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<16x15x40xi8>
 
 // CHECK-LABEL: func @copy_3d_16x15x40xi8
 func.func @copy_3d_16x15x40xi8(%t0: !tt, %out: !tt) -> !tt {
   // DP mapping: 5 mandated most minor, then 3 to allow 8 on the outermost.
   // CHECK: scf.forall {{.*}} in (8, 3, 5) {{.*}}
   // CHECK:   linalg.copy {{.*}} -> tensor<2x5x8xi8>
   // CHECK: {mapping = [#gpu.linear<z>, #gpu.linear<y>, #gpu.linear<x>]}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 128 desired_bit_alignment = 64
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 
 ////////////////////////////////////////////////////////////////////////////////
 // Tests below are expected to fail.
 ////////////////////////////////////////////////////////////////////////////////
 
 // -----

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{too few threads to map copy op to threads on the most minor dimension, given alignment and vector size constraints}}
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<257xf16>
 
 // NO-CHECK-LABEL-ON-EXPECTED-ERROR
 func.func @copy_1d_257xf16(%t0: !tt, %out: !tt) -> !tt {
   /// Too much data for all threads, we do not try to recover here, this is the
   /// job of higher-level transformations to select better tile sizes and number
   /// of threads.
   
   // expected-note @below {{target op}}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{too few threads to map copy op to threads on the most minor dimension, given alignment and vector size constraints}}
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 128
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<512xi8>
 
 // NO-CHECK-LABEL-ON-EXPECTED-ERROR
 func.func @copy_1d_512xi8(%t0: !tt, %out: !tt) -> !tt {
   /// Too much data for all threads given the forced alignment to 8b, 
   /// we do not try to recover here, this is the job of higher-level 
   /// transformations to select better tile sizes and number of threads.
   // expected-note @below {{target op}}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{too few threads to map copy op to threads on the most minor dimension, given alignment and vector size constraints}}
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 8
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }
 
 // -----
 
 !tt = tensor<16x32x64xi8>
 
 // NO-CHECK-LABEL-ON-EXPECTED-ERROR
 func.func @copy_3d_16x32x64xi8(%t0: !tt, %out: !tt) -> !tt {
   /// Too much data for all threads given the forced alignment to 8b, 
   /// we do not try to recover here, this is the job of higher-level 
   /// transformations to select better tile sizes and number of threads.
   // expected-note @below {{target op}}
   %0 = linalg.copy ins(%t0: !tt) outs(%out: !tt) -> !tt 
   return %0 : !tt
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{too few threads to map copy op to threads on the most minor dimension, given alignment and vector size constraints}}
   transform.structured.gpu.map_copy_to_threads %0 
     total_num_threads = 32 desired_bit_alignment = 8
-      : (!transform.any_op) -> (!transform.op<"linalg.copy">)
+      : (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.copy">)
 }