[mlir] Use processing unit names for `thread_dim_map` and `mapped to dims`

[guraypp]

Summary

Currently using thread_dim_map and a directory array map to dims. This is very confusing in many ways. This change uses meaningful words in these structures.

For now there is thread_x/y/z and block_x/y/z. They cannot be mixed in the same foreach_thread. However, it is possible to mix them or use them together in some cases in the future.

The change is almost NFC other than changing the names.

[guraypp]

update the description

[guraypp]

remove comments

[ftynse]

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

61

Nit: do not specify the number of stack elements in the vector unless you have a strong reason to.

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

298–299

This seems orthogonal to the purpose of the change, can you please comment why this is necessary?

mlir/lib/Dialect/SCF/IR/SCF.cpp

1073–1081

Why does it have to be a bit map? Are we expecting to use a mapping that is simultaneously thread-x, thread-y and block-z, for example?

[nicolasvasilache]

In D136851#3890998, @ftynse wrote:

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

One aspect here is that the op semantics and transformations should be retargetable.
Is there a way to have something like a BaseForeachThreadTargetEnumAttr for which the GPU dialect could define Thread/Block; the IREE xxx dialect could define Workgroupyy and future ACME dialect could define its own?

A (far from) ideal solution would be to have users define their own names "foo", "bar", "baz" and explicitly have the map_to_target specify the map ["foo" -> threadIdx.x ...].
This is so ugly that I don't think it should be considered but that should give a feeling on what a better preferred way to interact with the system could look like.

[guraypp]

remove size of the vector

[guraypp]

In D136851#3890998, @ftynse wrote:

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

I totally agree with you. Actually I created the enum array at the beginning. But`foreach_thread` is manually parsed, printed, verified. So I had to parse enum array manually here, it looked also unsafe to me. Therefore, I choose string array option.
But If you have a strong argument, I can work on changing to enum.

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

61

I removed the size here.

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

298–299

Good catch :) This is a minor bug fix.

Without this line, blockDim always required 3 dimensions. It was annoying when foreach_thread contains less than 3 loops. This line will set 1 to unused dimensions.

mlir/lib/Dialect/SCF/IR/SCF.cpp

1073–1081

Exactly. One can think about mixings blocks and threads like below. This will be my future work.

loop1 --> blockIdx.x, threadIdx.x
loop2 --> blockIdx.y, threadIdx.y

[ftynse]

In D136851#3891068, @nicolasvasilache wrote:

In D136851#3890998, @ftynse wrote:

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

One aspect here is that the op semantics and transformations should be retargetable.
Is there a way to have something like a BaseForeachThreadTargetEnumAttr for which the GPU dialect could define Thread/Block; the IREE xxx dialect could define Workgroupyy and future ACME dialect could define its own?

A (far from) ideal solution would be to have users define their own names "foo", "bar", "baz" and explicitly have the map_to_target specify the map ["foo" -> threadIdx.x ...].
This is so ugly that I don't think it should be considered but that should give a feeling on what a better preferred way to interact with the system could look like.

If you don't constrain the kind of attributes in the op verifier, any downstream client can use whatever they want there. Lowerings can then check if they support the specific attribute kind and fail early if they don't. If you want more guarantees, introduce an attribute interface. There is no need to invent new infrastructure layers here.

[nicolasvasilache]

In D136851#3891267, @ftynse wrote:

In D136851#3891068, @nicolasvasilache wrote:

In D136851#3890998, @ftynse wrote:

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

One aspect here is that the op semantics and transformations should be retargetable.
Is there a way to have something like a BaseForeachThreadTargetEnumAttr for which the GPU dialect could define Thread/Block; the IREE xxx dialect could define Workgroupyy and future ACME dialect could define its own?

A (far from) ideal solution would be to have users define their own names "foo", "bar", "baz" and explicitly have the map_to_target specify the map ["foo" -> threadIdx.x ...].
This is so ugly that I don't think it should be considered but that should give a feeling on what a better preferred way to interact with the system could look like.

If you don't constrain the kind of attributes in the op verifier, any downstream client can use whatever they want there. Lowerings can then check if they support the specific attribute kind and fail early if they don't. If you want more guarantees, introduce an attribute interface. There is no need to invent new infrastructure layers here.

Cool thnks for the guidnce, I had not followed that part of the stack.

[guraypp]

In D136851#3891357, @nicolasvasilache wrote:

In D136851#3891267, @ftynse wrote:

In D136851#3891068, @nicolasvasilache wrote:

In D136851#3890998, @ftynse wrote:

Have you considered making it an actual attribute in the GPU dialect, e.g., an enum attribute? Passing strings around feels a lot like unsafe JSON even if it is a readability improvement.

One aspect here is that the op semantics and transformations should be retargetable.
Is there a way to have something like a BaseForeachThreadTargetEnumAttr for which the GPU dialect could define Thread/Block; the IREE xxx dialect could define Workgroupyy and future ACME dialect could define its own?

A (far from) ideal solution would be to have users define their own names "foo", "bar", "baz" and explicitly have the map_to_target specify the map ["foo" -> threadIdx.x ...].
This is so ugly that I don't think it should be considered but that should give a feeling on what a better preferred way to interact with the system could look like.

If you don't constrain the kind of attributes in the op verifier, any downstream client can use whatever they want there. Lowerings can then check if they support the specific attribute kind and fail early if they don't. If you want more guarantees, introduce an attribute interface. There is no need to invent new infrastructure layers here.

Cool thnks for the guidnce, I had not followed that part of the stack.

Then I have a proposal. As Alex said, let's relex verifier of foreach_thread, for example as follows (here I improved the syntax a bit).

scf.foreach_thread (%bi, %bj)  {
  scf.foreach_thread (%ti, %tj) {
  }  {map = "unit", dimensions = ["x", "y"]}
}  {map = "group", dimensions = ["x", "y"]}

Next, let's add a new generic op to the GPU dialect. One can map named parallelism unit to block and thread.

// transform dialect for GPU
transform.gpu.map_foreach %op "group" to blocks
transform.gpu.map_foreach %op "unit" to threads

Another accelerator transform dialect uses only the outer level parallelism for example.

