[mlir][Transform] Add support for mapping to GPU warps and to linear ids

[nicolasvasilache]

Summary

This revisions refactors the implementation of mapping to threads to additionally allow warps and linear ids to be specified.

warp_dims is currently specified along with block_dims as a transform attribute.

Linear ids on th other hand use the flattened block_dims to predicate on the first (linearized) k threads.
An additional GPULinearIdMappingAttr is added to the GPU dialect to allow specifying loops mapped to this new scheme.

Various implementation and transform op semantics cleanups are also applied.

[nicolasvasilache]

Merge mapping to warps and linear ids.

[nicolasvasilache]

Delete unused code and better handle insertion points.

[nicolasvasilache]

Add comment

[nicolasvasilache]

Rebase

[bondhugula]

Please split out the IR/ changes out from the TransformOps/ changes if appropriate; add test cases for linear_id/lane_id op and its lowering - it doesn't seem to itself appear in test cases. It looks odd that the only code manipulating a GPU IR dialect op appears in TransformOps/. Commit summary appears incomplete - doesn't seems to capture all the changes.

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

124 ↗

(On Diff #505973)

This doesn't seem to have any test cases here. Is this PR ready for review?

124 ↗

(On Diff #505973)

linear_id or GPU_LaneIdOp to be consistent.

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

79 ↗

(On Diff #505973)

This description needs to be completed (recommend focusing on "what" instead of what it allows in general).

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

36

Doc comment.

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

81–123

Missing code comments for these.

[nicolasvasilache]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

124 ↗

(On Diff #505973)

Good point, there is indeed no value in introducing new GPU ops that do not lower through the common pass for the purpose of the transformations I am implementing.

Refactored the implementation to avoid relying on a new op.

This op may make sense to be introduced separately and lowered through the existing schemes that are mindful of index type length. But this is indeed a separate concern and a separate PR.

[nicolasvasilache]

Address comments and refactor.

[bondhugula]

Dismissing change request that has now been addressed. This is now purely a GPU transform ops related revision. Any of the other reviewers familiar/interested in transforms ops may be able to review it.

[ThomasRaoux]

LGTM, added a few nits

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

68 ↗

(On Diff #506093)

Doesn't have to be addressed in this patch but DimX, DimY and DimZ don't really make sense for linearID, I wonder if we could just have an index to give the order of distribution?

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

52

nit: replace A the -> The?

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

482

nit: doesn't feel like a very useful comment

[nicolasvasilache]

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

68 ↗

(On Diff #506093)

ack, we can also go higher-D than 3

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

 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h.inc"
 
 namespace mlir {
 class DialectRegistry;
 namespace transform {
 namespace gpu {
 
+constexpr int64_t kWarpSize = 32;
+
+/// Helper type for functions that generate ids for the mapping of a
+/// scf.forall.
+using GpuIdBuilderFnType = llvm::function_ref<SmallVector<Value>(
+    RewriterBase &, scf::ForallOp, ArrayRef<int64_t> mappingDims)>;
+
+/// Helper struct for passing the mapping attributes and id generator to the
+/// common forall rewriter.
+struct GpuIdBuilder {
+  /// The mapping attributes targeted by this generator.
+  SmallVector<DeviceMappingAttrInterface> mappingAttributes;
+  /// The constructor that builds the concrete IR for mapping ids.
+  GpuIdBuilderFnType idBuilder;
+};
+
 /// Map the top level `scf.forall` op to GPU Thread Blocks.
 /// Mapping is one-to-one and the induction variables of `scf.forall` are
-/// rewritten to gpu.block_id according to the thread_dim_apping attribute.
+/// rewritten to gpu.block_id according to the thread_dim_mapping attribute.
 /// Dynamic, `scf.forall` trip counts are currently not supported.
 /// Dynamic block dim sizes are currently not supported.
-DiagnosedSilenceableFailure mapForallToBlocksImpl(
-    RewriterBase &rewriter, TransformOpInterface transformOp,
-    scf::ForallOp forallOp, SmallVectorImpl<int64_t> &gridDims,
-    const ArrayRef<DeviceMappingAttrInterface> &mappingAttributes,
-    function_ref<void(RewriterBase &, scf::ForallOp, SmallVectorImpl<Value> &)>
-        blockIdGenerator);
+DiagnosedSilenceableFailure
+mapForallToBlocksImpl(RewriterBase &rewriter, TransformOpInterface transformOp,
+                      scf::ForallOp forallOp,
+                      SmallVectorImpl<int64_t> &gridDims,
+                      const GpuIdBuilder &gpuIdBuilder);
 
 /// Search `scf.forall` ops nested under `target` and map each such op to GPU
 /// threads. Mapping is one-to-one and the induction variables of `scf.forall`
 /// are rewritten to gpu.thread_id according to the thread_dim_mapping
 /// attribute.
 /// Sibling `scf.forall` are supported in which case, the union of the number of
 /// threads is computed and may result in predication.
 /// Dynamic, `scf.forall` trip counts are currently not supported.
 /// Dynamic block dim sizes are currently not supported.
 DiagnosedSilenceableFailure mapNestedForallToThreadsImpl(
     RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,
     Operation *target, const SmallVectorImpl<int64_t> &kernelBlockDims,
-    bool syncAfterDistribute,
-    const ArrayRef<DeviceMappingAttrInterface> &threadMappingAttributes,
-    function_ref<void(RewriterBase &, scf::ForallOp, SmallVectorImpl<Value> &)>
-        threadIdGenerator);
+    bool syncAfterDistribute, const GpuIdBuilder &gpuIdBuilder);
 
 /// Find the unique top level scf::ForallOp within a given target op.
 DiagnosedSilenceableFailure
 findTopLevelForallOp(Operation *target, scf::ForallOp &topLevelForallOp,
                      TransformOpInterface transformOp);
 
 } // namespace gpu
 } // namespace transform
 
 namespace gpu {
 void registerTransformDialectExtension(DialectRegistry &registry);
 } // namespace gpu
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_TRANSFORMOPS_GPUTRANSFORMOPS_H

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/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.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/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 using namespace mlir::transform;
+using namespace mlir::transform::gpu;
 
 #define DEBUG_TYPE "gpu-transforms"
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
 namespace {
 
-/// Helper type for functions that generate ids for the mapping of a scf.forall.
-using IdGeneratorFnType = llvm::function_ref<void(RewriterBase &, scf::ForallOp,
-                                                  SmallVectorImpl<Value> &)>;
-
-struct MappingToGpuHelper {
-  MappingToGpuHelper(SmallVector<DeviceMappingAttrInterface> mappingAttributes,
-                     IdGeneratorFnType idGenerator)
-      : mappingAttributes(mappingAttributes), idGenerator(idGenerator) {}
-
-  SmallVector<DeviceMappingAttrInterface> mappingAttributes;
-  IdGeneratorFnType idGenerator;
-};
+/// Return a flatten thread id for the workgroup with given sizes.
+static OpFoldResult getLinearThreadId(RewriterBase &rewriter, Location loc) {
+  AffineExpr tx, ty, tz, BDX, BDY;
+  bindDims(rewriter.getContext(), tx, ty, tz);
+  bindSymbols(rewriter.getContext(), BDX, BDY);
+  IndexType indexType = rewriter.getIndexType();
+  SmallVector<OpFoldResult> threadsAndWorkGroups{
+      rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x).getResult(),
+      rewriter.create<ThreadIdOp>(loc, indexType, Dimension::y).getResult(),
+      rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z).getResult()};
+  threadsAndWorkGroups.push_back(
+      rewriter.create<BlockDimOp>(loc, indexType, Dimension::x).getResult());
+  threadsAndWorkGroups.push_back(
+      rewriter.create<BlockDimOp>(loc, indexType, Dimension::y).getResult());
+  return makeComposedFoldedAffineApply(
+      rewriter, loc, tx + ty * BDX + tz * BDX * BDY, threadsAndWorkGroups);
+}
 
-struct MappingToGpuBlocksHelper : public MappingToGpuHelper {
+struct GpuBlockIdBuilder : public GpuIdBuilder {
 
-  MappingToGpuBlocksHelper(MLIRContext *ctx)
-      : MappingToGpuHelper(
-            SmallVector<DeviceMappingAttrInterface>{
-                GPUBlockMappingAttr::get(ctx, Blocks::DimX),
+  GpuBlockIdBuilder(MLIRContext *ctx) : GpuIdBuilder() {
+    mappingAttributes = {GPUBlockMappingAttr::get(ctx, Blocks::DimX),
                          GPUBlockMappingAttr::get(ctx, Blocks::DimY),
                          GPUBlockMappingAttr::get(ctx, Blocks::DimZ)},
-            IdGeneratorFnType{[](RewriterBase &rewriter, scf::ForallOp forallOp,
-                                 SmallVectorImpl<Value> &ids) {
-              OpBuilder::InsertionGuard guard(rewriter);
-              rewriter.setInsertionPoint(forallOp);
-              IndexType indexType = rewriter.getIndexType();
-              auto loc = forallOp->getLoc();
-              ids.assign(
-                  {rewriter.create<BlockIdOp>(loc, indexType, Dimension::x),
-                   rewriter.create<BlockIdOp>(loc, indexType, Dimension::y),
-                   rewriter.create<BlockIdOp>(loc, indexType, Dimension::z)});
-            }}) {}
+    idBuilder = [](RewriterBase &rewriter, scf::ForallOp forallOp,
+                   ArrayRef<int64_t> mappingDims) {
+      return SmallVector<Value>{
+          rewriter.create<BlockIdOp>(loc, indexType, Dimension::x),
+          rewriter.create<BlockIdOp>(loc, indexType, Dimension::z)};
+    };
+  }
 };
 
-struct MappingToGpuThreadsHelper : public MappingToGpuHelper {
-  MappingToGpuThreadsHelper(MLIRContext *ctx)
-      : MappingToGpuHelper(
-            SmallVector<DeviceMappingAttrInterface>{
-                GPUThreadMappingAttr::get(ctx, Threads::DimX),
+struct GpuThreadIdBuilder : public GpuIdBuilder {
+  GpuThreadIdBuilder(MLIRContext *ctx, SmallVector<Value> *idCaptures = nullptr)
+      : GpuIdBuilder() {
+    mappingAttributes = {GPUThreadMappingAttr::get(ctx, Threads::DimX),
                          GPUThreadMappingAttr::get(ctx, Threads::DimY),
-                GPUThreadMappingAttr::get(ctx, Threads::DimZ)},
-            IdGeneratorFnType{[](RewriterBase &rewriter, scf::ForallOp forallOp,
-                                 SmallVectorImpl<Value> &ids) {
-              OpBuilder::InsertionGuard guard(rewriter);
-              rewriter.setInsertionPoint(forallOp);
-              IndexType indexType = rewriter.getIndexType();
-              auto loc = forallOp->getLoc();
-              ids.assign(
-                  {rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x),
-                   rewriter.create<ThreadIdOp>(loc, indexType, Dimension::y),
-                   rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z)});
-            }}) {}
+                         GPUThreadMappingAttr::get(ctx, Threads::DimZ)};
+    idBuilder = [](RewriterBase &rewriter, scf::ForallOp forallOp,
+                   ArrayRef<int64_t> mappingDims) {
+      return SmallVector<Value>{
+          rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x),
+          rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z)};
+    };
+  }
+};
+
+struct GpuWarpIdBuilder : public GpuIdBuilder {
+  GpuWarpIdBuilder(MLIRContext *ctx, SmallVector<Value> *idCaptures = nullptr)
+      : GpuIdBuilder() {
+    mappingAttributes = {GPUWarpMappingAttr::get(ctx, Warps::DimX),
+                         GPUWarpMappingAttr::get(ctx, Warps::DimY),
+                         GPUWarpMappingAttr::get(ctx, Warps::DimZ)};
+    idBuilder = [](RewriterBase &rewriter, scf::ForallOp forallOp,
+                   ArrayRef<int64_t> mappingDims) {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(forallOp);
+      Location loc = forallOp.getLoc();
+      Value warpId = rewriter.create<SubgroupIdOp>(loc);
+      SmallVector<int64_t> reverseBlockDims(llvm::reverse(mappingDims));
+      SmallVector<int64_t> strides = computeStrides(reverseBlockDims);
+      AffineExpr d0;
+      bindDims(rewriter.getContext(), d0);
+      SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
+      SmallVector<Value> ids;
+      for (AffineExpr e : delinearizingExprs)
+        ids.push_back(makeComposedAffineApply(rewriter, loc, e, warpId));
+      return ids;
+    };
+  }
 };
 
 } // namespace
 
 static DiagnosedSilenceableFailure
 failureHelper(std::optional<TransformOpInterface> transformOp,
               scf::ForallOp forallOp, const Twine &message) {
   if (transformOp.has_value())
-    return emitDefiniteFailure(*transformOp, message);
+    return transformOp->emitSilenceableError() << message;
   return emitDefiniteFailure(forallOp, message);
 }
 
 /// Check if given mapping attributes are one of the desired attributes
 static DiagnosedSilenceableFailure
 checkMappingAttributeTypes(std::optional<TransformOpInterface> transformOp,
                            scf::ForallOp forallOp) {
   if (!forallOp.getMapping().has_value())
     return failureHelper(transformOp, forallOp, "mapping must be present");
 
   bool hasBlockMapping =
       llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
         return attr.isa<GPUBlockMappingAttr>();
       });
   bool hasThreadMapping =
       llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
         return attr.isa<GPUThreadMappingAttr>();
       });
