[Mlir-commits] [mlir] fc06190 - [mlir][scf] Revisit affected loops during range folding (#230961)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 20:08:59 PDT 2026
Author: 陈子昂
Date: 2026-10-30T11:08:55+08:00
New Revision: fc061907e0a379003f6d6120763f53020936a19e
URL: https://github.com/llvm/llvm-project/commit/fc061907e0a379003f6d6120763f53020936a19e
DIFF: https://github.com/llvm/llvm-project/commit/fc061907e0a379003f6d6120763f53020936a19e.diff
LOG: [mlir][scf] Revisit affected loops during range folding (#230961)
Folding an outer loop can expose a range-folding opportunity in an inner
loop that has already been visited. Convert loop-range folding to an
`OpRewritePattern<scf::ForOp>` and use `applyPatternsGreedily` on the
pass root so the greedy driver revisits affected loops within one pass
invocation.
Fixes #228463
Added:
Modified:
mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
mlir/test/Dialect/SCF/loop-range.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp
index 0d7d310cfaafd0..e097bb3cec8ddc 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 "mlir/Transforms/GreedyPatternRewriteDriver.h"
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
@@ -28,86 +29,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 Base::Base;
-void ForLoopRangeFolding::runOnOperation() {
- getOperation()->walk([&](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);
- 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));
- }
+ // 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();
+ }
- ValueRange wrapIndvar(indVar);
- user->replaceAllUsesWith(wrapIndvar);
- user->erase();
+ 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() {
+ RewritePatternSet patterns(&getContext());
+ patterns.add<ForLoopRangeFoldingPattern>(&getContext());
+ GreedyRewriteConfig config;
+ config.setStrictness(GreedyRewriteStrictness::ExistingOps)
+ .enableFolding(false)
+ .enableConstantCSE(false)
+ .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
+ if (failed(
+ applyPatternsGreedily(getOperation(), std::move(patterns), config)))
+ signalPassFailure();
}
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 3cf4c857e3f8ce..b97d8a805303c8 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]])
More information about the Mlir-commits
mailing list