// transform dialect for other accelerator
transform.other_accelerator.map_foreach %op "group" to threads

[guraypp]

Then I have a proposal. As Alex said, let's relex verifier of foreach_thread, for example as follows (here I improved the syntax a bit).

scf.foreach_thread (%bi, %bj)  {
  scf.foreach_thread (%ti, %tj) {
  }  {map = "unit", dimensions = ["x", "y"]}
}  {map = "group", dimensions = ["x", "y"]}

Next, let's add a new generic op to the GPU dialect. One can map named parallelism unit to block and thread.

// transform dialect for GPU
transform.gpu.map_foreach %op "group" to blocks
transform.gpu.map_foreach %op "unit" to threads

Another accelerator transform dialect uses only the outer level parallelism for example.

// transform dialect for other accelerator
transform.other_accelerator.map_foreach %op "group" to threads

I will need to also change tile_to_foreach_thread_op to something like below.

transform.structured.tile_to_foreach_thread_op %op (map = "unit", dimensions = ["x", "y"])

[ftynse]

In D136851#3891440, @guraypp wrote:

Then I have a proposal. As Alex said, let's relex verifier of foreach_thread, for example as follows (here I improved the syntax a bit).

scf.foreach_thread (%bi, %bj)  {
  scf.foreach_thread (%ti, %tj) {
  }  {map = "unit", dimensions = ["x", "y"]}
}  {map = "group", dimensions = ["x", "y"]}

Next, let's add a new generic op to the GPU dialect. One can map named parallelism unit to block and thread.

// transform dialect for GPU
transform.gpu.map_foreach %op "group" to blocks
transform.gpu.map_foreach %op "unit" to threads

Another accelerator transform dialect uses only the outer level parallelism for example.

// transform dialect for other accelerator
transform.other_accelerator.map_foreach %op "group" to threads

I think having string attributes and then binding them to something in extra operations is an overkill and also doesn't advance us in the right direction. If you watch some of the earlier MLIR presentations, we specifically wanted to avoid MLIR becoming a "JSON of compiler IRs". String attributes loosely bound by an operation located elsewhere has exactly that JSON feeling for me. What I propose concretely is to introduce a DeviceMappingAttrInterface, initially with no methods at all, just as a unification mechanism. Then, have a GPUThreadMappingAttr and GPUBlockMappingAttr, potentially sharing parts of the implementation, that implement the DeviceMappingAttrInterface. These can print as #gpu.threads<y, z> and #gpu.blocks<x, y>, respectively. Other accelerators can then introduce their own attributes #other_accelerator.mapping<parallel_hw_feature_no_7> that also implement the interface. The verifier of foreach_thread can then check that the mapping attribute is DeviceMappingAttrInterface (or an array thereof depending on the desired model for mapping to more than one parallel HW dimension). Initially, the lowering pattern to GPU can do an additional isa<GPUThreadMappingAttr, GPUBlockMappingAttr>() on the attribute and fail to match otherwise, with a helpful debug message. Other accelerators can do the same. In a slightly longer term, we can consider moving HW-specific parts of the lowering into interface methods so we can keep a common lowering, but this is not a priority.

[guraypp]

@ftynse thanks for clear explanation. As it is different that this current PR, I implement in another in D137413. Let me know what do you think.

[nicolasvasilache]

I think we can now abandon this ?

[guraypp]

We implemented a better solution in https://reviews.llvm.org/D137413

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

