[Mlir-commits] [mlir] [mlir] Add shouldPromoteIfSingleIteration option to loopUnrollByFactor (PR #215080)

[llvmlistbot at llvm.org]

https://github.com/davidlerner96 created https://github.com/llvm/llvm-project/pull/215080

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.

>From 54e96cb913f1e1c924f5b543173a5fe83faf7a51 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 +++++++++-----
 2 files changed, 13 insertions(+), 6 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..78f1ced15c97b 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;
 }