[mlir] Add map_nested_foreach_thread_to_gpu_blocks op to transform dialect

[guraypp]

Summary

This revision adds a new op map_nested_foreach_thread_to_gpu_blocks to transform dialect.
If generate_gpu_launch argument is given, the op first generates gpu_launch. Otherwise, target must be gpu_launch. The op searches top level scf.foreach_threads inside the gpu_launch and distributes them with gpu.block_id attribute.
Loop mapping is explicit and given by the map_nested_foreach_thread_to_gpu_blocks op. Mapping is done one-to-one, therefore the loops disappear.
It also adds gpu dialect as dependent since the new op can create gpu::LaunchOp for given scf::ForeachThreadOp.

[nicolasvasilache]

It is debatable whether this transform should 1) create a gpu.launch from a scf.foreach_thread or 2) expect the scf.foreach_thread to be already nested within one.
The current implementation is 2) and is coming from IREE which has special constraints about what IR comes into existence when and how async is handled.

1. seems better suited to upstream MLIR. Instead of

// Step 6. Erase old op.
rewriter.eraseOp(foreachThreadOp);

you want to create a gpu::LaunchOp with the outlined body.
I.e. targetBlock = foreachThreadOp->getBlock(); should be the block of the gpu::LaunchOp.

Things will be synchronous for now and we will look into async launches later.

Most of the implementation should be factored out in a helper that we can reuse on the IREE side, as before.

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

1218

nit double comment

[nicolasvasilache]

Also, please be proactive in good doc comments everywhere.

[guraypp]

update

[nicolasvasilache]

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

136 ↗