Op<Transform_Dialect, "gpu.map_nested_foreach_to_threads",
      TransformOpInterface]> {
   let description = [{
       Target the `gpu.launch op` and rewrite all `scf.foreach_thread`
       nested in it to distributed `gpu.thread_id` attribute.
 
       The operation searches for `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_mapping` attribute. 
+      `gpu.thread_id` according to the `mapping` attribute. 
       
       Sibling `scf.foreach_thread` are supported in which case, the union of 
       the number of threads is computed and may result in predication.
 
       Multiple scf.foreach_thread are supported per `gpu.launch` in which case, 
       the max of all the threads is computed and taken for the global 
       `gpu.thread_id`. If necessary, `scf.foreach_thread` that do not use the 
       whole thread range result in predicated computations.
let description = [{
 
       #### Example:
 
       ```
       gpu.launch blocks(%bx, %by, %bz) in (%x = %0, %y = %1, %z = %2)
                  threads(%tx, %ty, %tz) in (%tx = %3, %ty = %4, %tz = %5) {
         scf.foreach_thread (%i, %j) in (7, 9) {
           ... // body 1
-        } {thread_dim_mapping = [1, 0, 2]}
+        } {mapping = ["thread_y", "thread_x"]}
         scf.foreach_thread (%i) in (12) {
           ... // body 2
         }
         gpu.terminator
       }
       ```
       is translated to:
 

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

let description = [{
     ```
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    Variadic<PDL_Operation>:$num_threads,
                    Variadic<PDL_Operation>:$tile_sizes,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$static_num_threads,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$static_tile_sizes,
-                   OptionalAttr<I64ArrayAttr>:$thread_dim_mapping);
+                   OptionalAttr<StrArrayAttr>:$mapping);
   let results = (outs PDL_Operation:$foreach_thread_op,
                       PDL_Operation:$tiled_op);
   let assemblyFormat = [{
     $target oilist(
         `num_threads` custom<DynamicIndexList>($num_threads,
                                                $static_num_threads,
                                                "ShapedType::kDynamicSize") |
          `tile_sizes` custom<DynamicIndexList>($tile_sizes,
                                                $static_tile_sizes,
                                                "ShapedType::kDynamicSize"))
-    (`(` `mapped` `to` `dims` $thread_dim_mapping^ `)`)? attr-dict
+    (`(` `mapped` `to` `dims` $mapping^ `)`)? attr-dict
   }];
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure apply(
         ::mlir::transform::TransformResults &transformResults,
         ::mlir::transform::TransformState &state);
 

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

 /// dimensions, e.g. in the Linalg case).
 struct ForeachThreadTilingResult {
   Operation *tileOp;
   Operation *tiledOp;
 };
 FailureOr<ForeachThreadTilingResult>
 tileToForeachThreadOp(RewriterBase &builder, TilingInterface op,
                       ArrayRef<OpFoldResult> numThreads,
-                      ArrayRef<int64_t> threadDimMapping = {});
+                      ArrayRef<StringRef> threadDimMapping = {});
 
 /// Same as `tileToForeachThreadOp`, but calculate the number of threads
 /// required using the given tileSizes.
 FailureOr<ForeachThreadTilingResult>
 tileToForeachThreadOpUsingTileSizes(RewriterBase &builder, TilingInterface op,
                                     ArrayRef<OpFoldResult> tileSizes,
-                                    ArrayRef<int64_t> threadDimMapping = {});
+                                    ArrayRef<StringRef> threadDimMapping = {});
 
 /// All indices returned by IndexOp should be invariant with respect to tiling.
 /// Therefore, if an operation is tiled, we have to transform the indices
 /// accordingly, i.e. offset them by the values of the corresponding induction
 /// variables that are captured implicitly in the body of the op.
 ///
 /// Example. `linalg.generic` before tiling:
 ///

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

let description = [{
     used. This ensures that memory side effects of a thread are not visible to
     other threads (or in the parent body), apart from explicitly shared tensors.
 
     The name "thread" conveys the fact that the parallel execution is mapped
     (i.e. distributed) to a set of virtual threads of execution, one function
     application per thread. Further lowerings are responsible for specifying
     how this is materialized on concrete hardware resources.
 
-    An optional thread_dim_mapping index array attribute specifies for each
-    virtual thread dimension, how it remaps 1-1 to a set of concrete processing
-    element resources (e.g. a CUDA grid dimension or a level of concrete nested
-    async parallelism). At this time, the specification is backend-dependent and
-    is not verified by the op, beyond being an index array attribute.
-    It is the reponsibility of the lowering to interpret the index array in the
-    context of the concrete target the op is lowered to, or to ignore it when
-    the specification is ill-formed or unsupported for a particular target.
+    An optional `mapping` string array attribute specifies for each virtual 
+    thread dimension, how it remaps 1-1 to a set of concrete processing element
+    resources (e.g. a CUDA grid dimension or a level of concrete nested async 
+    parallelism). Mapping can be one of following thread_x, thread_y, thread_z, 
+    block_x, block_y, and block_z. in this way, one can use 3 dimensional threads 
+    and blocks that contains threads. The threads and blocks cannot be mixed at 
+    this point or used together. 
 
     The only allowed terminator is `scf.foreach_thread.perform_concurrently`.
     `scf.foreach_thread` returns one value per `shared_out` operand. The
     actions of the `perform_concurrently` terminators specify how to combine the
     partial results of all parallel invocations into a full value, in some
     unspecified order. The "destination" of each such op must be a `shared_out`
     block argument of the `scf.foreach_thread` op.
 
%matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in
           tensor<?xT> into tensor<?xT>
       }
     }
     // Implicit synchronization point.
     // Sequential context.
     //
     ```
 
-    Example with thread_dim_mapping attribute:
+    Example with mapping attribute:
 
     ```mlir
     //
     // Sequential context.
     //
     %matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in
         (%num_threads_1, %numthread_id_2) shared_outs(...)
       -> (tensor<?x?xT>, tensor<?xT>) {
       //
       // Parallel context, each thread with id = **(%thread_id_2, %thread_id_1)**
       // runs its version of the code.
       //
        scf.foreach_thread.perform_concurrently {
          ...
       }
-    } { thread_dim_mapping = [1, 0] }
+    } { mapping = ["thread_y", "thread_x"] }
     // Implicit synchronization point.
     // Sequential context.
     //
     ```
 
     Example with privatized tensors:
 
     ```mlir
     %t0 = ...
     %t1 = ...
     %r = scf.foreach_thread ... shared_outs(%o = t0) -> tensor<?xf32> {
       // %t0 and %t1 are privatized. %t0 is definitely copied for each thread
       // because the scf.foreach_thread op's %t0 use bufferizes to a memory
       // write. In the absence of other conflicts, %t1 is copied only if there
       // are uses of %t1 in the body that bufferize to a memory read and to a
       // memory write.
       "some_use"(%t0)
       "some_use"(%t1)
     }
     ```
   }];
   let arguments = (ins Variadic<Index>:$num_threads,
                        Variadic<AnyRankedTensor>:$outputs,
-                       DefaultValuedAttr<I64ArrayAttr, "{}">:$thread_dim_mapping);
+                       DefaultValuedAttr<StrArrayAttr, "{}">:$mapping);
 
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 
   // The default builder does not add the proper body BBargs, roll our own.
   let skipDefaultBuilders = 1;
   let builders = [
     // Bodyless builder, outputs must be specified.
     OpBuilder<(ins "ValueRange":$outputs, "ValueRange":$num_threads,
-                   CArg<"ArrayRef<int64_t>", "{}">:$thread_dim_mapping)>,
+                   CArg<"ArrayRef<StringRef>", "{}">:$mapping)>,
     // Builder that takes a bodyBuilder lambda.
     OpBuilder<(ins "ValueRange":$outputs, "ValueRange":$num_threads,
-                   "ArrayRef<int64_t>":$thread_dim_mapping,
+                   "ArrayRef<StringRef>":$mapping,
                    "function_ref<void(OpBuilder &, Location, ValueRange)>":$bodyBuilder)>
   ];
   let extraClassDeclaration = [{
     int64_t getRank() { return getNumThreads().size(); }
 
     OpResult getTiedOpResult(OpOperand *opOperand) {
       assert(opOperand->getOperandNumber() >= getRank() && "invalid operand");
       return getOperation()->getOpResult(
::mlir::Value getThreadIndex(int64_t idx) {
       return getThreadIndices()[idx];
     }
 
     ::mlir::Block::BlockArgListType getRegionOutArgs() {
       return getBody()->getArguments().drop_front(getRank());
     }
 
     /// Return the thread indices in the order specified by the
-    /// thread_dim_mapping attribute. Return failure is
-    /// thread_dim_mapping is not a valid permutation.
+    /// mapping attribute. Return failure is mapping is not a valid permutation.
     FailureOr<SmallVector<Value>> getPermutedThreadIndices();
 
     /// Return the number of threads in the order specified by the
-    /// thread_dim_mapping attribute.
-    /// Return failure is thread_dim_mapping is not a valid permutation.
+    /// mapping attribute.
+    /// Return failure is mapping is not a valid permutation.
     FailureOr<SmallVector<OpFoldResult>> getPermutedNumThreads(OpBuilder &b);
 
     // The ensureTerminator method generated by SingleBlockImplicitTerminator is
     // unaware of the fact that our terminator also needs a region to be
     // well-formed. We override it here to ensure that we do the right thing.
     static void ensureTerminator(Region &region, OpBuilder &builder, Location loc);
 
     PerformConcurrentlyOp getTerminator();

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

 void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
                                 SmallVectorImpl<Value> &dynamicVec,
                                 SmallVectorImpl<int64_t> &staticVec,
                                 int64_t sentinel);
 
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr);
 
+/// Extract StringRef values from the assumed ArrayAttr of StringAttr.
+SmallVector<StringRef> extractFromStringArrayAttr(Attribute attr);
+
 /// Given a value, try to extract a constant Attribute. If this fails, return
 /// the original value.
 OpFoldResult getAsOpFoldResult(Value val);
 
 /// Given an array of values, try to extract a constant Attribute from each
 /// value. If this fails, return the original value.
 SmallVector<OpFoldResult> getAsOpFoldResult(ValueRange values);
 

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

if (getGenerateGpuLaunch()) {
     topLevelForeachThreadOp = cast<scf::ForeachThreadOp>(newForeachThreadOp);
   }
 
   SmallVector<int64_t> gridDim = extractFromI64ArrayAttr(getGridDim());
   diag = mlir::transform::gpu::mapForeachToBlocksImpl(
       rewriter, topLevelForeachThreadOp, generateGpuBlockIds, gridDim,
       transformOp);
   if (diag.succeeded()) {
+    // Set 1 to dim that aren't used
+    gridDim.resize(3, 1);
     diag = alterGpuLaunch(rewriter, gpuLaunch,
                           cast<TransformOpInterface>(getOperation()),
                           gridDim[0], gridDim[1], gridDim[2]);
   }
 
   results.assign({gpuLaunch});
   return diag;
 }
if (failed(potentialBlockDim) ||
       })) {
     return failureHelper("unsupported dynamic blockdim size");
   }
 
   SmallVector<int64_t> blockDim =
       llvm::to_vector(llvm::map_range(*potentialBlockDim, [](OpFoldResult ofr) {
         return getConstantIntValue(ofr).value();
       }));
+  blockDim.resize(3, 1);
 
   // Step 1. Create the gpu.thread ops
   Location loc = foreachThreadOp.getLoc();
   IndexType indexType = rewriter.getIndexType();
 
   SmallVector<Dimension> gpuDims{Dimension::x, Dimension::y, Dimension::z};
   SmallVector<Value> threadOps;
   for (int64_t idx : llvm::seq<int64_t>(0, blockDim.size())) {

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

if (!tilableOp) {
           << "only TilingInterface ops are supported";
       diag.attachNote(target->getLoc()) << "target op";
       return diag;
     }
     rewriter.setInsertionPoint(tilableOp);
     auto maybeThreadDimMappingAttr = threadDimMapping;
     auto dimMapping = llvm::to_vector(
         maybeThreadDimMappingAttr
-            ? extractFromI64ArrayAttr(*maybeThreadDimMappingAttr)
-            : ArrayRef<int64_t>{});
+            ? extractFromStringArrayAttr(*maybeThreadDimMappingAttr)
+            : ArrayRef<StringRef>{});
 
     FailureOr<linalg::ForeachThreadTilingResult> tilingResult = failure();
     if (!mixedNumThreads.empty()) {
       tilingResult = linalg::tileToForeachThreadOp(rewriter, tilableOp,
                                                    numThreads, dimMapping);
     } else {
       tilingResult = linalg::tileToForeachThreadOpUsingTileSizes(
           rewriter, tilableOp, tileSizes, dimMapping);
DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
   ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
 
   // Result payload ops.
   SmallVector<Operation *> tileOps;
   SmallVector<Operation *> tiledOps;
 
   DiagnosedSilenceableFailure diag = tileToForeachThreadOpImpl(
       rewriter, state, cast<TransformOpInterface>(getOperation()), targets,
-      getMixedNumThreads(), getMixedTileSizes(), getThreadDimMapping(), tileOps,
+      getMixedNumThreads(), getMixedTileSizes(), getMapping(), tileOps,
       tiledOps);
 
   if (!diag.succeeded())
     return diag;
 
   transformResults.set(getForeachThreadOp().cast<OpResult>(), tileOps);
   transformResults.set(getTiledOp().cast<OpResult>(), tiledOps);
 

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

 /// size of data.
 /// It is the user's responsibility to ensure that `numThreads` is a valid
 /// tiling specification (i.e. that only tiles parallel dimensions, e.g. in the
 /// Linalg case). If `omitTileOffsetBoundsCheck` is true, then the function will
 /// assume that `tileSize[i] * (numThread[i] -1) <= dimSize[i]` holds.
 static FailureOr<ForeachThreadTilingResult> tileToForeachThreadOpImpl(
     RewriterBase &b, TilingInterface op, ArrayRef<OpFoldResult> numThreads,
     Optional<ArrayRef<OpFoldResult>> nominalTileSizes,
-    ArrayRef<int64_t> threadDimMapping, bool omitTileOffsetBoundsCheck) {
+    ArrayRef<StringRef> threadDimMapping, bool omitTileOffsetBoundsCheck) {
   Location loc = op->getLoc();
   OpBuilder::InsertionGuard g(b);
   SmallVector<Range> loopRanges = op.getIterationDomain(b);
   if (loopRanges.empty())
     return op->emitOpError("expected non-empty loop ranges");
   auto hasStrideOne = [](Range r) { return !isConstantIntValue(r.stride, 1); };
   if (llvm::any_of(loopRanges, hasStrideOne))
     return op->emitOpError("only stride-1 supported atm");
b.create<tensor::ParallelInsertSliceOp>(loc, std::get<1>(it),
                                             resultSizes, strides);
   }
   return ForeachThreadTilingResult{foreachThreadOp, tiledOp};
 }
 
 FailureOr<ForeachThreadTilingResult>
 linalg::tileToForeachThreadOp(RewriterBase &b, TilingInterface op,
                               ArrayRef<OpFoldResult> numThreads,
-                              ArrayRef<int64_t> threadDimMapping) {
+                              ArrayRef<StringRef> threadDimMapping) {
   return tileToForeachThreadOpImpl(b, op, numThreads, /*nominalTileSizes=*/None,
                                    threadDimMapping,
                                    /*omitTileOffsetBoundsCheck=*/false);
 }
 
 FailureOr<ForeachThreadTilingResult>
 linalg::tileToForeachThreadOpUsingTileSizes(
     RewriterBase &b, TilingInterface op, ArrayRef<OpFoldResult> tileSizes,
-    ArrayRef<int64_t> threadDimMapping) {
+    ArrayRef<StringRef> threadDimMapping) {
   SmallVector<Range> loopRanges = op.getIterationDomain(b);
   unsigned nLoops = loopRanges.size();
   SmallVector<OpFoldResult> numThreads;
   numThreads.reserve(nLoops);
   AffineExpr s0, s1;
   bindSymbols(b.getContext(), s0, s1);
   AffineExpr divExpr = s0.ceilDiv(s1);
   for (const auto &it : llvm::zip(tileSizes, loopRanges)) {

mlir/lib/Dialect/SCF/IR/SCF.cpp

void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<ForOpIterArgsFolder, SimplifyTrivialLoops,
               LastTensorLoadCanonicalization, ForOpTensorCastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ForeachThreadOp
 //===----------------------------------------------------------------------===//
 
+// Thread mapping
+enum class ThreadMap {
+  none = 0,
+  tidx = 1,
+  tidy = 2,
+  tidz = 4,
+  bidx = 8,
+  bidy = 10,
+  bidz = 20
+};
+
+static Optional<ThreadMap> symbolizeThreadMap(StringRef str) {
+  return StringSwitch<Optional<ThreadMap>>(str)
+      .Case("none", ThreadMap::none)
+      .Case("thread_x", ThreadMap::tidx)
+      .Case("thread_y", ThreadMap::tidy)
+      .Case("thread_z", ThreadMap::tidz)
+      .Case("block_x", ThreadMap::bidx)
+      .Case("block_y", ThreadMap::bidy)
+      .Case("block_z", ThreadMap::bidz)
+      .Default(None);
+}
+
 LogicalResult ForeachThreadOp::verify() {
   // Call terminator's verify to produce most informative error messages.
   if (failed(getTerminator().verify()))
     return failure();
 
   // Check number of outputs.
   if (getNumResults() != getOutputs().size())
     return emitOpError("produces ")
            << getNumResults() << " results, but has only "
            << getOutputs().size() << " outputs";
 
   // Check that the body defines block arguments for thread indices and outputs.
   auto *body = getBody();
   if (body->getNumArguments() != getRank() + getOutputs().size())
     return emitOpError("region expects ") << getRank() << " arguments";
   for (int64_t i = 0; i < getRank(); ++i)
     if (!body->getArgument(i).getType().isIndex())
       return emitOpError("expects ")
              << i << "-th block argument to be an index";
   for (unsigned i = 0; i < getOutputs().size(); ++i)
     if (body->getArgument(i + getRank()).getType() != getOutputs()[i].getType())
       return emitOpError("type mismatch between ")
              << i << "-th output and corresponding block argument";
 
+  bool mapBlock = false, mapThread = false;
+  for (auto map : getMapping()) {
+    auto threadMap = symbolizeThreadMap(map.cast<StringAttr>().getValue());
+    if (!threadMap.has_value())
+      return emitOpError("expected ")
+             << "mapping to be one of: none, thread_x, thread_y, thread_z, "
+                "block_x, block_y, and block_z";
+    if (threadMap.value() == ThreadMap::tidx ||
+        threadMap.value() == ThreadMap::tidy ||
+        threadMap.value() == ThreadMap::tidz)
+      mapThread = true;
+    if (threadMap.value() == ThreadMap::bidx ||
+        threadMap.value() == ThreadMap::bidy ||
+        threadMap.value() == ThreadMap::bidz)
+      mapBlock = true;
+    if (mapBlock && mapThread)
+      return emitOpError("Mapping blocks and threads cannot be mixed in the "
+                         "same foreach_thread");
+  }
+
   return success();
 }
 
 void ForeachThreadOp::print(OpAsmPrinter &p) {
   p << " (";
   llvm::interleaveComma(getThreadIndices(), p);
   p << ") in (";
   llvm::interleaveComma(getNumThreads(), p);
result.addAttribute("operand_segment_sizes",
                            static_cast<int32_t>(outOperands.size())}));
   return success();
 }
 
 // Bodyless builder, outputs must be specified.
 void ForeachThreadOp::build(mlir::OpBuilder &builder,
                             mlir::OperationState &result, ValueRange outputs,
                             ValueRange numThreads,
-                            ArrayRef<int64_t> threadDimMapping) {
+                            ArrayRef<StringRef> mapping) {
   result.addOperands(numThreads);
   result.addOperands(outputs);
-  result.addAttribute(ForeachThreadOp::getThreadDimMappingAttrName(result.name),
-                      builder.getI64ArrayAttr(threadDimMapping));
+  result.addAttribute(ForeachThreadOp::getMappingAttrName(result.name),
+                      builder.getStrArrayAttr(mapping));
   result.addAttribute(
       "operand_segment_sizes",
       builder.getDenseI32ArrayAttr({static_cast<int32_t>(numThreads.size()),
                                     static_cast<int32_t>(outputs.size())}));
   result.addTypes(TypeRange(outputs));
 
   Region *bodyRegion = result.addRegion();
   OpBuilder::InsertionGuard g(builder);
bodyBlock.addArguments(
       TypeRange(outputs),
       SmallVector<Location>(outputs.size(), result.location));
   ForeachThreadOp::ensureTerminator(*bodyRegion, builder, result.location);
 }
 
 // Builder that takes a bodyBuilder lambda.
 void ForeachThreadOp::build(
     mlir::OpBuilder &builder, mlir::OperationState &result, ValueRange outputs,
-    ValueRange numThreads, ArrayRef<int64_t> threadDimMapping,
+    ValueRange numThreads, ArrayRef<StringRef> mapping,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
   result.addOperands(numThreads);
   result.addOperands(outputs);
-  result.addAttribute(ForeachThreadOp::getThreadDimMappingAttrName(result.name),
-                      builder.getI64ArrayAttr(threadDimMapping));
+  result.addAttribute(ForeachThreadOp::getMappingAttrName(result.name),
+                      builder.getStrArrayAttr(mapping));
   result.addAttribute(
       "operand_segment_sizes",
       builder.getDenseI32ArrayAttr({static_cast<int32_t>(numThreads.size()),
                                     static_cast<int32_t>(outputs.size())}));
   result.addTypes(TypeRange(outputs));
 
   Region *bodyRegion = result.addRegion();
   OpBuilder::InsertionGuard g(builder);

 }
 
 PerformConcurrentlyOp ForeachThreadOp::getTerminator() {
   return cast<PerformConcurrentlyOp>(getBody()->getTerminator());
 }
 
 template <typename T>
 static FailureOr<SmallVector<T>> permute(const SmallVector<T> &vals,
-                                         ArrayRef<int64_t> perm) {
+                                         ArrayRef<StringRef> perm) {
   if (vals.size() != perm.size())
     return failure();
   SmallVector<T> result(vals.size());
   SmallVector<bool> seen(vals.size());
-  for (auto [idx, val] : llvm::zip(perm, vals)) {
-    // Already seen, invalid thread_dim_mapping.
-    if (seen[idx])
+  for (auto [mappingName, val] : llvm::zip(perm, vals)) {
+    auto maybeMap = symbolizeThreadMap(mappingName);
+    if (!maybeMap.has_value())
+      return failure();
+    ThreadMap mapping = maybeMap.value();
+    int idx = 0;
+    if (mapping == ThreadMap::tidx || mapping == ThreadMap::bidx)
+      idx = 0;
+    else if (mapping == ThreadMap::tidy || mapping == ThreadMap::bidy)
+      idx = 1;
+    else if (mapping == ThreadMap::tidz || mapping == ThreadMap::bidz)
+      idx = 2;
+    // Already seen, invalid mapping.
+    if (seen[idx]) {
       return failure();
+    }
+    if (mapping != ThreadMap::none)
       result[idx] = val;
     seen[idx] = true;
   }
-  // Some not seen, invalid thread_dim_mapping.
+  // Some not seen, invalid mapping.
   if (!llvm::all_of(seen, [](bool b) { return b; }))
     return failure();
   return result;
 }
 
-/// Helper to get apply the `thread_dim_mapping` permutation of a
+/// Helper to get apply the `mapping` permutation of a
 /// `foreachThreadOp` to `values`.
 template <typename T>
 static FailureOr<SmallVector<T>>
 getValuesPermutedByThreadMapping(scf::ForeachThreadOp foreachThreadOp,
                                  const SmallVector<T> &values) {
   // Apply mapping permutation if specified.
-  auto mapping = foreachThreadOp.getThreadDimMapping();
+  auto mapping = foreachThreadOp.getMapping();
+  auto maps = extractFromStringArrayAttr(mapping);
   if (mapping && !mapping.empty()) {
-    auto maybePermuted = permute(values, extractFromI64ArrayAttr(mapping));
+    auto maybePermuted = permute(values, maps);
     if (failed(maybePermuted))
       return foreachThreadOp->emitError("invalid permutation");
     return *maybePermuted;
   }
   return values;
 }
 
-/// Return the thread indices in the order specified by the thread_dim_mapping
-/// attribute. Return failure is thread_dim_mapping is not a valid permutation.
+/// Return the thread indices in the order specified by the mapping
+/// attribute. Return failure is mapping is not a valid permutation.
 FailureOr<SmallVector<Value>> ForeachThreadOp::getPermutedThreadIndices() {
   SmallVector<Value> threadCountValues = this->getThreadIndices();
-  threadCountValues.resize(3, Value());
   return getValuesPermutedByThreadMapping(*this, threadCountValues);
 }
 
 /// Return the number of threads in the order specified by the
-/// thread_dim_mapping attribute.
-/// Return failure is thread_dim_mapping is not a valid permutation.
+/// mapping attribute.
+/// Return failure is mapping is not a valid permutation.
 FailureOr<SmallVector<OpFoldResult>>
 ForeachThreadOp::getPermutedNumThreads(OpBuilder &b) {
   SmallVector<OpFoldResult> threadCountValues = this->getNumThreads();
-  threadCountValues.resize(3, b.getIndexAttr(1));
   return getValuesPermutedByThreadMapping(*this, threadCountValues);
 }
 
 ForeachThreadOp mlir::scf::getForeachThreadOpThreadIndexOwner(Value val) {
   auto tidxArg = val.dyn_cast<BlockArgument>();
   if (!tidxArg)
     return ForeachThreadOp();
   assert(tidxArg.getOwner() && "unlinked block argument");

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     }
 
     // Create new ForeachThreadOp without any results and drop the automatically
     // introduced terminator.
     rewriter.setInsertionPoint(foreachThreadOp);
     auto newForeachThreadOp = rewriter.create<ForeachThreadOp>(
         foreachThreadOp.getLoc(), /*outputs=*/ValueRange(),
         foreachThreadOp.getNumThreads(),
-        extractFromI64ArrayAttr(foreachThreadOp.getThreadDimMapping()));
+        extractFromStringArrayAttr(foreachThreadOp.getMapping()));
     newForeachThreadOp.getBody()->getTerminator()->erase();
 
     // Move over block contents of the old op.
     SmallVector<Value> replacementBbArgs;
     replacementBbArgs.append(
         newForeachThreadOp.getBody()->getArguments().begin(),
         newForeachThreadOp.getBody()->getArguments().end());
     replacementBbArgs.append(foreachThreadOp.getOutputs().size(), Value());

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
   return llvm::to_vector<4>(
       llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
         return a.cast<IntegerAttr>().getInt();
       }));
 }
 
