[Mlir-commits] [mlir] [mlir][scf-to-gpu] Don't assert when a muli upper bound has no derivable factor (PR #237725)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Stanislav Bardyuk (kodlan)

<details>
<summary>Changes</summary>

deriveStaticUpperBound handles arith.muli by deriving a bound for each factor and reading the constant out of the result. When a factor has no derivable bound (for example the induction variable of the enclosing scf.parallel), the recursive call returns a null Value, and getDefiningOp on it hits the dyn_cast on a non-existent value assertion. The other branches already return an empty value in that case, so the caller's "cannot derive loop-invariant upper bound" failure path never got a chance to run.

This checks both recursive results before using them. The nest is then left as scf.parallel, the same thing that already happens when the inner bound is the outer induction variable directly.

Fixes #<!-- -->231871

---
Full diff: https://github.com/llvm/llvm-project/pull/237725.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp (+16-12) 
- (modified) mlir/test/Conversion/SCFToGPU/parallel_loop.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index 38e49fc0b71d7e..65eba8da0562ed 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
@@ -339,18 +339,22 @@ static Value deriveStaticUpperBound(Value upperBound,
   }
 
   if (auto multiplyOp = upperBound.getDefiningOp<arith::MulIOp>()) {
-    if (auto lhs = deriveStaticUpperBound(multiplyOp.getOperand(0), rewriter)
-                       .getDefiningOp<arith::ConstantIndexOp>())
-      if (auto rhs = deriveStaticUpperBound(multiplyOp.getOperand(1), rewriter)
-                         .getDefiningOp<arith::ConstantIndexOp>()) {
-        // Assumptions about the upper bound of minimum computations no longer
-        // work if multiplied by mixed signs, so abort in this case.
-        if ((lhs.value() < 0) != (rhs.value() < 0))
-          return {};
-
-        return arith::ConstantIndexOp::create(rewriter, multiplyOp.getLoc(),
-                                              lhs.value() * rhs.value());
-      }
+    // Either factor may have no derivable bound (e.g. an outer parallel IV).
+    Value lhsBound = deriveStaticUpperBound(multiplyOp.getOperand(0), rewriter);
+    if (!lhsBound)
+      return {};
+    Value rhsBound = deriveStaticUpperBound(multiplyOp.getOperand(1), rewriter);
+    if (!rhsBound)
+      return {};
+    auto lhs = cast<arith::ConstantIndexOp>(lhsBound.getDefiningOp());
+    auto rhs = cast<arith::ConstantIndexOp>(rhsBound.getDefiningOp());
+    // Assumptions about the upper bound of minimum computations no longer
+    // work if multiplied by mixed signs, so abort in this case.
+    if ((lhs.value() < 0) != (rhs.value() < 0))
+      return {};
+
+    return arith::ConstantIndexOp::create(rewriter, multiplyOp.getLoc(),
+                                          lhs.value() * rhs.value());
   }
 
   return {};
diff --git a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir b/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir
index 8eccf0b66a60d5..c9ed6b3e127f37 100644
--- a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir
+++ b/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir
@@ -734,3 +734,26 @@ func.func @parallel_loop_unsigned(%arg0 : index, %arg1 : index, %arg2 : index,
   } { mapping = [#gpu.loop_dim_map<processor = block_x, map = (d0) -> (d0), bound = (d0) -> (d0)>] }
   return
 }
+
+// -----
+
+// Loop-variant upper bound through arith.muli. Cannot be mapped (used to crash).
+
+func.func @parallel_loop_loop_variant_muli_bound() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  scf.parallel (%i) = (%c0) to (%c4) step (%c1) {
+    %ub = arith.muli %i, %c4 : index
+    scf.parallel (%j) = (%c0) to (%ub) step (%c1) {
+      scf.reduce
+    } {mapping = [#gpu.loop_dim_map<processor = thread_x, map = (d0) -> (d0), bound = (d0) -> (d0)>]}
+    scf.reduce
+  } {mapping = [#gpu.loop_dim_map<processor = block_x, map = (d0) -> (d0), bound = (d0) -> (d0)>]}
+  return
+}
+
+// CHECK-LABEL: @parallel_loop_loop_variant_muli_bound
+// CHECK-NOT: gpu.launch
+// CHECK: scf.parallel
+// CHECK: scf.parallel

``````````

</details>


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