[Mlir-commits] [mlir] [mlir][scf] Fold Self-addition of the induction variable into the loop range (PR #228725)

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/228725

>From 79e052e7cfb1dc824be84db011fb5b3126917b8d Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 3 Oct 2026 21:04:23 +0530
Subject: [PATCH] [mlir][scf] Fold Self-addition of the induction variable into
 the loop range

scf-for-loop-range-folding only folded an arith op into the loop range if the induction variable had a single use (hasOneUse()). That counts operand slots rather than users, so `arith.addi %i, %i`, a single
operation that uses the induction variable twice, was rejected.

now we check whether the induction variable is used in a single operation (user).

Fixes #228474

Assisted-by: Claude
---
 .../SCF/Transforms/LoopRangeFolding.cpp       | 20 ++++-
 mlir/test/Dialect/SCF/loop-range.mlir         | 73 +++++++++++++++++++
 2 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index d4b4c4cfc063f..0d7d310cfaafd 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
@@ -45,12 +45,18 @@ void ForLoopRangeFolding::runOnOperation() {
     // Fold until a fixed point is reached
     while (true) {
 
-      // If the induction variable is used more than once, we can't fold its
-      // arith ops into the loop range
-      if (!indVar.hasOneUse())
+      // If the induction variable is used by more than one operation, we can't
+      // fold its arith ops into the loop range. Note that a single operation
+      // may use the induction variable several times (e.g. `arith.addi %i,
+      // %i`), so check for a single user rather than a single use.
+      if (indVar.use_empty())
         break;
 
       Operation *user = *indVar.getUsers().begin();
+      if (!llvm::all_of(indVar.getUsers(),
+                        [&](Operation *u) { return u == user; }))
+        break;
+
       if (!isa<arith::AddIOp, arith::MulIOp>(user))
         break;
 
@@ -65,13 +71,19 @@ void ForLoopRangeFolding::runOnOperation() {
       IRMapping stepMap;
       stepMap.map(indVar, op.getStep());
 
-      if (isa<arith::AddIOp>(user)) {
+      if (auto addOp = dyn_cast<arith::AddIOp>(user)) {
         Operation *lbFold = b.clone(*user, lbMap);
         Operation *ubFold = b.clone(*user, ubMap);
 
         op.setLowerBound(lbFold->getResult(0));
         op.setUpperBound(ubFold->getResult(0));
 
+        // `arith.addi %i, %i` is `2 * %i`, so the step has to be doubled too.
+        if (addOp.getLhs() == indVar && addOp.getRhs() == indVar) {
+          Operation *stepFold = b.clone(*user, stepMap);
+          op.setStep(stepFold->getResult(0));
+        }
+
       } 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
diff --git a/mlir/test/Dialect/SCF/loop-range.mlir b/mlir/test/Dialect/SCF/loop-range.mlir
index 79e74f430618a..3cf4c857e3f8c 100644
--- a/mlir/test/Dialect/SCF/loop-range.mlir
+++ b/mlir/test/Dialect/SCF/loop-range.mlir
@@ -218,3 +218,76 @@ func.func @no_fold_runtime_mul(%multiplier: index) {
 // CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C10]] step %[[C1]] {
 // CHECK:         %[[R:.*]] = arith.muli %[[I]], {{.*}} : index
 // CHECK:         "test.sink"(%[[R]])
+
+// `arith.addi %i, %i` is `2 * %i`, so it can be folded into the loop range like
+// a multiplication by 2: the lower bound, the upper bound and the step are all
+// doubled. The step has to be doubled too, otherwise the folded loop would visit
+// every index instead of every second one.
+func.func @fold_self_add(%arg0: memref<16xi32>, %arg1: i32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  scf.for %i = %c0 to %c8 step %c1 {
+    %0 = arith.addi %i, %i : index
+    memref.store %arg1, %arg0[%0] : memref<16xi32>
+  }
+  return
+}
+
+// CHECK-LABEL: func @fold_self_add
+// CHECK-SAME:  (%[[ARG0:[a-z0-9]+]]: {{.*}}, %[[ARG1:[a-z0-9]+]]: {{.*}})
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       %[[C8:.*]] = arith.constant 8 : index
+// CHECK:       %[[LB:.*]] = arith.addi %[[C0]], %[[C0]] : index
+// CHECK:       %[[UB:.*]] = arith.addi %[[C8]], %[[C8]] : index
+// CHECK:       %[[STEP:.*]] = arith.addi %[[C1]], %[[C1]] : index
+// CHECK:       scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+// CHECK-NEXT:    memref.store %[[ARG1]], %[[ARG0]]{{\[}}%[[I]]
+
+// The induction variable has a user other than the `arith.addi`. Folding would
+// change what that user sees (`2 * %i` instead of `%i`), so the loop must be
+// left unchanged.
+func.func @no_fold_self_add_extra_user() {
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c8 step %c1 {
+    %0 = arith.addi %i, %i : index
+    "test.sink"(%0) : (index) -> ()
+    "test.sink"(%i) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_self_add_extra_user
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C8:.*]] = arith.constant 8 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.addi %[[I]], %[[I]] : index
+// CHECK:         "test.sink"(%[[R]])
+// CHECK:         "test.sink"(%[[I]])
+
+// `arith.muli %i, %i` is not a scaling of the induction variable by a constant,
+// so it cannot be expressed as a new loop range. The pass only folds a
+// multiplication by a known positive constant, so the loop must be left
+// unchanged.
+func.func @no_fold_self_mul() {
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c8 step %c1 {
+    %0 = arith.muli %i, %i : index
+    "test.sink"(%0) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @no_fold_self_mul
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       %[[C8:.*]] = arith.constant 8 : index
+// CHECK:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
+// CHECK:         %[[R:.*]] = arith.muli %[[I]], %[[I]] : index
+// CHECK:         "test.sink"(%[[R]])