+/// Extract StringRef values from the assumed ArrayAttr of StringAttr.
+SmallVector<StringRef> extractFromStringArrayAttr(Attribute attr) {
+  return llvm::to_vector(
+      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> StringRef {
+        return a.cast<StringAttr>().getValue();
+      }));
+}
+
 /// Given a value, try to extract a constant Attribute. If this fails, return
 /// the original value.
 OpFoldResult getAsOpFoldResult(Value val) {
   if (!val)
     return OpFoldResult();
   Attribute attr;
   if (matchPattern(val, m_Constant(&attr)))
     return attr;

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

func.func @map_nested_foreach_to_threads_excessive_threads(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> {
   %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, %c900) {
         %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
 
   %name2 = 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] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
 
   return %y : memref<2 x 32 x f32>
 }
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
func.func @map_nested_foreach_to_threads_fewer_threads(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> {
   %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, %c900) {
         %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
 
   %name2 = 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] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
 
   return %y : memref<2 x 32 x f32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
func.func @map_nested_foreach_to_threads_dynamic_trip_count(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token, %c9 : index, %c7 : index) -> memref<2 x 32 x f32> {
   %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, %c900) {
         %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
   return %y : memref<2 x 32 x f32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
   // expected-error @below {{unsupported dynamic blockdim size}}
   transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [128, 4, 1] }
 }
 
 // -----
 
 func.func @map_nested_foreach_to_threads_4d_loop(%x: memref<2x32x32x32xf32>, %y: memref<2x32x32x32xf32>, %stream : !gpu.async.token) -> memref<2x32x32x32xf32> {
   %one = arith.constant 1 : index
   %c2 = arith.constant 1 : index
   %c32 = arith.constant 32 : index
   %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, %k, %l) in (%c2, %c32,%c32,%c32) {
         %4 = memref.load %x[%i, %j, %k, %l] : memref<2x32x32x32xf32>        
         memref.store %4, %y[%i, %j, %k, %l] : memref<2x32x32x32xf32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x", "thread_z"] }
     gpu.terminator
   }
   return %y : memref<2x32x32x32xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
func.func @map_foreach_to_blocks_not_unique(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> {
   %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, %c900) {
         %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
 
     scf.foreach_thread (%i, %j) in (%c7, %c9) {
         %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
         %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
         %6 = math.fma %alpha, %4, %5 : f32
         memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-     }  {thread_dim_mapping = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
 
   return %y : memref<2 x 32 x f32>
 }
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):  
func.func @map_foreach_to_blocks_large_loop(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> {
   %c9 = arith.constant 9 : index
   %c7 = arith.constant 7 : index
 
   scf.foreach_thread (%i, %j) in (%c7, %c65537) {
       %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
       %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
       %6 = math.fma %alpha, %4, %5 : f32
       memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-  }  {thread_dim_mapping = [0, 1, 2]}
+  }  { mapping = ["thread_x", "thread_y"] }
 
   scf.foreach_thread (%i, %j) in (%c7, %c9) {
       %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
       %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
       %6 = math.fma %alpha, %4, %5 : f32
       memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-  }  {thread_dim_mapping = [1, 0, 2]}
+  }  { mapping = ["thread_y", "thread_x"] }
   
   return %y : memref<2 x 32 x f32>
 }
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):  
   %funcop = transform.structured.match ops{["func.func"]} in %arg0   
   // expected-error @below {{could not find a unique topLevel scf.foreach_thread}}   
   %1 = transform.gpu.map_foreach_to_blocks %funcop { generate_gpu_launch }
 }
 
 // -----
 
 func.func @map_foreach_to_blocks_large_loop(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> {
   %one = arith.constant 1 : index
   %c65535 = arith.constant 65535 : index
   scf.foreach_thread (%i, %j) in (%c65535, %c65535) {
       %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32>
       %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32>
       %6 = math.fma %alpha, %4, %5 : f32
       memref.store %6, %y[%i, %j] : memref<2 x 32 x f32>
-  }  {thread_dim_mapping = [0, 1, 2]}
+  }  { mapping = ["thread_x", "thread_y"] }
   return %y : memref<2 x 32 x f32>
 }
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):  
   %funcop = transform.structured.match ops{["func.func"]} in %arg0   
   // expected-error @below {{Trying to launch a GPU kernel with gridDim = (65535, 65535, 1) blockDim = (1, 1, 1). It is larger than the limits.}}
   %1 = transform.gpu.map_foreach_to_blocks %funcop { generate_gpu_launch }
 }
 
 // -----
 

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

