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

[Stanislav Bardyuk]

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

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

>From 11e249f47fec5a621b4e5e3427e9374d98b1ec6e Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk at google.com>
Date: Sun, 1 Nov 2026 12:53:28 +0000
Subject: [PATCH] [mlir][scf-to-gpu] Don't assert when a muli upper bound has
 no derivable factor

deriveStaticUpperBound derives a bound for each factor of an arith.muli and
reads the constant out of the result. When a factor has no derivable bound
(e.g. 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. Check both results first, so the caller's
existing "cannot derive loop-invariant upper bound" path leaves the nest as
scf.parallel instead.

Fixes #231871
---
 mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp     | 28 +++++++++++--------
 .../Conversion/SCFToGPU/parallel_loop.mlir    | 23 +++++++++++++++
 2 files changed, 39 insertions(+), 12 deletions(-)

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