+  bool hasWarpMapping =
+      llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
+        return attr.isa<GPUWarpMappingAttr>();
+      });
   int64_t countMappingTypes = 0;
   countMappingTypes += hasBlockMapping ? 1 : 0;
   countMappingTypes += hasThreadMapping ? 1 : 0;
+  countMappingTypes += hasWarpMapping ? 1 : 0;
   if (countMappingTypes > 1) {
     return failureHelper(transformOp, forallOp,
                          "cannot mix different mapping types, use nesting");
   }
 
   DenseSet<Attribute> seen;
   for (Attribute map : forallOp.getMapping()->getValue()) {
     if (llvm::is_contained(seen, map)) {
if (llvm::any_of(forallOp.getMixedUpperBound(), [&](OpFoldResult ofr) {
         return !getConstantIntValue(ofr).has_value();
       })) {
     return failureHelper(transformOp, forallOp,
                          "unsupported dynamic sizes in forall op");
   }
   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.
+/// 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, std::optional<int64_t> gridDimX,
     std::optional<int64_t> gridDimY, std::optional<int64_t> gridDimZ,
     std::optional<int64_t> blockDimX, std::optional<int64_t> blockDimY,
     std::optional<int64_t> blockDimZ) {
 
   static constexpr int maxTotalBlockdim = 1024;
   static constexpr int maxBlockdimx = 1024;
return transformOp.emitSilenceableError()
            << gridDimX.value_or(1) << ", " << gridDimY.value_or(1) << ", "
            << gridDimZ.value_or(1) << ") blockDim = (" << blockDimX.value_or(1)
            << ", " << blockDimY.value_or(1) << ", " << blockDimZ.value_or(1)
            << "). It is larger than the limits.";
   }
   return DiagnosedSilenceableFailure::success();
 }
 