//      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]}
+     }  { mapping = ["block_x", "block_y"]}
     gpu.terminator
   }
   return %y : !type
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
-  transform.gpu.map_foreach_to_blocks %funcop { blockDim = [12, 9, 1]}
+  transform.gpu.map_foreach_to_blocks %funcop { gridDim = [12, 9]}
 }
 
 // -----
 
 !type = memref<2 x 32 x f32>
 !type1d = memref<32 x f32>
 
 // CHECK-LABEL: func.func @saxpy2d(
//      CHECK:   gpu.barrier
   %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 = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"]}
      scf.foreach_thread (%i) in (%c12) {
         %7 = memref.load %t[%i] : !type1d
         %8 = arith.addf %alpha, %7 : f32
         memref.store %8, %t[%i] : !type1d
-     }  {thread_dim_mapping = [0, 1, 2]}
+     }  {mapping = ["thread_x"] }
     gpu.terminator
   }
   return %y : !type
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
-  transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [12, 9, 1] }
+  transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [12, 9] }
 }
 
 // -----
 
 !type4d = memref<32x64x4x32xf32>
 
 // CHECK-LABEL: func.func @saxpy4d(
 // CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<32x64x4x32xf32>

 //      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]}
+    }  { mapping =  ["thread_y", "thread_x"] }
+  }  { mapping = ["block_x", "block_y"] }
   return %y : !type4d
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["func.func"]} in %arg0
   %gpuLaunch = transform.gpu.map_foreach_to_blocks %funcop { generate_gpu_launch }
