[Mlir-commits] [mlir] [mlir][scf] Revisit affected loops during range folding (PR #230961)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 12 07:54:06 PDT 2026
https://github.com/Michael-Chen-NJU updated https://github.com/llvm/llvm-project/pull/230961
>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 1/3] [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]])
>From 2ebd4c3ec42b6b692469399ac771d6696d466abd Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <2802328816 at qq.com>
Date: Sun, 11 Oct 2026 23:48:49 +0800
Subject: [PATCH 2/3] [mlir][scf] Use a pattern for loop range folding
---
.../SCF/Transforms/LoopRangeFolding.cpp | 189 ++++++++----------
1 file changed, 83 insertions(+), 106 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index fc1f0aad798ef..a265c6796ea4d 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
@@ -18,7 +18,8 @@
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
-#include "llvm/ADT/SetVector.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
@@ -29,119 +30,95 @@ using namespace mlir;
using namespace mlir::scf;
namespace {
-struct ForLoopRangeFolding
- : public impl::SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
- void runOnOperation() override;
-};
-} // namespace
+struct ForLoopRangeFoldingPattern : public OpRewritePattern<ForOp> {
+ using OpRewritePattern<ForOp>::OpRewritePattern;
-void ForLoopRangeFolding::runOnOperation() {
- 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) {
+ void initialize() { setHasBoundedRewriteRecursion(); }
+
+ LogicalResult matchAndRewrite(ForOp op,
+ PatternRewriter &rewriter) const override {
Value indVar = op.getInductionVar();
auto canBeFolded = [&](Value value) {
return op.isDefinedOutsideOfLoop(value) || value == indVar;
};
- // Fold until a fixed point is reached
- while (true) {
-
- // 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;
-
- if (!llvm::all_of(user->getOperands(), canBeFolded))
- 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;
- ubMap.map(indVar, op.getUpperBound());
- IRMapping stepMap;
- stepMap.map(indVar, op.getStep());
-
- 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
- // 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);
-
- op.setLowerBound(lbFold->getResult(0));
- op.setUpperBound(ubFold->getResult(0));
- 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();
+ // 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())
+ return failure();
+
+ Operation *user = *indVar.getUsers().begin();
+ if (!llvm::all_of(indVar.getUsers(),
+ [&](Operation *u) { return u == user; }))
+ return failure();
+
+ if (!isa<arith::AddIOp, arith::MulIOp>(user))
+ return failure();
+
+ if (!llvm::all_of(user->getOperands(), canBeFolded))
+ return failure();
+
+ 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)
+ return failure();
}
- };
- while (!worklist.empty())
- foldLoop(worklist.pop_back_val());
+
+ rewriter.setInsertionPoint(op);
+ IRMapping lbMap;
+ lbMap.map(indVar, op.getLowerBound());
+ IRMapping ubMap;
+ ubMap.map(indVar, op.getUpperBound());
+ IRMapping stepMap;
+ stepMap.map(indVar, op.getStep());
+
+ Operation *lbFold = rewriter.clone(*user, lbMap);
+ Operation *ubFold = rewriter.clone(*user, ubMap);
+ Operation *stepFold = nullptr;
+ if (auto addOp = dyn_cast<arith::AddIOp>(user)) {
+ // `arith.addi %i, %i` is `2 * %i`, so the step has to be doubled too.
+ if (addOp.getLhs() == indVar && addOp.getRhs() == indVar)
+ stepFold = rewriter.clone(*user, stepMap);
+ } else {
+ stepFold = rewriter.clone(*user, stepMap);
+ }
+
+ rewriter.modifyOpInPlace(op, [&]() {
+ op.setLowerBound(lbFold->getResult(0));
+ op.setUpperBound(ubFold->getResult(0));
+ if (stepFold)
+ op.setStep(stepFold->getResult(0));
+ });
+ rewriter.replaceOp(user, indVar);
+ return success();
+ }
+};
+
+struct ForLoopRangeFolding
+ : public impl::SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
+ void runOnOperation() override;
+};
+} // namespace
+
+void ForLoopRangeFolding::runOnOperation() {
+ SmallVector<Operation *> loops;
+ getOperation()->walk<WalkOrder::PreOrder>(
+ [&](ForOp op) { loops.push_back(op); });
+ RewritePatternSet patterns(&getContext());
+ patterns.add<ForLoopRangeFoldingPattern>(&getContext());
+ GreedyRewriteConfig config;
+ config.setStrictness(GreedyRewriteStrictness::ExistingOps)
+ .enableFolding(false);
+ if (failed(applyOpPatternsGreedily(loops, std::move(patterns), config)))
+ signalPassFailure();
}
std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
>From b321d4648c83bc3f71f1e25593eef0627f27282a Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <2802328816 at qq.com>
Date: Mon, 12 Oct 2026 22:47:59 +0800
Subject: [PATCH 3/3] [mlir][scf] Run loop range folding with root-based greedy
rewriting
---
mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp | 11 +++++------
1 file changed, 5 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index a265c6796ea4d..c40923ea8d59c 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
@@ -19,7 +19,6 @@
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/ADT/SmallVector.h"
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
@@ -109,15 +108,15 @@ struct ForLoopRangeFolding
} // namespace
void ForLoopRangeFolding::runOnOperation() {
- SmallVector<Operation *> loops;
- getOperation()->walk<WalkOrder::PreOrder>(
- [&](ForOp op) { loops.push_back(op); });
RewritePatternSet patterns(&getContext());
patterns.add<ForLoopRangeFoldingPattern>(&getContext());
GreedyRewriteConfig config;
config.setStrictness(GreedyRewriteStrictness::ExistingOps)
- .enableFolding(false);
- if (failed(applyOpPatternsGreedily(loops, std::move(patterns), config)))
+ .enableFolding(false)
+ .enableConstantCSE(false)
+ .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
+ if (failed(
+ applyPatternsGreedily(getOperation(), std::move(patterns), config)))
signalPassFailure();
}
More information about the Mlir-commits
mailing list