[Mlir-commits] [mlir] [mlir][gpu] Fix crash in transform.gpu.map_nested_forall_to_threads with zero iterations (PR #170282)

[Fabian Mora]

================
@@ -773,22 +773,30 @@ static DiagnosedSilenceableFailure checkMappingSpec(
     std::optional<TransformOpInterface> transformOp, scf::ForallOp forallOp,
     ArrayRef<int64_t> numParallelIterations, ArrayRef<int64_t> blockOrGridSizes,
     int factor, bool useLinearMapping = false) {
+  if (llvm::any_of(blockOrGridSizes, [](int64_t i) { return i <= 0; })) {
+    return definiteFailureHelper(transformOp, forallOp,
+                                 "block/grid sizes must be strictly positive");
+  }
   if (!useLinearMapping && blockOrGridSizes.front() % factor != 0) {
     auto diag = definiteFailureHelper(
         transformOp, forallOp,
         Twine("3-D mapping: size of threadIdx.x must be a multiple of ") +
             Twine(factor));
     return diag;
   }
-  if (computeProduct(numParallelIterations) * factor >
-      computeProduct(blockOrGridSizes)) {
+  bool hasZeroParallelIteration =
+      llvm::any_of(numParallelIterations, [](int64_t i) { return i == 0; });
+  int64_t required = hasZeroParallelIteration
+                         ? 0
+                         : computeProduct(numParallelIterations) * factor;
+  int64_t available = computeProduct(blockOrGridSizes);
----------------
fabianmcg wrote:

Please use get more descriptive names, and add a comment

https://github.com/llvm/llvm-project/pull/170282