-  transform.gpu.map_nested_foreach_to_threads %gpuLaunch { blockDim = [32, 4, 1] }
+  transform.gpu.map_nested_foreach_to_threads %gpuLaunch { blockDim = [32, 4] }
 }
 
 // -----
 
 !type = memref<2 x 32 x f32>
 !type1d = memref<32 x f32>
 
 // CHECK-LABEL: func.func @saxpy2d_no_barrier(
 func.func @saxpy2d_no_barrier(%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-NOT:   gpu.barrier
 //      CHECK:   return
   %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 = [1, 0, 2]}
+     }  { mapping = ["thread_y", "thread_x"] }
     gpu.terminator
   }
   return %y : !type
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0
-  transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [12, 9, 1], syncAfterDistribute = false }
+  transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [12, 9], syncAfterDistribute = false }
 }

mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir

//  CHECK-SAME:   %[[C:[0-9a-z]+]]: tensor<?x?xf32>
   //      CHECK:   %[[tC:.*]] = tensor.extract_slice %[[C_BLK]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
   //      CHECK:   %[[RES:.*]] = linalg.matmul
   // CHECK-SAME:      ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>)
   // CHECK-SAME:     outs(%[[tC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
   //      CHECK:   scf.foreach_thread.perform_concurrently {
   // CHECK-NEXT:     tensor.parallel_insert_slice %[[RES]] into %[[C_BLK]]{{.*}} :
   // CHECK-SAME:       tensor<?x?xf32> into tensor<?x?xf32>
   // CHECK-NEXT:   }
