[Mlir-commits] [mlir] [mlir][scf] Revisit affected loops during range folding (PR #230961)

[llvmlistbot at llvm.org]

https://github.com/Michael-Chen-NJU created https://github.com/llvm/llvm-project/pull/230961

Folding an outer loop can expose a range-folding opportunity in an inner loop that has already been visited. Revisit affected loops with a worklist so one pass can fold both.

Fixes #228463


>From 67437526ecdc3e26a79cf5895af8db7f95efd3db Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <2802328816 at qq.com>
Date: Sun, 11 Oct 2026 21:21:10 +0800
Subject: [PATCH] [mlir][scf] Revisit affected loops during range folding

---
 .../SCF/Transforms/LoopRangeFolding.cpp       | 38 ++++++++++++++-
 mlir/test/Dialect/SCF/loop-range.mlir         | 48 +++++++++++++++++++
 2 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index 0d7d310cfaafd..fc1f0aad798ef 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/IRMapping.h"
+#include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
@@ -35,7 +36,21 @@ struct ForLoopRangeFolding
 } // namespace
 
 void ForLoopRangeFolding::runOnOperation() {
-  getOperation()->walk([&](ForOp op) {
+  llvm::SmallSetVector<ForOp, 16> worklist;
+  // Popping from the back processes nested loops before their parents.
+  getOperation()->walk<WalkOrder::PreOrder>(
+      [&](ForOp op) { worklist.insert(op); });
+
+  auto enqueueInductionVarOwner = [&](Value value, ForOp current) {
+    auto arg = dyn_cast<BlockArgument>(value);
+    if (!arg)
+      return;
+    auto loop = dyn_cast<ForOp>(arg.getOwner()->getParentOp());
+    if (loop && loop != current && loop.getInductionVar() == value)
+      worklist.insert(loop);
+  };
+
+  auto foldLoop = [&](ForOp op) {
     Value indVar = op.getInductionVar();
 
     auto canBeFolded = [&](Value value) {
@@ -64,6 +79,9 @@ void ForLoopRangeFolding::runOnOperation() {
         break;
 
       OpBuilder b(op);
+      Value oldLowerBound = op.getLowerBound();
+      Value oldUpperBound = op.getUpperBound();
+      Value oldStep = op.getStep();
       IRMapping lbMap;
       lbMap.map(indVar, op.getLowerBound());
       IRMapping ubMap;
@@ -103,11 +121,27 @@ void ForLoopRangeFolding::runOnOperation() {
         op.setStep(stepFold->getResult(0));
       }
 
+      // Revisit loops whose induction variables may have lost or gained uses.
+      for (Value value : {oldLowerBound, oldUpperBound, oldStep})
+        enqueueInductionVarOwner(value, op);
+      for (Value value : user->getOperands())
+        enqueueInductionVarOwner(value, op);
+
+      // Replacing this arithmetic result may make a nested loop foldable.
+      for (Operation *consumer : user->getUsers()) {
+        if (!isa<arith::AddIOp, arith::MulIOp>(consumer))
+          continue;
+        for (Value value : consumer->getOperands())
+          enqueueInductionVarOwner(value, op);
+      }
+
       ValueRange wrapIndvar(indVar);
       user->replaceAllUsesWith(wrapIndvar);
       user->erase();
     }
-  });
+  };
+  while (!worklist.empty())
+    foldLoop(worklist.pop_back_val());
 }
 
 std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
diff --git a/mlir/test/Dialect/SCF/loop-range.mlir b/mlir/test/Dialect/SCF/loop-range.mlir
index 3cf4c857e3f8c..b97d8a805303c 100644
--- a/mlir/test/Dialect/SCF/loop-range.mlir
+++ b/mlir/test/Dialect/SCF/loop-range.mlir
@@ -291,3 +291,51 @@ func.func @no_fold_self_mul() {
 // CHECK:       scf.for %[[I:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
 // CHECK:         %[[R:.*]] = arith.muli %[[I]], %[[I]] : index
 // CHECK:         "test.sink"(%[[R]])
+
+// Folding the outer loop exposes the inner loop's addition for folding in the
+// same pass invocation.
+func.func @fold_nested_loops(%out: memref<10xi32>, %v: i32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  scf.for %i = %c0 to %c4 step %c1 {
+    scf.for %j = %c0 to %c4 step %c1 {
+      %scaled = arith.muli %i, %c2 : index
+      %idx = arith.addi %j, %scaled : index
+      memref.store %v, %out[%idx] : memref<10xi32>
+    }
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @fold_nested_loops
+// CHECK: scf.for %[[OUTER:.*]] =
+// CHECK: scf.for %[[INNER:.*]] =
+// CHECK-NOT: arith.addi
+// CHECK: memref.store %{{.*}}, %{{.*}}{{\[}}%[[INNER]]]
+
+// Each fold exposes a folding opportunity in the next nested loop.
+func.func @fold_three_nested_loops() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  scf.for %i = %c0 to %c4 step %c1 {
+    scf.for %j = %c0 to %c4 step %c1 {
+      scf.for %k = %c0 to %c4 step %c1 {
+        %scaled = arith.muli %i, %c2 : index
+        %middle = arith.addi %j, %scaled : index
+        %inner = arith.addi %k, %middle : index
+        "test.sink"(%inner) : (index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @fold_three_nested_loops
+// CHECK: scf.for %[[OUTER:.*]] =
+// CHECK: scf.for %[[MIDDLE:.*]] =
+// CHECK: scf.for %[[INNER:.*]] =
+// CHECK-NEXT: "test.sink"(%[[INNER]])