-/// Creates an empty-body gpu::LaunchOp using the provided kernel settings and
-/// put a terminator within.
+/// Creates an empty-body gpu::LaunchOp using the provided kernel settings
+/// and put a terminator within.
 static DiagnosedSilenceableFailure
 createGpuLaunch(RewriterBase &rewriter, Location loc,
                 TransformOpInterface transformOp, LaunchOp &launchOp,
                 std::optional<int64_t> gridDimX = std::nullopt,
                 std::optional<int64_t> gridDimY = std::nullopt,
                 std::optional<int64_t> gridDimZ = std::nullopt,
                 std::optional<int64_t> blockDimX = std::nullopt,
                 std::optional<int64_t> blockDimY = std::nullopt,
if (blockDimY.has_value())
     gpuLaunch.getBlockSizeYMutable().assign(
         createConstValue(blockDimY.value()));
   if (blockDimZ.has_value())
     gpuLaunch.getBlockSizeZMutable().assign(
         createConstValue(blockDimZ.value()));
   return DiagnosedSilenceableFailure::success();
 }
 
+/// Struct to return the result of the rewrite of a forall operation.
+struct ForallRewriteResult {
+  SmallVector<int64_t> mappingSizes;
+  SmallVector<Value> mappingIds;
+};
+
+/// Helper to replace ids of dimensions known to be 1 by 0 to simplify the IR.
+static void
+replaceUnitMappingIdsHelper(RewriterBase &rewriter, Location loc,
+                            ValueRange mappingIds,
+                            ArrayRef<int64_t> availableMappingSizes) {
+  assert(!mappingIds.empty() && "expected some mapping");
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(mappingIds.front().getDefiningOp());
+  Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+  for (auto [dim, id] : llvm::zip_equal(availableMappingSizes, mappingIds)) {
+    if (dim == 1)
+      rewriter.replaceAllUsesWith(id, zero);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // MapForallToBlocks
 //===----------------------------------------------------------------------===//
 
-static FailureOr<SmallVector<int64_t>> rewriteOneForallCommonImpl(
+static FailureOr<ForallRewriteResult> rewriteOneForallCommonImpl(
     RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,
     scf::ForallOp forallOp,
     const SmallVectorImpl<int64_t> &availableMappingSizes,
-    const ArrayRef<DeviceMappingAttrInterface> &allMappingAttributes,
-    IdGeneratorFnType idGenerator) {
+    const GpuIdBuilder &gpuIdBuilder) {
   LDBG("Start rewriteOneForallCommonImpl");
 
   // Step 0. GPU-specific verifications. There is no better place to anchor
-  // those right now: the ForallOp is target-independent and the transform op
-  // does not apply to individual ForallOp.
+  // those right now: the ForallOp is target-independent and the transform
+  // op does not apply to individual ForallOp.
   DiagnosedSilenceableFailure diag = verifyGpuMapping(transformOp, forallOp);
   if (!diag.succeeded())
     return failure();
 
   // Step 1. Complete the mapping to a full mapping (with 1s) if necessary.
   SmallVector<int64_t> tmpMappingSizes = llvm::to_vector(
       llvm::map_range(forallOp.getMixedUpperBound(), [](OpFoldResult ofr) {
         auto maybeStaticValue = getConstantIntValue(ofr);
         assert(maybeStaticValue && "expected static value");
         return maybeStaticValue.value();
       }));
   SmallVector<Attribute> forallMappings =
       llvm::to_vector(forallOp.getMapping()->getValue());
-  for (auto attr : allMappingAttributes) {
+  for (auto attr : gpuIdBuilder.mappingAttributes) {
     if (llvm::is_contained(forallMappings, attr))
       continue;
     forallMappings.push_back(attr);
     tmpMappingSizes.push_back(1);
   }
 
-  // Step 2. sort the values by the corresponding DeviceMappingAttrInterface.
+  // Step 2. sort the values by the corresponding
+  // DeviceMappingAttrInterface.
   auto comparator = [&](DeviceMappingAttrInterface a,
                         DeviceMappingAttrInterface b) -> bool {
     return a.getMappingId() < b.getMappingId();
   };
   SmallVector<int64_t> mappingSizes =
       getValuesSortedByKey(forallMappings, tmpMappingSizes, comparator);
   LLVM_DEBUG(llvm::interleaveComma(mappingSizes, DBGS() << "mappingSizes: ");
              llvm::dbgs() << "\n";
              llvm::interleaveComma(forallMappings, DBGS() << "mappingAttrs: ");
              llvm::dbgs() << "\n");
 
-  // Step 3. Generate the mappingIdOps using the provided generator and map the
-  // induction variables to the newly created ops. Replace ids of dimension
-  // known to be of size 1 by zero to simplify the IR.
-  SmallVector<Value> mappingIdOps;
-  Location loc = forallOp.getLoc();
-  idGenerator(rewriter, forallOp, mappingIdOps);
+  // Step 3. Generate the mappingIdOps using the provided generator and map
+  // the induction variables to the newly created ops.
+  SmallVector<Value> mappingIdOps =
+      gpuIdBuilder.idBuilder(rewriter, forallOp, mappingSizes);
   LLVM_DEBUG(llvm::interleaveComma(mappingIdOps, DBGS() << "mappingIdOps: ");
              llvm::dbgs() << "\n");
   assert(mappingIdOps.size() == mappingSizes.size() && "expect equal sizes");
-  Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-  if (!availableMappingSizes.empty()) {
-    for (size_t i : llvm::seq(size_t(0), availableMappingSizes.size())) {
-      if (availableMappingSizes[i] == 1)
-        mappingIdOps[i] = zero;
-    }
-  }
 
   IRMapping bvm;
   for (auto [iv, dim] :
        llvm::zip_equal(forallOp.getInductionVars(),
                        ArrayRef<Attribute>{forallMappings}.take_front(
                            forallOp.getInductionVars().size()))) {
     Value peIdOp = mappingIdOps[static_cast<int64_t>(
         dim.cast<DeviceMappingAttrInterface>().getMappingId())];
     bvm.map(iv, peIdOp);
   }
 
   // Step 4. Maybe create conditionals to predicate the region.
   // Skip this step when availableMappingSizes is empty.
+  Location loc = forallOp.getLoc();
   Value predicate;
   if (!availableMappingSizes.empty()) {
     LLVM_DEBUG(llvm::interleaveComma(availableMappingSizes,
                                      DBGS() << "availableMappingSizes: ");
                llvm::dbgs() << "\n");
     for (auto [id, mappingSize, availableMappingSize] :
          llvm::zip_equal(mappingIdOps, mappingSizes, availableMappingSizes)) {
       if (mappingSize > availableMappingSize) {
static FailureOr<ForallRewriteResult> rewriteOneForallCommonImpl(
 
   // Step 5. Move the body of forallOp.
   // Erase the terminator first, it will not be used.
   rewriter.eraseOp(forallOp.getTerminator());
   Block *targetBlock;
   Block::iterator insertionPoint;
   if (predicate) {
     // Step 5.a. If predicated, move at the beginning.
-    auto ifOp =
-        rewriter.create<scf::IfOp>(loc, predicate, /*withElseRegion=*/false);
+    auto ifOp = rewriter.create<scf::IfOp>(loc, predicate,
+                                           /*withElseRegion=*/false);
     targetBlock = ifOp.thenBlock();
     insertionPoint = ifOp.thenBlock()->begin();
   } else {
-    // Step 5.b. Otherwise, move inline just at the rewriter insertion point.
+    // Step 5.b. Otherwise, move inline just at the rewriter insertion
+    // point.
     targetBlock = forallOp->getBlock();
     insertionPoint = rewriter.getInsertionPoint();
   }
   Block &sourceBlock = forallOp.getRegion().front();
   targetBlock->getOperations().splice(insertionPoint,
                                       sourceBlock.getOperations());
 
   // Step 6. RAUW thread indices to thread ops.
   for (Value loopIndex : forallOp.getInductionVars()) {
     Value threadIdx = bvm.lookup(loopIndex);
     rewriter.replaceAllUsesWith(loopIndex, threadIdx);
   }
 
   // Step 7. Erase old op.
   rewriter.eraseOp(forallOp);
 
-  return mappingSizes;
+  return ForallRewriteResult{mappingSizes, mappingIdOps};
 }
 
 DiagnosedSilenceableFailure mlir::transform::gpu::mapForallToBlocksImpl(
     RewriterBase &rewriter, TransformOpInterface transformOp,
     scf::ForallOp forallOp, SmallVectorImpl<int64_t> &gridDims,
-    const ArrayRef<DeviceMappingAttrInterface> &allMappingAttributes,
-    IdGeneratorFnType idGenerator) {
+    const GpuIdBuilder &gpuIdBuilder) {
+
   // Pass an empty anyAvailableMappingSizes.
+  Location loc = forallOp.getLoc();
   SmallVector<int64_t> anyAvailableMappingSizes;
-  FailureOr<SmallVector<int64_t>> maybeMappingSizes =
-      rewriteOneForallCommonImpl(rewriter, transformOp, forallOp,
-                                 anyAvailableMappingSizes, allMappingAttributes,
-                                 idGenerator);
-  if (failed(maybeMappingSizes))
+  FailureOr<ForallRewriteResult> rewriteResult = rewriteOneForallCommonImpl(
+      rewriter, transformOp, forallOp, anyAvailableMappingSizes, gpuIdBuilder);
+
+  // Fail if anything goes wrong.
+  if (failed(rewriteResult))
     return DiagnosedSilenceableFailure::definiteFailure();
-  gridDims = *maybeMappingSizes;
+  gridDims = rewriteResult->mappingSizes;
+
+  // Replace ids of dimensions known to be 1 by 0 to simplify the IR.
+  // Here, the result of mapping determines the available mapping sizes.
+  replaceUnitMappingIdsHelper(rewriter, loc, rewriteResult->mappingIds,
+                              gridDims);
+
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 mlir::transform::gpu::findTopLevelForallOp(Operation *target,
                                            scf::ForallOp &topLevelForallOp,
                                            TransformOpInterface transformOp) {
   auto walkResult = target->walk([&](scf::ForallOp forallOp) {
if (getGenerateGpuLaunch()) {
     rewriter.eraseOp(topLevelForallOp);
     topLevelForallOp = cast<scf::ForallOp>(newForallOp);
   }
 
   diag = verifyGpuMapping(transformOp, topLevelForallOp);
   if (!diag.succeeded())
     return diag;
 
-  MappingToGpuBlocksHelper helper(getContext());
+  GpuBlockIdBuilder gpuBlockIdBuilder(getContext());
   diag = mlir::transform::gpu::mapForallToBlocksImpl(
-      rewriter, transformOp, topLevelForallOp, gridDims,
-      helper.mappingAttributes, helper.idGenerator);
+      rewriter, transformOp, topLevelForallOp, gridDims, gpuBlockIdBuilder);
   if (!diag.succeeded())
     return diag;
 
   diag = alterGpuLaunch(rewriter, gpuLaunch,
                         cast<TransformOpInterface>(getOperation()), gridDims[0],
                         gridDims[1], gridDims[2]);
 
   results.push_back(gpuLaunch);
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // MapNestedForallToThreads
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(
     RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,
     Operation *target, const SmallVectorImpl<int64_t> &kernelBlockDims,
-    bool syncAfterDistribute,
-    const ArrayRef<DeviceMappingAttrInterface> &allMappingAttributes,
-    IdGeneratorFnType idGenerator) {
+    bool syncAfterDistribute, const GpuIdBuilder &gpuIdBuilder) {
   DiagnosedSilenceableFailure diag = DiagnosedSilenceableFailure::success();
   target->walk([&](scf::ForallOp forallOp) {
-    // Ignore cases with different attributes.
+    // Ignore cases with different attributes than this builder supports.
     for (Attribute map : forallOp.getMapping()->getValue()) {
-      if (!llvm::is_contained(allMappingAttributes, map)) {
+      if (!llvm::is_contained(gpuIdBuilder.mappingAttributes, map)) {
         return WalkResult::skip();
       }
     }
     diag = verifyGpuMapping(transformOp, forallOp);
     if (diag.succeeded()) {
-      // Take the loc ahead of time
       Location loc = forallOp.getLoc();
       OpBuilder::InsertionGuard g(rewriter);
+      // Insert after to allow for syncthreads after `forall` is erased.
       rewriter.setInsertionPointAfter(forallOp);
-      if (failed(rewriteOneForallCommonImpl(rewriter, transformOp, forallOp,
-                                            kernelBlockDims,
-                                            allMappingAttributes, idGenerator)))
+      FailureOr<ForallRewriteResult> rewriteResult = rewriteOneForallCommonImpl(
+          rewriter, transformOp, forallOp, kernelBlockDims, gpuIdBuilder);
+
+      // Fail if anything goes wrong.
+      if (failed(rewriteResult))
         diag = DiagnosedSilenceableFailure::definiteFailure();
+
       // Add a syncthreads if needed. TODO: warpsync
       if (syncAfterDistribute)
         rewriter.create<BarrierOp>(loc);
+
+      // Replace ids of dimensions known to be 1 by 0 to simplify the IR.
+      // Here, the available mapping sizes are given by `kernelBlockDims`.
+      replaceUnitMappingIdsHelper(rewriter, loc, rewriteResult->mappingIds,
+                                  kernelBlockDims);
     }
     return diag.succeeded() ? WalkResult::advance() : WalkResult::interrupt();
   });
   return diag;
 }
 
 DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(
     Operation *target, ApplyToEachResultList &results, TransformState &state) {
DiagnosedSilenceableFailure diag =
                      blockDims[0], blockDims[1], blockDims[2]);
   if (diag.isSilenceableFailure()) {
     diag.attachNote(getLoc()) << getBlockDimAttrName() << " is too large";
     return diag;
   }
 
   MLIRContext *ctx = getContext();
   IRRewriter rewriter(ctx);
-  MappingToGpuThreadsHelper helper(ctx);
+  SmallVector<Value> idCaptures;
+  GpuThreadIdBuilder gpuThreadIdBuilder(ctx, &idCaptures);
   diag = mlir::transform::gpu::mapNestedForallToThreadsImpl(
       rewriter, transformOp, target, blockDims, getSyncAfterDistribute(),
-      helper.mappingAttributes, helper.idGenerator);
+      gpuThreadIdBuilder);
 
   if (!diag.succeeded())
     return diag;
 
   diag = alterGpuLaunch(rewriter, gpuLaunch, transformOp, std::nullopt,
                         std::nullopt, std::nullopt, blockDims[0], blockDims[1],
                         blockDims[2]);