-  // CHECK-NEXT: } {thread_dim_mapping = [1, 0]}
+  // CHECK-NEXT: } {mapping = ["thread_y", "thread_x"]}
     %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
                       outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
     return %0 : tensor<?x?xf32>
   }
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapped to dims [1, 0])
+    %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapped to dims ["thread_y", "thread_x"])
   }
 }
 
 // -----
 
 // Tests that dimension 0 can eliminate affine.min/max, dimension 1 cannot.
 
 // CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * -15 + 300, 15)>
%result = linalg.generic {indexing_maps = [
         linalg.yield %2 : f32
     } -> tensor<4xf32>
     return %result : tensor<4xf32>
   }
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-    %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [2] (mapped to dims [0])
+    %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [2] (mapped to dims ["thread_x"])
   }
 }
 // CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2)>
 
 // CHECK-LABEL: extract_source(
 //       CHECK:  %[[C2:.*]] = arith.constant 2 : index
 //       CHECK:  scf.foreach_thread (%[[ARG:.*]]) in (%[[C2]]) shared_outs(%{{.*}} = %{{.*}}) -> (tensor<4xf32>) {
 //       CHECK:    %[[OFF:.*]] = affine.apply #[[$map0]](%[[ARG]])

mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir

%4 = scf.foreach_thread (%arg0) in (%c320) shared_outs(%arg1 = %2) -> (tensor<320xf32>) {
     %7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<1xf32>) -> tensor<1xf32>
     // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
     %8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<1xf32>) -> tensor<1xf32>
 
     scf.foreach_thread.perform_concurrently {
       // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]}
       tensor.parallel_insert_slice %8 into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32>
     }
