[Mlir-commits] [mlir] [MLIR][SCF] Fix ForLoopRangeFolding miscompile with non-positive MulIOp multiplier (PR #188995)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188995

The scf-for-loop-range-folding pass transforms loops of the form

  for (i = lb; i < ub; i += step) { use(i * c) }

into

  for (j = lb*c; j < ub*c; j += step*c) { use(j) }

This transformation is only valid when c is strictly positive, since scf.for requires a positive step. When c is zero or negative, the new step becomes zero or effectively negative (wrapping in unsigned arithmetic for index type), producing an incorrect loop.

Add a guard that restricts the MulIOp folding to cases where the loop-invariant multiplier is a statically known positive integer constant. Non-constant loop-invariant multipliers are also excluded since their sign cannot be determined at compile time.

Fixes #56235
Fixes #116664

Assisted-by: Claude Code

>From 6b89ee850af8f74284befd2e952437aa8b166740 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 16:58:21 -0700
Subject: [PATCH] [MLIR][SCF] Fix ForLoopRangeFolding miscompile with
 non-positive MulIOp multiplier

The scf-for-loop-range-folding pass transforms loops of the form

  for (i = lb; i < ub; i += step) { use(i * c) }

into

  for (j = lb*c; j < ub*c; j += step*c) { use(j) }

This transformation is only valid when c is strictly positive, since
scf.for requires a positive step. When c is zero or negative, the new
step becomes zero or effectively negative (wrapping in unsigned
arithmetic for index type), producing an incorrect loop.

Add a guard that restricts the MulIOp folding to cases where the
loop-invariant multiplier is a statically known positive integer
constant. Non-constant loop-invariant multipliers are also excluded
since their sign cannot be determined at compile time.

Fixes #56235
Fixes #116664

Assisted-by: Claude Code
---
 .../SCF/Transforms/LoopRangeFolding.cpp       | 12 ++-
 mlir/test/Dialect/SCF/loop-range.mlir         | 87 +++++++++++++++++++
 2 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index 0c2a71eb25fe5..d4b4c4cfc063f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/IRMapping.h"
 
 namespace mlir {
@@ -71,7 +72,16 @@ void ForLoopRangeFolding::runOnOperation() {
         op.setLowerBound(lbFold->getResult(0));
         op.setUpperBound(ubFold->getResult(0));
 
-      } else if (isa<arith::MulIOp>(user)) {
+      } else if (auto mulOp = dyn_cast<arith::MulIOp>(user)) {
+        // Only fold if the multiplier is a known strictly positive constant.
+        // Multiplying by zero or a negative value would produce an invalid
+        // step (scf.for requires a strictly positive step).
+        Value multiplier =
+            (mulOp.getLhs() == indVar) ? mulOp.getRhs() : mulOp.getLhs();
+        std::optional<int64_t> multiplierVal = getConstantIntValue(multiplier);
+        if (!multiplierVal || *multiplierVal <= 0)
+          break;
+
         Operation *lbFold = b.clone(*user, lbMap);
         Operation *ubFold = b.clone(*user, ubMap);
         Operation *stepFold = b.clone(*user, stepMap);
diff --git a/mlir/test/Dialect/SCF/loop-range.mlir b/mlir/test/Dialect/SCF/loop-range.mlir
index e50db8ae9f35b..79e74f430618a 100644
--- a/mlir/test/Dialect/SCF/loop-range.mlir
+++ b/mlir/test/Dialect/SCF/loop-range.mlir
@@ -131,3 +131,90 @@ func.func @fold_only_first_add(%arg0: memref<?xi32>, %arg1: index, %arg2: index)
 // CHECK:         %[[I4:.*]] = memref.load %[[ARG0]]{{\[}}%[[I3]]
 // CHECK:         %[[I5:.*]] = arith.muli %[[I4]], %[[I4]] : i32
 // CHECK:         memref.store %[[I5]], %[[ARG0]]{{\[}}%[[I3]]
+
+// Do not fold arith.muli with a negative multiplier: the resulting loop step
+// would be negative, which is invalid for scf.for.
+func.func @no_fold_negative_mul() {
+  %cm1 = arith.constant -1 : index
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c10 step %c1 {
+    %0 = arith.muli %i, %cm1 : index
+    "test.sink"(%0) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_negative_mul
+// CHECK:       %[[CM1:.*]] = arith.constant -1 : index
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C10:.*]] = arith.constant 10 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.muli %[[I]], %[[CM1]] : index
+// CHECK:         "test.sink"(%[[R]])
+
+// Do not fold arith.muli with a negative multiplier when the induction variable
+// is on the RHS.
+func.func @no_fold_negative_mul_indvar_rhs() {
+  %cm1 = arith.constant -1 : index
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c10 step %c1 {
+    %0 = arith.muli %cm1, %i : index
+    "test.sink"(%0) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_negative_mul_indvar_rhs
+// CHECK:       %[[CM1:.*]] = arith.constant -1 : index
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C10:.*]] = arith.constant 10 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.muli %[[CM1]], %[[I]] : index
+// CHECK:         "test.sink"(%[[R]])
+
+// Do not fold arith.muli with a zero multiplier.
+func.func @no_fold_zero_mul() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c10 step %c1 {
+    %0 = arith.muli %i, %c0 : index
+    "test.sink"(%0) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_zero_mul
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C10:.*]] = arith.constant 10 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.muli %[[I]], %[[C0]] : index
+// CHECK:         "test.sink"(%[[R]])
+
+// Do not fold arith.muli when the multiplier is a non-constant loop-invariant
+// value, since its sign is unknown at compile time.
+func.func @no_fold_runtime_mul(%multiplier: index) {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c10 step %c1 {
+    %0 = arith.muli %i, %multiplier : index
+    "test.sink"(%0) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_runtime_mul
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C10:.*]] = arith.constant 10 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.muli %[[I]], {{.*}} : index
+// CHECK:         "test.sink"(%[[R]])