(On Diff #461846)

This will break in various ways at link time when used from outside, please use a proper type / typedef.
E.g. something similar coming from scf:

LoopNest mlir::scf::buildLoopNest(
    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
    ValueRange steps, ValueRange iterArgs,
    function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
        bodyBuilder) {

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

135

nice!

[guraypp]

Remove template type for lambda and add a proper type

[csigg]

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

1465

I'm not familiar with transform ops, is this how this is done (rather than rewriter.notifyMatchFailure())?

1482–1483

Would it be cleaner to clone the region with a block arg mapping so that it can be rolled back?

1489

This looks like a left-over.

1594

Should there be some check somewhere that the current gridDims are 1 before rewriting them?

[nicolasvasilache]

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

1465

Transform ops are more directive, the equivalent of the pattern driver is the tranform IR itself where, if things fail, to apply we want a louder signal (at least for now).

1482–1483

Rollback is a thing only for DialectConversion today, we do not use that.
Transforms maintain operation handles, moving ops allows to keep the handles valid which is a behavior we want here.

1594

Yes, it would make sense to add an error at the top.

[guraypp]

Address csigg's comments

[guraypp]

rebase

[guraypp]

fix mistakes in alterGpuLaunch

[nicolasvasilache]

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

21

please drop spurious include (that I imagine was insert automatically by e.g. vscode ?)

24

please drop spurious include (that I imagine was insert automatically by e.g. vscode ?)

1404

almost never use unsigned unless doing bit manipulations, just use int64_t everywhere plz.

1404

Definitely use an OpBuilder::InsertionGuard here, never let a function leak a modification of insertion point

1405

Please don't use the magic constant 1 to detect whether or not to modify, pass an explicit Optional<int64_t>

1411

optional for less magic constant-y tests

1429

Note: it is good practice OpBuilder::InsertionGuard each time you modify insertion points in a scoped piece of code.
Here is does not matter much because it won't leak above, but it does matter a lot at function boundaries.

1435

0, 0, 0, feels fishy here, what is the intent ?
Use None if you want it to be optional.

1527

we use int64_t everywhere we can, please don't deviate from that

1541

Definitely use an OpBuilder::InsertionGuard here, never let a function leak a modification of insertion point

1564

Note: it is good practice OpBuilder::InsertionGuard each time you modify insertion points in a scoped piece of code.
Here is does not matter much because it won't leak above, but it does matter a lot at function boundaries.

1579

Note: it is good practice OpBuilder::InsertionGuard each time you modify insertion points in a scoped piece of code.
Here is does not matter much because it won't leak above, but it does matter a lot at function boundaries.

I would still use one here given the fact that this lambda may have a chance of turning into a standalone function in the future.

[guraypp]

address nicolasvasilache's comments

[ftynse]

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

876–877

Please do not use === that indicates the page title. Use the same-level header as the operation operations in this file.

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

140 ↗

(On Diff #462457)

What does "top-level" mean here? Is this going considering ascendants or descendants of the given operation?

136 ↗

(On Diff #461846)

Overall, when you need an immutable cheap view of a sequential construct such as vector, take an ArrayRef, when you need a mutable vector, take a SmallVectorImpl so it is compatible with any stack-size and not only the hardcoded/inferred one.

[nicolasvasilache]

Please address @ftynse 's comments and let's land this

[guraypp]

remove ==== in tablegen file

[guraypp]

rebase

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

let description = [{
     This transformation supports the following attributes:
       - `vectorize_padding`: a UnitAttr to activate the vectorization of 
       `tensor.pad` ops. Different pipelines may prefer to lower such ops to 
       loops.
       - `disable_multi_reduction_to_contract_patterns`: a UnitAttr to deactivate
       the rewrite of `vector.multi_reduction` to `vector.contract`. This is 
       intended to be used in tests only.
       - `disable_transfer_permutation_map_lowering_patterns`: a UnitAttr to
       deactivate the rewrite of `vector.transfer` with permutation maps into
       explicit `vector.transpose` operations. This is intended to be used in
       tests only but may be promotoed to a first class attribute in the future.
 
     #### Return modes:
     
     This operation produces `definiteFailure` if vectorization fails for any
     reason.
     The operation always returns the handle to the target op that is expected 
     to be isolated from above.
def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target, 
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def MapNestedForeachThreadToGpuBlocks : Op<Transform_Dialect, 
+    "structured.map_nested_foreach_thread_to_gpu_blocks",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformOpInterface,
+     TransformEachOpTrait]> {
+  let description = [{
+    Target the gpu_launch op and rewrite all scf.foreach_thread
+      to distributed gpu.block_id attribute.
+
+      The operation searches `scf.foreach_thread` ops nested under `target` 
+      and maps each such op to GPU thread blocks. Mapping is one-to-one and the
+      induction variables of `scf.foreach_thread` are rewritten to 
+      gpu.thread_id according to the gridDim attribute. 
+      
+      It requires a top-level `scf.foreach_thread`, siblings are not supported.
+
+      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.
+    
+      The operation alters the block size of the given gpu_launch using
+      gridDim argument.
+
+      Return modes:
+      =============
+      This operation ignores non-gpu_launch ops and drops them in the return.
+
+      If any scf.foreach_thread with tensors is found, the transform definitely 
+      fails.    
+
+      If all the scf.foreach_thread operations contained within the LaunchOp 
+      referred to by the `target` PDLOperation lower to GPU properly, the 
+      transform succeeds. Otherwise the transform definitely fails.
+      
+      The returned handle points to the same LaunchOp operand, consuming it and
+      producing a new SSA value to satisfy chaining and linearity of the IR 
+      properties.
+    ...
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$gridDim);
+  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::TransformState &state);
+  }];
+}
+
 #endif // LINALG_TRANSFORM_OPS

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

 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/StringSet.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
 /// Extracts a vector of unsigned from an array attribute. Asserts if the
 /// attribute contains values other than intergers. May truncate.
 static SmallVector<unsigned> extractUIntArray(ArrayAttr attr) {
   SmallVector<unsigned> result;
void transform::TileOp::getEffects(
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getDynamicSizes(), effects);
   producesHandle(getTiledLinalgOp(), effects);
   producesHandle(getLoops(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// MapNestedForeachThreadToGpuBlocks
+//===----------------------------------------------------------------------===//
+static LogicalResult rewriteTopLevelForeachThreadToGpuBlocks(RewriterBase &rewriter, 
+scf::ForeachThreadOp foreachThreadOp) {
+
+  if (foreachThreadOp->getNumResults() > 0)
+    return foreachThreadOp.emitOpError(
+        "only bufferized scf.foreach_thread lowers to gpu.block_id");
+  if (foreachThreadOp.getNumThreads().size() > 3)
+    return foreachThreadOp.emitOpError(
+        "scf.foreach_thread with rank > 3 does not lower to gpu.block_id");
+
+  // Step 0. Outline the compute workload region and set up the workload
+  // operands.
+  auto potentialGridDim = foreachThreadOp.getPermutedNumThreads(rewriter);
+  if (failed(potentialGridDim) ||
+      llvm::any_of(*potentialGridDim, [](OpFoldResult ofr) {
+        return !getConstantIntValue(ofr).has_value();
+      }))
+     return foreachThreadOp.emitOpError("unsupported dynamic gridDim size");
+
+  SmallVector<int64_t> gridDims =
+    llvm::to_vector(llvm::map_range(*potentialGridDim, [](OpFoldResult ofr) {
+      return getConstantIntValue(ofr).value();
+    }));
+
+  Location loc = foreachThreadOp.getLoc();
+  IndexType indexType = rewriter.getIndexType();
+  rewriter.setInsertionPoint(foreachThreadOp);
+  SmallVector<gpu::Dimension> gpuDims{gpu::Dimension::x, gpu::Dimension::y,
+                                      gpu::Dimension::z};
+  SmallVector<Value> blockOps;
+  for (int64_t idx : llvm::seq<int64_t>(0, gridDims.size())) {
+    blockOps.push_back(
+        rewriter.create<gpu::BlockIdOp>(loc, indexType, gpuDims[idx]));
+  }
+
+  // Step 3. Move the body of foreachThreadOp.
+  // Erase the terminator first, it will not be used since we are on buffers.
+  rewriter.eraseOp(foreachThreadOp.getTerminator());
+  Block *targetBlock;
+  Block::iterator insertionPoint;
+  targetBlock = foreachThreadOp->getBlock();
+  insertionPoint = Block::iterator(foreachThreadOp);
+  Block &sourceBlock = foreachThreadOp.getRegion().front();
+  targetBlock->getOperations().splice(insertionPoint,
+                                      sourceBlock.getOperations());
+
+  // // Step 4. RAUW thread indices to thread ops.
+  SmallVector<Value> threadIndices = *foreachThreadOp.getPermutedThreadIndices();
+  assert(blockOps.size() == 3 && "3 workgroup id ops are required");
+  assert(threadIndices.size() == 3 && "3 thread id dimensions are required");
+  for (auto it : llvm::zip(threadIndices, blockOps)) {
+    Value val = std::get<0>(it);
+    if (!val) continue;
+    for (Operation *user : llvm::make_early_inc_range(val.getUsers())) {
+      rewriter.updateRootInPlace(
+          user, [&]() { user->replaceUsesOfWith(val, std::get<1>(it)); });
+    }
+  }
+
+  // Step 6. Erase old op.
+  rewriter.eraseOp(foreachThreadOp);
+
+  return success();
+}
+
+
+DiagnosedSilenceableFailure
+transform::MapNestedForeachThreadToGpuBlocks::applyToOne(
+    Operation *target, SmallVectorImpl<Operation *> &results,
+    transform::TransformState &state) {
+
+  gpu::LaunchOp gpuLaunch = dyn_cast<gpu::LaunchOp>(target);
+  if (!gpuLaunch) {
+    target->emitError("Given target is not gpu.launch");
+    return DiagnosedSilenceableFailure::definiteFailure();
+  }
+
+  scf::ForeachThreadOp topLevelForeachThreadOp;
+  auto walkResult = target->walk([&](scf::ForeachThreadOp foreachThreadOp) {
+    if (foreachThreadOp->getParentOfType<scf::ForeachThreadOp>())
+      return WalkResult::advance();
+    if (topLevelForeachThreadOp) return WalkResult::interrupt();
+    topLevelForeachThreadOp = foreachThreadOp;
+    return WalkResult::advance();
+  });
+
+  if (walkResult.wasInterrupted()) {
+    target->emitOpError(
+        "could not find a unique topLevel scf.foreach_thread");
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+  }
+
+  SimpleRewriter rewriter(getContext());
+  rewriter.setInsertionPoint(target);
+
+  if (failed(rewriteTopLevelForeachThreadToGpuBlocks(rewriter, topLevelForeachThreadOp)))
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+  
+  results.assign({target});
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===----------------------------------------------------------------------===//
 // MapNestedForeachThreadToGpuThreads
 //===----------------------------------------------------------------------===//
 
 /// Searches `scf.foreach_thread` ops nested under `target` and maps each such
 /// op to GPU threads. Mapping is one-to-one and the induction variables of
 /// `scf.foreach_thread` are rewritten to gpu.thread_id according to the
 /// thread_dim_apping attribute. Sibling `scf.foreach_thread` are supported in
 /// which case, the union of the number of threads is computed and may result in
static LogicalResult alterGpuLaunchBlockDim(SimpleRewriter &rewriter,
                                             SmallVector<int64_t> blockDim) {
   gpu::KernelDim3 currentBlockdim = gpuLaunch.getBlockSizeOperandValues();
   if (blockDim[0] < 1 || blockDim[1] < 1 || blockDim[2] < 1) {
     gpuLaunch->emitError() << "Given blockDim(" << blockDim[0] << ","
                            << blockDim[1] << "," << blockDim[2]
                            << ") is invalid";
     return failure();
   }
   rewriter.setInsertionPointAfterValue(currentBlockdim.x);
   auto createBlockDimValue = [&](int64_t dim) {
     return rewriter.create<arith::ConstantIndexOp>(currentBlockdim.x.getLoc(),
                                                    dim);
   };
   gpuLaunch.blockSizeXMutable().assign(createBlockDimValue(blockDim[0]));
   gpuLaunch.blockSizeYMutable().assign(createBlockDimValue(blockDim[1]));
   gpuLaunch.blockSizeZMutable().assign(createBlockDimValue(blockDim[2]));
   return success();
 }
 
 DiagnosedSilenceableFailure
 transform::MapNestedForeachThreadToGpuThreads::applyToOne(
     Operation *target, SmallVectorImpl<Operation *> &results,
     transform::TransformState &state) {
 
   gpu::LaunchOp gpuLaunch = dyn_cast<gpu::LaunchOp>(target);
   if (!gpuLaunch) {
     target->emitError("Given target is not gpu.launch");
     return DiagnosedSilenceableFailure::definiteFailure();
   }
 
   SmallVector<int64_t> blockDim = extractFromI64ArrayAttr(getBlockDim());
   blockDim.resize(/*size=*/3, /*value=*/1);
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   auto walkResult = mlir::linalg::rewriteMapNestedForeachThreadToGpuThreads(
       rewriter, target, blockDim, true);
   if (walkResult.wasInterrupted())
     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
 
   LogicalResult result = alterGpuLaunchBlockDim(rewriter, gpuLaunch, blockDim);
   if (failed(result))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   results.assign({target});
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
if (targets.empty()) {
     return DiagnosedSilenceableFailure(success());
   }
 
   // Result payload ops.
   SmallVector<Operation *> tileOps;
   SmallVector<Operation *> tiledOps;
 
   // Given a list of OpFoldResults that are either index attrs or op handles,
   // return a list of OpFoldResults where all op handles are replaced with the
   // first (and only) OpResult of that payload op. (There must be exactly one
   // mapped payload op and it must have exactly one index result.)
   auto getOpResultsOrIndexAttrs =
       [&](SmallVector<OpFoldResult> &result,
           ArrayRef<OpFoldResult> opHandlesOrIndexAttrs) {
         for (OpFoldResult ofr : opHandlesOrIndexAttrs) {
           if (ofr.is<Attribute>()) {
             result.push_back(ofr);
             continue;
           }
           ArrayRef<Operation *> dynamicNumThreads =
               state.getPayloadOps(ofr.get<Value>());
           if (dynamicNumThreads.size() != 1) {
             DiagnosedSilenceableFailure diag =
                 emitSilenceableError()
                 << "handle must be mapped to exactly 1 payload op";
             diag.attachNote(ofr.get<Value>().getLoc())
                 << "mapped to " << dynamicNumThreads.size() << " ops";
             return diag;
           }
           Operation *op = dynamicNumThreads[0];
           if (op->getNumResults() != 1 ||
               !op->getResult(0).getType().isIndex()) {
             DiagnosedSilenceableFailure diag =
                 emitSilenceableError()
                 << "payload op must have exactly 1 index result";
             diag.attachNote(op->getLoc())
                 << "has " << op->getNumResults() << " results";
             return diag;
           }
           result.push_back(op->getResult(0));
         }
for (Operation *target : targets) {
     auto tilableOp = dyn_cast<TilingInterface>(target);
     if (!tilableOp) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "only TilingInterface ops are supported";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     rewriter.setInsertionPoint(tilableOp);
     auto maybeThreadDimMappingAttr = getThreadDimMapping();
     auto dimMapping = llvm::to_vector(
         maybeThreadDimMappingAttr
             ? extractFromI64ArrayAttr(*maybeThreadDimMappingAttr)
             : ArrayRef<int64_t>{});
 
     FailureOr<ForeachThreadTilingResult> tilingResult = failure();
     if (!getMixedNumThreads().empty()) {
       tilingResult = linalg::tileToForeachThreadOp(rewriter, tilableOp,
                                                    numThreads, dimMapping);
     } else {
       tilingResult = linalg::tileToForeachThreadOpUsingTileSizes(
           rewriter, tilableOp, tileSizes, dimMapping);
     }
 
     if (failed(tilingResult))
       return emitDefaultSilenceableFailure(tilableOp);
     rewriter.replaceOp(tilableOp, tilingResult->tileOp->getResults());
 
     tileOps.push_back(tilingResult->tileOp);
     tiledOps.push_back(tilingResult->tiledOp);
   }
 
   transformResults.set(getForeachThreadOp().cast<OpResult>(), tileOps);
   transformResults.set(getTiledOp().cast<OpResult>(), tiledOps);
 
   return DiagnosedSilenceableFailure(success());
 }
 
 void transform::TileToForeachThreadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getTileSizes(), effects);
   onlyReadsHandle(getNumThreads(), effects);
   producesHandle(getResults(), effects);
 }
 
 SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedNumThreads() {
   return getMixedSizes(getStaticNumThreads(), getNumThreads());
 }
 
 SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedTileSizes() {
   return getMixedSizes(getStaticTileSizes(), getTileSizes());
 }
 
 LogicalResult TileToForeachThreadOp::verify() {
   if (getMixedNumThreads().empty() == getMixedTileSizes().empty())
     return emitOpError("either num_threads or tile_sizes must be specified");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is an helper only to call vectorize via a pattern inside of
 /// VectorizeOp::applyToOne.
 struct VectorizationPattern : public RewritePattern {
   explicit VectorizationPattern(MLIRContext *context)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       return failure();
     return vectorize(rewriter, linalgOp);
   }
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {

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

 // RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
 
 !type = memref<2 x 32 x f32>
 !type1d = memref<32 x f32>
 
+// CHECK-LABEL: func.func @saxpy2dblock(
+// CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>
+// CHECK-SAME:    %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>
+// CHECK-SAME:    %[[ARGT:[0-9a-z]+]]: memref<32xf32>
+func.func @saxpy2dblock(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {
+  %one = arith.constant 1 : index
+  %c12 = arith.constant 12 : index
+  %c9 = arith.constant 9 : index
+  %c7 = arith.constant 7 : index
+//      CHECK:   gpu.launch
+//      CHECK:   %[[BLKX:.*]] = gpu.block_id  x
+//      CHECK:   %[[BLKY:.*]] = gpu.block_id  y
+//      CHECK:   memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]]]
+//      CHECK:   memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]]]
+  %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)
+            threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) 
+  {
+    scf.foreach_thread (%i, %j) in (%c7, %c9) {
+        %4 = memref.load %x[%i, %j] : !type        
+        %5 = memref.load %y[%i, %j] : !type
+        %6 = math.fma %alpha, %4, %5 : f32
+        memref.store %6, %y[%i, %j] : !type
+     }  {thread_dim_mapping = [0, 1, 2]}               
+    gpu.terminator
+  }
+  return %y : !type
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
+    transform.structured.map_nested_foreach_thread_to_gpu_blocks %funcop { blockDim = [12, 9, 1] }
+  }
+}
+
+// -----
+
+!type = memref<2 x 32 x f32>
+!type1d = memref<32 x f32>
+
 // CHECK-LABEL: func.func @saxpy2d(
 // CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGT:[0-9a-z]+]]: memref<32xf32>
 func.func @saxpy2d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {
   %one = arith.constant 1 : index
   %c12 = arith.constant 12 : index
   %c9 = arith.constant 9 : index

 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
     transform.structured.map_nested_foreach_thread_to_gpu_threads %funcop { blockDim = [12, 9, 1] }
   }
 }