-  } {thread_dim_mapping = []}
+  } {mapping = []}
   return %4 : tensor<320xf32>
 }
 
 // -----
 
 // CHECK-LABEL: different_repetitive_region_via_alias
 func.func @different_repetitive_region_via_alias(%arg0: tensor<4xf32>,
                                                  %arg1: tensor<4xf32>,

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> {
       // CHECK: tensor.extract_slice
       // CHECK: scf.foreach_thread.perform_concurrently
       // CHECK: tensor.parallel_insert_slice %{{.*}} into %[[o]]
       %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
       scf.foreach_thread.perform_concurrently {
         tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
           tensor<1xf32> into tensor<100xf32>
       }
-  // CHECK: } {thread_dim_mapping = [5]}
-  } {thread_dim_mapping = [5]}
+  // CHECK: } {mapping = ["none"]}
+  } {mapping = ["none"]}
   return
 }

mlir/test/Dialect/SCF/ops.mlir

func.func @simple_example(%in: tensor<100xf32>, %out: tensor<100xf32>) {
   return
 }
 
 // CHECK-LABEL: func.func @elide_terminator
 func.func @elide_terminator() -> () {
   %num_threads = arith.constant 100 : index
 
   //      CHECK:    scf.foreach_thread
-  // CHECK-NEXT:  } {thread_dim_mapping = [42]}
+  // CHECK-NEXT:  } {mapping = ["none"]}
   // CHECK-NEXT:  return
   scf.foreach_thread (%thread_idx) in (%num_threads) {
     scf.foreach_thread.perform_concurrently {
     }
-  } {thread_dim_mapping = [42]}
+  } {mapping = ["none"]}
   return
 }
 
 // CHECK-LABEL: @switch
 func.func @switch(%arg0: index) -> i32 {
   // CHECK: %{{.*}} = scf.index_switch %arg0 -> i32
   %0 = scf.index_switch %arg0 -> i32
   // CHECK-NEXT: case 2 {

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

struct RewriteExtractSliceFromCollapseShapeUsingScfForeach
   RewriteExtractSliceFromCollapseShapeUsingScfForeach(MLIRContext *context)
       : RewriteExtractSliceFromCollapseShapeBase(context) {}
   LogicalResult emitReplacement(tensor::ExtractSliceOp op, Value dest,
                                 tensor::ExtractSliceFromCollapseHelper &helper,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     auto foreachOp = rewriter.create<scf::ForeachThreadOp>(
         loc, /*outputs=*/dest, /*numThreads=*/helper.getIterationSpaceSizes(),
-        /*threadDimMapping=*/ArrayRef<int64_t>{},
+        /*mapping=*/ArrayRef<StringRef>{},
         [&](OpBuilder &nestedBuilder, Location loc, ValueRange regionArgs) {
           unsigned numThreadIdRegionArgs =
               helper.getIterationSpaceSizes().size();
           unsigned numOutputRegionArgs =
               regionArgs.size() - numThreadIdRegionArgs;
           ValueRange outputIvs = regionArgs.take_front(numThreadIdRegionArgs);
           ValueRange outputArgs = regionArgs.take_back(numOutputRegionArgs);
           assert(outputArgs.size() == 1 &&