[mlir][transform] Decouple GPUDeviceMapping attribute from the GPU transfrom dialect code generator

[guraypp]

Summary

DeviceMappingAttrInterface is implemented as unifiying mechanism for thread mapping. A code generator could use any attribute that implements this interface to lower scf.foreach_thread to device specific code. It is allowed to choose its own mapping and interpretation.

Currently, GPU transform dialect supports only GPUThreadMapping and GPUBlockMapping; however, other mappings should to be supported as well. This change addresses this issue. It decouples gpu transform dialect from the GPUThreadMapping and GPUBlockMapping. Now, they can work any other mapping.

[guraypp]

remove accidentally edited lines

[nicolasvasilache]

Thanks!

[guraypp]

rebase

mlir/include/mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td

 def DimZ : I64EnumAttrCase<"DimZ", 2, "z">;
 
 def ThreadsEnum : I64EnumAttr<"Threads", "threads for loop mapping", [
     DimX, DimY, DimZ]> {
   let cppNamespace = "::mlir::gpu";
 }
 
 def GPUThreadMappingAttr 
-    : GPU_Attr<"GPUThreadMapping", "thread", [ DeviceMappingAttrInterface ]> {
+    : GPU_Attr<"GPUThreadMapping", "thread", [ 
+      DeclareAttrInterfaceMethods<DeviceMappingAttrInterface> ]> {
   let parameters = (ins
     EnumParameter<ThreadsEnum>:$thread
   );
   let assemblyFormat = "`<` params `>`";
   let description = [{
     An attribute that allows defining thread parallelism for GPU devices.
 
     Thread (aka work item) are grouped into a thread blocks where block may be 
     described by a 1-, 2-, or 3-dimensional rectangle. This attribute indicates 
     that thread parallelism is desired. It can be consumed by lowering to 
     generate GPU.
   }];
 }
 
 def BlocksEnum : I64EnumAttr<"Blocks", "threads for loop mapping", [
     DimX, DimY, DimZ]> {
   let cppNamespace = "::mlir::gpu";
 }
 
-def GPUBlockMappingAttr : GPU_Attr<"GPUBlockMapping", "block", [ DeviceMappingAttrInterface ] >  {
+def GPUBlockMappingAttr : GPU_Attr<"GPUBlockMapping", "block", [ 
+  DeclareAttrInterfaceMethods<DeviceMappingAttrInterface> ] >  {
   let parameters = (ins
     EnumParameter<BlocksEnum>:$block
   );
   let assemblyFormat = "`<` params `>`";
   let description = [{
     An attribute that allows defining thread block parallelism for GPU devices.
 
     Thread blocks (aka work-group) are grouped into a grid where grid may be 
     described by a 1-, 2-, or 3-dimensional rectangle. This attribute indicates 
     that thread block parallelism is desired. It can be consumed by lowering to
     generate GPU code.
   }];
 }
 
 #endif // GPU_DEVICE_MAPPING_ATTR

mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h

 
 /// 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.
-DiagnosedSilenceableFailure
-mapNestedForeachToThreadsImpl(RewriterBase &rewriter, Operation *target,
-                              const SmallVectorImpl<int64_t> &blockDim,
-                              bool syncAfterDistribute,
-                              llvm::Optional<TransformOpInterface> transformOp);
+DiagnosedSilenceableFailure mapNestedForeachToThreadsImpl(
+    RewriterBase &rewriter, Operation *target,
+    const SmallVectorImpl<int64_t> &blockDim, bool syncAfterDistribute,
+    llvm::Optional<TransformOpInterface> transformOp,
+    const ArrayRef<DeviceMappingAttrInterface> &threadMappingAttributes);
 
 /// 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.
 DiagnosedSilenceableFailure mapForeachToBlocksImpl(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
     function_ref<void(RewriterBase &, scf::ForeachThreadOp,
                       SmallVectorImpl<Value> &)>
         blockIdGenerator,
-    SmallVectorImpl<int64_t> &gridDims, TransformOpInterface transformOp);
+    SmallVectorImpl<int64_t> &gridDims, TransformOpInterface transformOp,
+    const ArrayRef<DeviceMappingAttrInterface> &mappingAttributes);
 
 /// Finds the top level scf::ForeachThreadOp of given target.
 DiagnosedSilenceableFailure
 findTopLevelForeachThreadOp(Operation *target,
                             scf::ForeachThreadOp &topLevelForeachThreadOp,
                             TransformOpInterface transformOp);
 
 } // namespace gpu
 } // namespace transform
 
 namespace gpu {
 void registerTransformDialectExtension(DialectRegistry &registry);
 } // namespace gpu
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_TRANSFORMOPS_GPUTRANSFORMOPS_H

mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td

let description = [{
     mappings, and those mappings must implement this interface. These mappings 
     can be used by the device-specific code generators and the desired regions 
     can be connected to the given processing unit.
     
     Currently, `scf.foreach_thread` uses this interface to express the mapping 
     of the loops it contains to the GPU's parallelism units such as threads and 
     thread blocks.
   }];
+
+ let methods = [
+    InterfaceMethod<[{
+        Returns mapping as an integer from the attribute.
+      }],
+      "int64_t", "getMappingId", (ins)
+    >
+  ];
 }
 
 def DeviceMappingArrayAttr : 
   TypedArrayAttrBase<DeviceMappingAttrInterface, 
   "Device Mapping array attribute"> { }
 
 #endif // MLIR_DEVICEMAPPINGINTERFACE

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// GPU Device Mapping Attributes
+//===----------------------------------------------------------------------===//
+
+int64_t GPUBlockMappingAttr::getMappingId() const {
+  return static_cast<int64_t>(getBlock());
+}
+
+int64_t GPUThreadMappingAttr::getMappingId() const {
+  return static_cast<int64_t>(getThread());
+}
+
+//===----------------------------------------------------------------------===//
 // MMAMatrixType
 //===----------------------------------------------------------------------===//
 
 MMAMatrixType MMAMatrixType::get(ArrayRef<int64_t> shape, Type elementType,
                                  StringRef operand) {
   return Base::get(elementType.getContext(), shape, elementType, operand);
 }
 

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

 //===- GPUTransformOps.cpp - Implementation of GPU transform ops ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 using namespace mlir::transform;
 
 namespace {
 /// A simple pattern rewriter that implements no special logic.
 class SimpleRewriter : public PatternRewriter {
 public:
   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
 };
 } // namespace
 
