[Mlir-commits] [mlir] [mlir] Add shouldPromoteIfSingleIteration option to loopUnrollByFactor (PR #215080)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 13 04:08:44 PDT 2026
https://github.com/davidlerner96 updated https://github.com/llvm/llvm-project/pull/215080
>From a738108e4c3a74775153c7fec5ec7c563151a523 Mon Sep 17 00:00:00 2001
From: David Lerner <davidlerner96 at gmail.com>
Date: Sun, 9 Aug 2026 10:46:20 +0300
Subject: [PATCH] [mlir] Add shouldPromoteIfSingleIteration option to
loopUnrollByFactor
Add a shouldPromoteIfSingleIteration parameter to loopUnrollByFactor to control whether single-iteration loops are promoted during unrolling. When set to false, the function skips calls to promoteIfSingleIteration on the main loop, epilogue loop, and the unroll-factor-1 early-exit path.
The parameter defaults to true to preserve existing behavior.
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 5 ++-
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 14 ++++---
mlir/test/Dialect/SCF/loop-unroll.mlir | 40 +++++++++++++++++++
.../lib/Dialect/SCF/TestLoopUnrolling.cpp | 11 ++++-
4 files changed, 62 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index a758032ef69b4..96e99132eda35 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -123,9 +123,12 @@ struct UnrolledLoopInfo {
/// due to invalid unroll factors. Requires positive loop bounds and step. If
/// specified, annotates the Ops in each unrolled iteration by applying
/// `annotateFn`.
+/// If `shouldPromoteIfSingleIteration` is true, the function will promote the
+/// loop body up if this has turned into a single iteration loop.
FailureOr<UnrolledLoopInfo> loopUnrollByFactor(
scf::ForOp forOp, uint64_t unrollFactor,
- function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
+ function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr,
+ bool shouldPromoteIfSingleIteration = true);
/// Unrolls this loop completely.
LogicalResult loopUnrollFull(scf::ForOp forOp);
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c789b4c8904d3..c158e624002bd 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -367,7 +367,8 @@ void mlir::generateUnrolledLoop(
/// epilogue loop, if the loop is unrolled.
FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
scf::ForOp forOp, uint64_t unrollFactor,
- function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
+ function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
+ bool shouldPromoteIfSingleIteration) {
assert(unrollFactor > 0 && "expected positive unroll factor");
// Return if the loop body is empty.
@@ -408,7 +409,7 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
int64_t ubCst = getLoopBound(forOp.getUpperBound());
int64_t stepCst = getLoopBound(step);
if (unrollFactor == 1) {
- if (constTripCount->isOne() &&
+ if (shouldPromoteIfSingleIteration && constTripCount->isOne() &&
failed(forOp.promoteIfSingleIteration(rewriter)))
return failure();
return UnrolledLoopInfo{forOp, std::nullopt};
@@ -487,7 +488,8 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
}
epilogueForOp->setOperands(epilogueForOp.getNumControlOperands(),
epilogueForOp.getInitArgs().size(), results);
- if (epilogueForOp.promoteIfSingleIteration(rewriter).failed())
+ if (!shouldPromoteIfSingleIteration ||
+ epilogueForOp.promoteIfSingleIteration(rewriter).failed())
resultLoops.epilogueLoopOp = epilogueForOp;
}
@@ -509,8 +511,10 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
return arith::AddIOp::create(b, loc, iv, stride);
},
annotateFn, iterArgs, yieldedValues);
- // Promote the loop body up if this has turned into a single iteration loop.
- if (forOp.promoteIfSingleIteration(rewriter).failed())
+ // Promote the loop body up if this has turned into a single iteration loop
+ // and `shouldPromoteIfSingleIteration` is true.
+ if (!shouldPromoteIfSingleIteration ||
+ forOp.promoteIfSingleIteration(rewriter).failed())
resultLoops.mainLoopOp = forOp;
return resultLoops;
}
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index f764013ed50f9..a047d4661eb95 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -5,6 +5,8 @@
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=6 unroll-up-to-factor=true}))" | FileCheck %s --check-prefix UNROLL-UP-TO
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5
+// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3 promote-single-iteration=false' | FileCheck %s --check-prefix NO-PROMOTE-BY-3
+// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3 promote-single-iteration=true' | FileCheck %s --check-prefix PROMOTE-BY-3
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll))" --split-input-file | FileCheck %s
func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
@@ -660,3 +662,41 @@ func.func @unroll_unsigned_i2_step2_bug2() -> (i32, i32) {
// UNROLL-BY-2: arith.addi
// UNROLL-BY-2: arith.muli
// UNROLL-BY-2: return
+
+// -----
+
+// Test unrolling by 3 on a 10-iteration loop, which leaves a single-iteration
+// epilogue. Exercises the promote-single-iteration option: when false the
+// epilogue stays an scf.for; when true it is promoted to a memref.store.
+func.func @static_loop_unroll_by_3_no_promote_epilogue(%arg0 : memref<?xf32>) {
+ %0 = arith.constant 7.0 : f32
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 10 : index
+ %step = arith.constant 1 : index
+ scf.for %i0 = %lb to %ub step %step {
+ memref.store %0, %arg0[%i0] : memref<?xf32>
+ }
+ return
+}
+// Promotion disabled keeps the epilogue as an scf.for.
+//
+// NO-PROMOTE-BY-3-LABEL: func @static_loop_unroll_by_3_no_promote_epilogue
+// NO-PROMOTE-BY-3: scf.for %{{.*}} = %c0 to %c9 step %c3
+// NO-PROMOTE-BY-3: memref.store
+// NO-PROMOTE-BY-3: memref.store
+// NO-PROMOTE-BY-3: memref.store
+// NO-PROMOTE-BY-3: scf.for %{{.*}} = %c9 to %c10 step %c1
+// NO-PROMOTE-BY-3: memref.store
+
+// With promotion enabled, the single-iteration epilogue is promoted out of the
+// loop.
+//
+// PROMOTE-BY-3-LABEL: func @static_loop_unroll_by_3_no_promote_epilogue
+// PROMOTE-BY-3: scf.for %{{.*}} = %c0 to %c9 step %c3
+// PROMOTE-BY-3: memref.store
+// PROMOTE-BY-3: memref.store
+// PROMOTE-BY-3: memref.store
+// PROMOTE-BY-3-NOT: scf.for
+// PROMOTE-BY-3: memref.store
+
+
diff --git a/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp b/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
index e2d6996e435bb..bbeae9d39db8d 100644
--- a/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
@@ -42,10 +42,12 @@ struct TestLoopUnrollingPass
TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
unsigned loopDepthParam,
- bool annotateLoopParam) {
+ bool annotateLoopParam,
+ bool promoteSingleIterationParam) {
unrollFactor = unrollFactorParam;
loopDepth = loopDepthParam;
annotateLoop = annotateLoopParam;
+ promoteSingleIteration = promoteSingleIterationParam;
}
void getDependentDialects(DialectRegistry ®istry) const override {
@@ -73,7 +75,8 @@ struct TestLoopUnrollingPass
if (unrollFactor.getValue() == -1)
(void)loopUnrollFull(loop);
else
- (void)loopUnrollByFactor(loop, unrollFactor, annotateFn);
+ (void)loopUnrollByFactor(loop, unrollFactor, annotateFn,
+ promoteSingleIteration);
}
}
Option<int64_t> unrollFactor{
@@ -89,6 +92,10 @@ struct TestLoopUnrollingPass
llvm::cl::init(false)};
Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
llvm::cl::init(0)};
+ Option<bool> promoteSingleIteration{
+ *this, "promote-single-iteration",
+ llvm::cl::desc("Promote single-iteration loops after unrolling."),
+ llvm::cl::init(true)};
};
} // namespace
More information about the Mlir-commits
mailing list