[mlir][gpu][transform] Fix regressions in MapForallToBlocks.

[ingomueller-net]

Summary

This fixes two regressions introduced by https://reviews.llvm.org/D155941:
(1) it may be that more dimensions are extracted from an existing
scf.forall op than there are dimensions to map to, so these iterators
should be combined with zip and not zip_longest; and (2) if
generate_gpu_launch is set, the grid has to be 3D, but previously,
lower-dimensional grids were also accepted, so this patch trivially
extends these grids with unit dimensions.

[nicolasvasilache]

this does not seem right

[ingomueller-net]

Abandoning because 92f088d335e4a55f10e74c541ac52ec2a4d4ceeb fixes the issues and has landed.

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

LLVM_DEBUG(
         llvm::dbgs() << "\n"; 
         llvm::interleaveComma(
           availableMappingSizes, DBGS() << "----availableMappingSizes: ");
         llvm::dbgs() << "\n";
         llvm::interleaveComma(activeIdOps, DBGS() << "----activeIdOps: ");
         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,
             "Trying to map to fewer GPU threads than loop iterations but "
             "overprovisioning is not yet supported. "
             "Try additional tiling of the before mapping or map to more "
             "threads.");
       }
DiagnosedSilenceableFailure transform::MapForallToBlocks::applyToOne(
   DiagnosedSilenceableFailure diag = mlir::transform::gpu::findTopLevelForallOp(
       target, topLevelForallOp, transformOp);
   if (!diag.succeeded()) {
     diag.attachNote(target->getLoc()) << "when applied to this payload op";
     return diag;
   }
 
   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);
 
   // Generate gpu launch here and move the forall inside
   if (getGenerateGpuLaunch()) {
     DiagnosedSilenceableFailure diag =
         createGpuLaunch(rewriter, target->getLoc(), transformOp, gpuLaunch);