[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

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

134

nice!

[guraypp]

Remove template type for lambda and add a proper type

[csigg]

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

1366

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

1383–1384

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

1390

This looks like a left-over.

1495

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

[nicolasvasilache]

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

1366

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).

1383–1384

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.

1495

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 ?)

1295

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

1297

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

1298

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

1300–1311

optional for less magic constant-y tests

1329

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.

1335–1336

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

1428

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

1442

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

1465

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.

1480

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

136

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.

140

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

[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

def MapNestedForeachThreadToGpuThreads :
   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 the top level `scf.foreach_thread`
+    to distributed gpu.block_id attribute. If `generate_gpu_launch` attribute 
+    is set, then first generates `gpu_launch` and moves the top level 
+    `scf.foreach_thread` inside.
+
+    The operation searches top level `scf.foreach_thread` ops under 
+    `gpu_launch` and maps each such op to GPU blocks. Mapping is 
+    one-to-one and the induction variables of `scf.foreach_thread` are 
+    rewritten to gpu.block_id according to the `thread_dim_apping` attribute.
+
+    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,
+                   UnitAttr:$generate_gpu_launch);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let description = [{
     Indicates that the given `target` op all the ops it contains should be
     vectorized with the configuration specified by the attributes of this op.
     This vectorization only handles structured ops that operate on shaped types
     and does not vectorize loops or straight-line. Internally, it applies a

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

 bool areElementwiseOpsFusable(OpOperand *fusedOperand);
 
 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
 FailureOr<Operation *> fuseElementwiseOps(RewriterBase &rewriter,
                                           OpOperand *fusedOperand);
 
+/// Maps the top level `scf.foreach_thread` op to GPU Thread Blocks. Mapping is
+/// one-to-one and the induction variables of `scf.foreach_thread` are rewritten
+/// to gpu.block_id according to the thread_dim_apping attribute. Dynamic,
+/// `scf.foreach_thread` trip counts are currently not supported. Dynamic block
+/// dim sizes are currently not supported.
+template <typename BlockIdGenerator>
+LogicalResult rewriteTopLevelForeachThreadToGpuBlocks(
+    RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
+    BlockIdGenerator blockIdGenerator, SmallVector<int64_t> &gridDims);
+
+/// Finds the top level scf::ForeachThreadOp of given target.
+FailureOr<scf::ForeachThreadOp> findTopLevelForeachThreadOp(Operation *target);
+
 /// 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
 /// predication. Dynamic, `scf.foreach_thread` trip counts are currently not
 /// supported. Dynamic block dim sizes are currently not supported.
 mlir::WalkResult rewriteMapNestedForeachThreadToGpuThreads(

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;
for (int64_t idx : llvm::seq<int64_t>(0, blockDim.size())) {
     threadOps.push_back(
         rewriter.create<gpu::ThreadIdOp>(loc, indexType, gpuDims[idx]));
   }
   // Step 2. Maybe create conditionals to predicate the region.
   Value predicate;
   for (auto [threadId, blockDim, globalBlockDim] :
        llvm::zip(threadOps, blockDim, globalBlockDims)) {
     if (blockDim > globalBlockDim) {
       return foreachThreadOp.emitOpError("blockDim size overflow: ")
              << blockDim << " > " << globalBlockDim;
     }
     if (blockDim == globalBlockDim)
       continue;
     Value tmpPredicate = rewriter.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::ult, threadId,
         rewriter.create<arith::ConstantIndexOp>(loc, blockDim));
     predicate =
auto walkResult = target->walk([&](scf::ForeachThreadOp foreachThreadOp) {
     if (failed(rewriteOneForeachThreadToGpuThreads(rewriter, foreachThreadOp,
                                                    blockDim, true)))
       return WalkResult::interrupt();
     return WalkResult::advance();
   });
   return walkResult;
 }
 
-// Alter blockDim of the given kernel
-static LogicalResult alterGpuLaunchBlockDim(SimpleRewriter &rewriter,
-                                            gpu::LaunchOp gpuLaunch,
-                                            SmallVector<int64_t> blockDim) {
+/// Alter grid or block dimensions of the given kernel
+static LogicalResult alterGpuLaunch(SimpleRewriter &rewriter,
+                                    gpu::LaunchOp gpuLaunch, int gridDimX = 0,
+                                    int gridDimY = 0, int gridDimZ = 0,
+                                    int blockDimX = 0, int blockDimY = 0,
+                                    int blockDimZ = 0) {
   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) {
+  auto createGridDimValue = [&](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]));
+  if (gridDimX)
+    gpuLaunch.gridSizeXMutable().assign(createGridDimValue(gridDimX));
+  if (gridDimY)
+    gpuLaunch.gridSizeYMutable().assign(createGridDimValue(gridDimY));
+  if (gridDimZ)
+    gpuLaunch.gridSizeZMutable().assign(createGridDimValue(gridDimZ));
+  if (blockDimX)
+    gpuLaunch.gridSizeXMutable().assign(createGridDimValue(blockDimX));
+  if (blockDimY)
+    gpuLaunch.gridSizeYMutable().assign(createGridDimValue(blockDimY));
+  if (blockDimZ)
+    gpuLaunch.gridSizeZMutable().assign(createGridDimValue(blockDimZ));
   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);
+  LogicalResult result = alterGpuLaunch(rewriter, gpuLaunch, 0, 0, 0,
+                                        blockDim[0], blockDim[1], blockDim[2]);
   if (failed(result))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   results.assign({target});
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
+// MapNestedForeachThreadToGpuBlocks
+//===----------------------------------------------------------------------===//
+
+template <typename BlockIdGenerator>
+LogicalResult mlir::linalg::rewriteTopLevelForeachThreadToGpuBlocks(
+    RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
+    BlockIdGenerator blockIdGenerator, SmallVector<int64_t> &gridDims) {
+  if (foreachThreadOp.getNumResults() > 0)
+    return foreachThreadOp->emitError(
+        "only bufferized scf.foreach_thread lowers to gpu.block_id");
+  if (foreachThreadOp.getNumThreads().size() > 3)
+    return foreachThreadOp->emitError(
+        "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->emitError("unsupported dynamic gridDim");
+
+  for (OpFoldResult ofr : *potentialGridDim)
+    gridDims.push_back(getConstantIntValue(ofr).value());
+
+  IndexType indexType = rewriter.getIndexType();
+  SmallVector<Value> blockOps;
+  blockIdGenerator(foreachThreadOp, gridDims, indexType, blockOps);
+
+  // Step 1. 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 2. RAUW thread indices to thread ops.
+  SmallVector<Value> threadIndices =
+      *foreachThreadOp.getPermutedThreadIndices();
+  assert(blockOps.size() == 3 && "3 block 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 3. Erase old op.
+  rewriter.eraseOp(foreachThreadOp);
+
+  return success();
+}
+
+FailureOr<scf::ForeachThreadOp>
+mlir::linalg::findTopLevelForeachThreadOp(Operation *target) {
+  scf::ForeachThreadOp topLevelForeachThreadOp;
+  auto walkResult = target->walk([&](scf::ForeachThreadOp foreachThreadOp) {
+    if (foreachThreadOp->getParentOfType<scf::ForeachThreadOp>())
+      return WalkResult::advance();
+    if (topLevelForeachThreadOp)
+      // TODO Handle multiple foreach if there is no dependences between them
+      return WalkResult::interrupt();
+    topLevelForeachThreadOp = foreachThreadOp;
+    return WalkResult::advance();
+  });
+
+  if (walkResult.wasInterrupted())
+    return target->emitError(
+        "could not find a unique topLevel scf.foreach_thread");
+
+  return topLevelForeachThreadOp;
+}
+
+/// Create gpuLauncOp with given kernel configurations
+static gpu::LaunchOp createGpuLaunch(RewriterBase &rewriter, Location loc,
+                                     int gridDimX = 1, int gridDimY = 1,
+                                     int gridDimZ = 1, int blockDimX = 1,
+                                     int blockDimY = 1, int blockDimZ = 1) {
+  auto createConstant = [&](int dim) {
+    return rewriter.create<arith::ConstantIndexOp>(loc, dim);
+  };
+  Value gridSizeX = createConstant(gridDimX);
+  Value gridSizeY = createConstant(gridDimY);
+  Value gridSizeZ = createConstant(gridDimZ);
+  Value blockSizeX = createConstant(blockDimX);
+  Value blockSizeY = createConstant(blockDimY);
+  Value blockSizeZ = createConstant(blockDimZ);
+  auto launchOp = rewriter.create<gpu::LaunchOp>(
+      loc, gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ);
+  rewriter.setInsertionPointToEnd(&launchOp.body().front());
+  rewriter.create<gpu::TerminatorOp>(loc);
+  return launchOp;
+}
+
+DiagnosedSilenceableFailure
+transform::MapNestedForeachThreadToGpuBlocks::applyToOne(
+    Operation *target, SmallVectorImpl<Operation *> &results,
+    transform::TransformState &state) {
+  gpu::LaunchOp gpuLaunch = dyn_cast<gpu::LaunchOp>(target);
+  SimpleRewriter rewriter(getContext());
+
+  if (!getGenerateGpuLaunch() && !gpuLaunch) {
+    target->emitError("Given target is not gpu.launch, set "
+                      "`generate_gpu_launch` attribute");
+    return DiagnosedSilenceableFailure::definiteFailure();
+  }
+
+  auto res = mlir::linalg::findTopLevelForeachThreadOp(target);
+  if (failed(res))
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+
+  scf::ForeachThreadOp topLevelForeachThreadOp = *res;
+  rewriter.setInsertionPoint(topLevelForeachThreadOp);
+
+  // Generate gpu launch here and move the foreach_thread inside
+  if (getGenerateGpuLaunch()) {
+    gpuLaunch = createGpuLaunch(rewriter, target->getLoc());
+    rewriter.setInsertionPointToStart(&gpuLaunch.body().front());
+    Operation *newForeachThreadOp = rewriter.clone(*topLevelForeachThreadOp);
+    rewriter.eraseOp(topLevelForeachThreadOp);
+    topLevelForeachThreadOp =
+        dyn_cast<scf::ForeachThreadOp>(newForeachThreadOp);
+  }
+
+  auto generateBlocks = [&](Operation *op, const SmallVector<int64_t> &gridDims,
+                            IndexType indexType, SmallVector<Value> &blockOps) {
+    Location loc = op->getLoc();
+    rewriter.setInsertionPoint(op);
+    SmallVector<gpu::Dimension> gpuDims{gpu::Dimension::x, gpu::Dimension::y,
+                                        gpu::Dimension::z};
+    for (int64_t idx : llvm::seq<int64_t>(0, gridDims.size())) {
+      blockOps.push_back(
+          rewriter.create<gpu::BlockIdOp>(loc, indexType, gpuDims[idx]));
+    }
+  };
+
+  SmallVector<int64_t> gridDim = extractFromI64ArrayAttr(getGridDim());
+  if (failed(mlir::linalg::rewriteTopLevelForeachThreadToGpuBlocks(
+          rewriter, topLevelForeachThreadOp, generateBlocks, gridDim)))
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+
+  if (failed(alterGpuLaunch(rewriter, gpuLaunch, gridDim[0], gridDim[1],
+                            gridDim[2])))
+    return DiagnosedSilenceableFailure::definiteFailure();
+
+  results.assign({gpuLaunch});
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===----------------------------------------------------------------------===//
 // TileToForeachThreadOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
     transform::TransformResults &transformResults,
     transform::TransformState &state) {
   IRRewriter rewriter(getContext());
   ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
public:
 
   void init() {
     declareDependentDialect<pdl::PDLDialect>();
     declareDependentDialect<LinalgDialect>();
     declareGeneratedDialect<AffineDialect>();
     declareGeneratedDialect<arith::ArithmeticDialect>();
     declareGeneratedDialect<scf::SCFDialect>();
     declareGeneratedDialect<vector::VectorDialect>();
+    declareGeneratedDialect<gpu::GPUDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace

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 {
+  %c9 = arith.constant 9 : index
+  %c7 = arith.constant 7 : index
+  %one = arith.constant 1 : 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

 ^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] }
   }
 }
 
+// -----
+
+!type4d = memref<32x64x4x32xf32>
+
+// CHECK-LABEL: func.func @saxpy4d(
+// CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<32x64x4x32xf32>
+// CHECK-SAME:    %[[ARGY:[0-9a-z]+]]: memref<32x64x4x32xf32>
+func.func @saxpy4d(%x: !type4d, %y: !type4d, %alpha : f32) -> !type4d {
+  %c32 = arith.constant 32 : index  
+  %c64 = arith.constant 64 : index  
+  %c4 = arith.constant 4 : index  
+//      CHECK:   gpu.launch
+//      CHECK:   %[[BLKX:.*]] = gpu.block_id  x
+//      CHECK:   %[[BLKY:.*]] = gpu.block_id  y
+//      CHECK:   %[[TIDX:.*]] = gpu.thread_id  x
+//      CHECK:   %[[TIDY:.*]] = gpu.thread_id  y
+//      CHECK:   memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]]
+//      CHECK:   memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]]
+  scf.foreach_thread (%i, %j) in (%c32, %c64) {
+    scf.foreach_thread (%k, %l) in (%c4, %c32) {
+      %4 = memref.load %x[%i, %j, %k, %l] : !type4d        
+      %5 = memref.load %y[%i, %j, %k, %l] : !type4d
+      %6 = math.fma %alpha, %4, %5 : f32
+      memref.store %6, %y[%i, %j, %k, %l] : !type4d
+    }  {thread_dim_mapping = [1, 0, 2]}
+  }  {thread_dim_mapping = [0, 1, 2]}
+  return %y : !type4d
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %funcop = transform.structured.match ops{["func.func"]} in %arg0
+    %gpuLaunch = transform.structured.map_nested_foreach_thread_to_gpu_blocks %funcop { generate_gpu_launch }
+    transform.structured.map_nested_foreach_thread_to_gpu_threads %gpuLaunch { blockDim = [32, 4, 1] }
+  }
+}