[Mlir-commits] [mlir] [mlir][affine] add remove-single-iteration-loop pass. (PR #129270)

[lonely eagle]

https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/129270

>From 9b99abb5a8bbabb2aa23a8373b036a4f664b0093 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 20 May 2026 13:23:38 +0000
Subject: [PATCH 1/2] imporve affine promoteIfSingleIteration.

---
 mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 66 ++++++++++++++++++++-
 1 file changed, 65 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 8f1249e3afaf0..30f1f710b901a 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -22,10 +22,12 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/OperationSupport.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/DebugLog.h"
+#include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
@@ -114,12 +116,74 @@ static void replaceIterArgsAndYieldResults(AffineForOp forOp) {
     std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
 }
 
+/// Return true if we can prove that the we always run at least the first
+/// iteration of the ForOp.
+static bool alwaysRunsFirstIteration(AffineForOp op) {
+  // Can't perform the analysis if the loops's bounds aren't index-typed.
+  if (!op.getInductionVar().getType().isIndex())
+    return false;
+  if (op.getLowerBoundMap().getNumResults() != 1 ||
+      op.getUpperBoundMap().getNumResults() != 1)
+    return false;
+
+  SmallVector<Value> lowerMapOperands = op.getLowerBoundOperands();
+  SmallVector<Value> upperMapOperands = op.getUpperBoundOperands();
+  ValueBoundsConstraintSet::Variable lower(op.getLowerBoundMap(),
+                                           lowerMapOperands);
+  ValueBoundsConstraintSet::Variable upper(op.getUpperBoundMap(),
+                                           upperMapOperands);
+  FailureOr<bool> isLb = ValueBoundsConstraintSet::compare(
+      lower, ValueBoundsConstraintSet::LT, upper);
+  return isLb.value_or(false);
+}
+
+/// Return true if we can prove that the we never run more than one iteration of
+/// the ForOp.
+static bool neverRunsSecondIteration(AffineForOp op) {
+  // Can't perform the analysis if the loops's bounds aren't index-typed.
+  if (!op.getInductionVar().getType().isIndex())
+    return false;
+
+  if (op.getLowerBoundMap().getNumResults() != 1 ||
+      op.getUpperBoundMap().getNumResults() != 1)
+    return false;
+
+  // The loop will only loop once if the inducation variable for the next time
+  // in the loop is greater than or equal to upper.
+  MLIRContext *context = op.getContext();
+  SmallVector<Value> lowerMapOperands = op.getLowerBoundOperands();
+  SmallVector<Value> upperMapOperands = op.getUpperBoundOperands();
+  SmallVector<AffineExpr> results;
+  AffineMap lowerMap = op.getLowerBoundMap();
+  for (AffineExpr expr : lowerMap.getResults())
+    results.push_back(expr + op.getStep().getSExtValue());
+
+  AffineMap nextItMap = AffineMap::get(
+      lowerMap.getNumDims(), lowerMap.getNumSymbols(), results, context);
+  ValueBoundsConstraintSet::Variable nextItVar(nextItMap, lowerMapOperands);
+  ValueBoundsConstraintSet::Variable upperVar(op.getUpperBoundMap(),
+                                              upperMapOperands);
+  FailureOr<bool> isUpperUnderNextIter = ValueBoundsConstraintSet::compare(
+      nextItVar, ValueBoundsConstraintSet::GE, upperVar);
+  return isUpperUnderNextIter.value_or(false);
+}
+
 /// Promotes the loop body of a forOp to its containing block if the forOp
 /// was known to have a single iteration.
 LogicalResult mlir::affine::promoteIfSingleIteration(AffineForOp forOp) {
   std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
-  if (!tripCount || *tripCount != 1)
+
+  // Only allow loops that are guaranteed to execute exactly once. If the trip
+  // count is constant, it must be exactly. If the trip count is dynamic, verify
+  // via affine analysis that it always runs the first iteration but never
+  // reaches the second.
+  if (tripCount && *tripCount != 1) {
     return failure();
+  }
+  if (!tripCount &&
+      !(alwaysRunsFirstIteration(forOp) && neverRunsSecondIteration(forOp))) {
+    return failure();
+  }
 
   // TODO: extend this for arbitrary affine bounds.
   if (forOp.getLowerBoundMap().getNumResults() != 1)

>From 9f29a65196ed4fe5731d0632549ebb84439cfedd Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 20 May 2026 13:41:42 +0000
Subject: [PATCH 2/2] add test.

---
 .../Dialect/Affine/affine-loop-normalize.mlir  | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 7d90efec0c21b..8dcd206237f3c 100644
--- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
@@ -323,3 +323,21 @@ func.func @multi_level_tiled_matmul() {
   }
   return
 }
+
+// -----
+
+// PROMOTE-SINGLE-ITER-LABEL: func @bound_value_promote_single_iter
+
+func.func @bound_value_promote_single_iter() -> index {
+  %c0 = arith.constant 0 :index
+  %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+  %res = affine.for %iv = %bound to 2 step 2 iter_args(%arg = %c0) -> index {
+      %sum = arith.addi %arg, %c0 : index
+      affine.yield %sum : index
+  }
+  return %res : index
+}
+// PROMOTE-SINGLE-ITER-NEXT: %[[C0:.*]] = arith.constant 0 : index
+// PROMOTE-SINGLE-ITER-NEXT: %[[VALUE_WITH_BOUNDS_0:.*]] = test.value_with_bounds
+// PROMOTE-SINGLE-ITER-NEXT: %[[ADD:.*]] = arith.addi %[[C0]], %[[C0]] : index
+// PROMOTE-SINGLE-ITER-NEXT: return %[[ADD]] : index