+/// Check if given mapping attributes are one of the desired attributes
+static DiagnosedSilenceableFailure checkAttributeType(
+    const ArrayRef<DeviceMappingAttrInterface> &threadMappingAttributes,
+    const Optional<ArrayAttr> &foreachMapping,
+    llvm::Optional<TransformOpInterface> transformOp) {
+  if (!foreachMapping.has_value())
+    return transformOp->emitSilenceableError() << "mapping must be present";
+
+  if (llvm::any_of(foreachMapping->getValue(),
+                   [&](DeviceMappingAttrInterface map) {
+                     return llvm::find(threadMappingAttributes, map) ==
+                            threadMappingAttributes.end();
+                   }))
+    return transformOp->emitDefiniteFailure()
+           << "mapping must be one of " << threadMappingAttributes;
+  return DiagnosedSilenceableFailure::success();
+}
+
 /// Determines if the size of the kernel configuration is supported by the GPU
 /// architecture being used. It presently makes use of CUDA limitations, however
 /// that aspect may be enhanced for other GPUs.
 static DiagnosedSilenceableFailure
 checkGpuLimits(TransformOpInterface transformOp, Optional<int64_t> gridDimX,
                Optional<int64_t> gridDimY, Optional<int64_t> gridDimZ,
                Optional<int64_t> blockDimX, Optional<int64_t> blockDimY,
                Optional<int64_t> blockDimZ) {

 // MapForeachToBlocks
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure mlir::transform::gpu::mapForeachToBlocksImpl(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
     function_ref<void(RewriterBase &, scf::ForeachThreadOp,
                       SmallVectorImpl<Value> &)>
         blockIdGenerator,
-    SmallVectorImpl<int64_t> &gridDims, TransformOpInterface transformOp) {
+    SmallVectorImpl<int64_t> &gridDims, TransformOpInterface transformOp,
+    const ArrayRef<DeviceMappingAttrInterface> &mappingAttributes) {
   // Step 0. Target-specific verifications. There is no good place to anchor
   // those right now: the ForeachThreadOp is target-independent and the
   // transform op does not apply to individual ForeachThreadOp.
-  MLIRContext *ctx = foreachThreadOp->getContext();
   Location loc = foreachThreadOp->getLoc();
-  Attribute bX = GPUBlockMappingAttr::get(ctx, Blocks::DimX);
-  Attribute bY = GPUBlockMappingAttr::get(ctx, Blocks::DimY);
-  Attribute bZ = GPUBlockMappingAttr::get(ctx, Blocks::DimZ);
+
   if (foreachThreadOp.getNumResults() > 0)
     return transformOp.emitSilenceableError()
            << "only bufferized scf.foreach_thread lowers to "
               "gpu.block_id";
   if (foreachThreadOp.getNumThreads().size() > 3)
     return transformOp.emitSilenceableError()
            << "scf.foreach_thread with rank > 3 does not lower to "
               "gpu.block_id";
   if (llvm::any_of(foreachThreadOp.getNumThreads(), [](Value v) {
         return !v.getDefiningOp<arith::ConstantIndexOp>();
       })) {
     return transformOp.emitSilenceableError()
            << "unsupported dynamic griddim size";
   }
-  if (!foreachThreadOp.getMapping().has_value())
-    return transformOp.emitSilenceableError() << "mapping must be present";
   SmallVector<Attribute> blockMapping =
       llvm::to_vector(foreachThreadOp.getMapping()->getValue());
-  if (llvm::any_of(blockMapping, [](DeviceMappingAttrInterface map) {
-        return !map.isa<GPUBlockMappingAttr>();
-      })) {
-    return transformOp.emitSilenceableError()
-           << "mapping must be #gpu.block<x/y/z/>";
-  }
 
   // Step 1. Complete the blockMapping to a full mapping (with 1s) if necessary.
   SmallVector<Value> numBlocks =
       llvm::to_vector(foreachThreadOp.getNumThreads());
   // Ensure we have 3 block sizes, one for each id.
   Value one;
-  for (auto attr : {bX, bY, bZ}) {
+  for (auto attr : mappingAttributes) {
     if (std::find(blockMapping.begin(), blockMapping.end(), attr) ==
         blockMapping.end()) {
       blockMapping.push_back(attr);
       one = one ? one : rewriter.create<arith::ConstantIndexOp>(loc, 1);
       numBlocks.push_back(one);
     }
   }
 
-  // Step 2. sort the values by the corresponding GPUBlockMappingAttr.
-  auto comparator = [](Attribute a, Attribute b) -> bool {
-    return static_cast<int64_t>(a.cast<GPUBlockMappingAttr>().getBlock()) <
-           static_cast<int64_t>(b.cast<GPUBlockMappingAttr>().getBlock());
+  // Step 2. sort the values by the corresponding DeviceMappingAttrInterface.
+  auto comparator = [&](DeviceMappingAttrInterface a,
+                        DeviceMappingAttrInterface b) -> bool {
+    return a.getMappingId() < b.getMappingId();
   };
   SmallVector<Value> gridDimValues = scf::ForeachThreadOp::getValuesSortedByKey(
       blockMapping, numBlocks, comparator);
   for (Value v : gridDimValues)
     gridDims.push_back(v.getDefiningOp<arith::ConstantIndexOp>().value());
 
   // Step 3. Generate the blockIds using the provided generator and map the
   // induction variables to the newly created ops.
   SmallVector<Value> blockOps;
   blockIdGenerator(rewriter, foreachThreadOp, blockOps);
   BlockAndValueMapping bvm;
   for (auto [blockIdx, blockDim] :
        llvm::zip(foreachThreadOp.getThreadIndices(), blockMapping)) {
-    bvm.map(blockIdx, blockOps[static_cast<int64_t>(
-                          blockDim.cast<GPUBlockMappingAttr>().getBlock())]);
+    bvm.map(blockIdx,
+            blockOps[static_cast<int64_t>(
+                blockDim.cast<DeviceMappingAttrInterface>().getMappingId())]);
   }
 
   // Step 4. 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 = foreachThreadOp->getBlock();
   Block::iterator insertionPoint = Block::iterator(foreachThreadOp);
   Block &sourceBlock = foreachThreadOp.getRegion().front();
if (getGenerateGpuLaunch()) {
     }
     rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());
     Operation *newForeachThreadOp = rewriter.clone(*topLevelForeachThreadOp);
     rewriter.eraseOp(topLevelForeachThreadOp);
     topLevelForeachThreadOp = cast<scf::ForeachThreadOp>(newForeachThreadOp);
   }
 
   SmallVector<int64_t> gridDim = extractFromI64ArrayAttr(getGridDim());
+  SmallVector<DeviceMappingAttrInterface> blockMappingAttributes = {
+      GPUBlockMappingAttr::get(getContext(), Blocks::DimX),
+      GPUBlockMappingAttr::get(getContext(), Blocks::DimY),
+      GPUBlockMappingAttr::get(getContext(), Blocks::DimZ)};
+
+  diag = checkAttributeType(blockMappingAttributes,
+                            topLevelForeachThreadOp.getMapping(), transformOp);
+  if (diag.succeeded())
     diag = mlir::transform::gpu::mapForeachToBlocksImpl(
         rewriter, topLevelForeachThreadOp, generateGpuBlockIds, gridDim,
-      transformOp);
+        transformOp, blockMappingAttributes);
   if (diag.succeeded()) {
     diag = alterGpuLaunch(rewriter, gpuLaunch,
                           cast<TransformOpInterface>(getOperation()),
                           gridDim[0], gridDim[1], gridDim[2]);
   }
 
   results.assign({gpuLaunch});
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // MapNestedForeachToThreads
 //===----------------------------------------------------------------------===//
 
 /// 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.
 static DiagnosedSilenceableFailure rewriteOneForeachThreadToGpuThreads(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
     const SmallVectorImpl<int64_t> &globalBlockDims, bool syncAfterDistribute,
-    llvm::Optional<TransformOpInterface> transformOp) {
+    llvm::Optional<TransformOpInterface> transformOp,
+    const ArrayRef<DeviceMappingAttrInterface> &threadMappingAttributes) {
   // Step 0. Target-specific verifications. There is no good place to anchor
   // those right now: the ForeachThreadOp is target-independent and the
   // transform op does not apply to individual ForeachThreadOp.
   auto failureHelper =
       [&](const Twine &message) -> DiagnosedSilenceableFailure {
     if (transformOp.has_value()) {
       return transformOp->emitSilenceableError() << message;
     }
     return emitDefiniteFailure(foreachThreadOp, message);
   };
-  MLIRContext *ctx = foreachThreadOp->getContext();
   Location loc = foreachThreadOp->getLoc();
-  Attribute tX = GPUThreadMappingAttr::get(ctx, Threads::DimX);
-  Attribute tY = GPUThreadMappingAttr::get(ctx, Threads::DimY);
-  Attribute tZ = GPUThreadMappingAttr::get(ctx, Threads::DimZ);
   if (foreachThreadOp.getNumResults() > 0)
     return failureHelper(
         "only bufferized scf.foreach_thread lowers to gpu.thread_id");
   if (foreachThreadOp.getNumThreads().size() > 3)
     return failureHelper(
         "scf.foreach_thread with rank > 3 does not lower to gpu.thread_id");
   if (llvm::any_of(foreachThreadOp.getNumThreads(), [](Value v) {
         return !v.getDefiningOp<arith::ConstantIndexOp>();
       })) {
     return failureHelper("unsupported dynamic blockdim size");
   }
   if (!foreachThreadOp.getMapping().has_value())
     return failureHelper("mapping must be present");
   SmallVector<Attribute> threadMapping =
       llvm::to_vector(foreachThreadOp.getMapping()->getValue());
-  if (llvm::any_of(threadMapping, [](DeviceMappingAttrInterface map) {
-        return !map.isa<GPUThreadMappingAttr>();
-      })) {
-    return transformOp->emitSilenceableError()
-           << "mapping must be #gpu.thread<x/y/z/>";
-  }
 
   // Step 1. Complete the threadMapping to a full mapping (with 1s) if
   // necessary.
   SmallVector<Value> numThreads =
       llvm::to_vector(foreachThreadOp.getNumThreads());
   // Ensure we have 3 block sizes, one for each id.
   Value one;
-  for (auto attr : {tX, tY, tZ}) {
+  for (auto attr : threadMappingAttributes) {
     if (std::find(threadMapping.begin(), threadMapping.end(), attr) ==
         threadMapping.end()) {
       threadMapping.push_back(attr);
       one = one ? one : rewriter.create<arith::ConstantIndexOp>(loc, 1);
       numThreads.push_back(one);
     }
   }
 
-  // Step 2. sort the values by the corresponding GPUThreadMappingAttr.
-  auto comparator = [](Attribute a, Attribute b) -> bool {
-    return static_cast<int64_t>(a.cast<GPUThreadMappingAttr>().getThread()) <
-           static_cast<int64_t>(b.cast<GPUThreadMappingAttr>().getThread());
+  // Step 2. sort the values by the corresponding DeviceMappingAttrInterface.
+  auto comparator = [&](DeviceMappingAttrInterface a,
+                        DeviceMappingAttrInterface b) -> bool {
+    return a.getMappingId() < b.getMappingId();
   };
   SmallVector<Value> blockDimValues =
       scf::ForeachThreadOp::getValuesSortedByKey(threadMapping, numThreads,
                                                  comparator);
   SmallVector<int64_t> blockDims =
       llvm::to_vector(llvm::map_range(blockDimValues, [](Value v) {
         return v.getDefiningOp<arith::ConstantIndexOp>().value();
       }));
 
   // Step 3. Create the gpu.thread ops and map the induction variables to the
   // newly created ops.
   IndexType indexType = rewriter.getIndexType();
   SmallVector<Value> threadOps{
       rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x),
       rewriter.create<ThreadIdOp>(loc, indexType, Dimension::y),
       rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z)};
   BlockAndValueMapping bvm;
   for (auto [blockIdx, blockDim] :
        llvm::zip(foreachThreadOp.getThreadIndices(), threadMapping)) {
-    bvm.map(blockIdx, threadOps[static_cast<int64_t>(
-                          blockDim.cast<GPUThreadMappingAttr>().getThread())]);
+    bvm.map(
+        blockIdx,
+        threadOps[blockDim.cast<DeviceMappingAttrInterface>().getMappingId()]);
   }
 
   // Step 4. Maybe create conditionals to predicate the region.
   Value predicate;
   for (auto [threadId, blockDim, globalBlockDim] :
        llvm::zip(threadOps, blockDims, globalBlockDims)) {
     if (blockDim > globalBlockDim) {
       return failureHelper(
static DiagnosedSilenceableFailure rewriteOneForeachThreadToGpuThreads(
   rewriter.eraseOp(foreachThreadOp);
 
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForeachToThreadsImpl(
     RewriterBase &rewriter, Operation *target,
     const SmallVectorImpl<int64_t> &blockDim, bool syncAfterDistribute,
-    llvm::Optional<TransformOpInterface> transformOp) {
+    llvm::Optional<TransformOpInterface> transformOp,
+    const ArrayRef<DeviceMappingAttrInterface> &threadMappingAttributes) {
   DiagnosedSilenceableFailure diag = DiagnosedSilenceableFailure::success();
   target->walk([&](scf::ForeachThreadOp foreachThreadOp) {
+    diag = checkAttributeType(threadMappingAttributes,
+                              foreachThreadOp.getMapping(), transformOp);
+    if (diag.succeeded()) {
       rewriter.setInsertionPoint(foreachThreadOp);
       diag = rewriteOneForeachThreadToGpuThreads(
-        rewriter, foreachThreadOp, blockDim, syncAfterDistribute, transformOp);
+          rewriter, foreachThreadOp, blockDim, syncAfterDistribute, transformOp,
+          threadMappingAttributes);
+    }
     return diag.succeeded() ? WalkResult::advance() : WalkResult::interrupt();
   });
   return diag;
 }
 
 DiagnosedSilenceableFailure transform::MapNestedForeachToThreads::applyToOne(
     ::mlir::Operation *target,
     ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
DiagnosedSilenceableFailure diag =
       checkGpuLimits(transformOp, llvm::None, llvm::None, llvm::None,
                      blockDim[0], blockDim[1], blockDim[2]);
   if (diag.isSilenceableFailure()) {
     results.assign({target});
     diag.attachNote(getLoc()) << getBlockDimAttrName() << " is very large";
     return diag;
   }
 
-  SimpleRewriter rewriter(getContext());
+  MLIRContext *ctx = getContext();
+  SimpleRewriter rewriter(ctx);
   rewriter.setInsertionPoint(target);
 
+  SmallVector<DeviceMappingAttrInterface> threadMappingAttributes = {
+      GPUThreadMappingAttr::get(ctx, Threads::DimX),
+      GPUThreadMappingAttr::get(ctx, Threads::DimY),
+      GPUThreadMappingAttr::get(ctx, Threads::DimZ)};
+
   diag = mlir::transform::gpu::mapNestedForeachToThreadsImpl(
-      rewriter, target, blockDim, getSyncAfterDistribute(), transformOp);
+      rewriter, target, blockDim, getSyncAfterDistribute(), transformOp,
+      threadMappingAttributes);
+
   if (diag.succeeded()) {
     diag =
         alterGpuLaunch(rewriter, gpuLaunch, transformOp, llvm::None, llvm::None,
                        llvm::None, blockDim[0], blockDim[1], blockDim[2]);
   }
 
   results.assign({gpuLaunch});
   return diag;

mlir/test/Dialect/GPU/transform-gpu-failing.mlir

%name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)
     gpu.terminator
   }  
   return 
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0
-  %foreach, %tiled = transform.structured.tile_to_foreach_thread_op %matmul num_threads [10, 20, 30]
+  %foreach, %tiled = transform.structured.tile_to_foreach_thread_op %matmul num_threads [10, 20, 30] (mapping = [ #gpu.thread<y>, #gpu.thread<x>, #gpu.thread<z> ] )
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
   // expected-error @below {{only bufferized scf.foreach_thread lowers to gpu.thread_id}}    
   transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [128, 4, 1] }
 }
 
 // -----