diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -985,8 +985,7 @@
         llvm::dbgs() << "\n");
     // clang-format on
     for (auto [activeId, activeMappingSize, availableMappingSize] :
-         llvm::zip_equal(activeIdOps, activeMappingSizes,
-                         availableMappingSizes)) {
+         llvm::zip(activeIdOps, activeMappingSizes, availableMappingSizes)) {
       if (activeMappingSize > availableMappingSize) {
         return definiteFailureHelper(
             transformOp, forallOp,
@@ -1139,8 +1138,15 @@
   }
 
   SmallVector<int64_t> gridDims{getGridDims()};
-  if (!getGenerateGpuLaunch() && gridDims.size() != 3)
-    return transformOp.emitDefiniteFailure("transform require size-3 mapping");
+  if (gridDims.size() != 3) {
+    if (!getGenerateGpuLaunch() || gridDims.size() > 3)
+      return transformOp.emitDefiniteFailure(
+          "transform requires a size-3 mapping");
+
+    // Fill up grid with unit dimensions to 3D.
+    while (gridDims.size() < 3)
+      gridDims.push_back(1);
+  }
 
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(topLevelForallOp);