[Mlir-commits] [mlir] [mlir][scf] Fix front peeling skipped due to float rounding in trip count (PR #231100)

[llvmlistbot at llvm.org]

https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/231100

>From 83a4ade9d59187188a242920bb51133597f94963 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Mon, 12 Oct 2026 14:27:42 +0530
Subject: [PATCH 1/4] [mlir][scf] Fix front peeling skipped due to float
 rounding in trip count

---
 .../SCF/Transforms/LoopSpecialization.cpp     |  2 +-
 .../Dialect/SCF/for-loop-peeling-front.mlir   | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 685c57319a216..d0719643603c9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -239,7 +239,7 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
   auto stepInt = getConstantIntValue(forOp.getStep());
 
   // Peeling is not needed if there is one or less iteration.
-  if (lbInt && ubInt && stepInt && ceil(float(*ubInt - *lbInt) / *stepInt) <= 1)
+  if (lbInt && ubInt && stepInt && *ubInt - *lbInt <= *stepInt)
     return failure();
 
   // The peeling bound (%lb + %step) is computed with affine.apply, which
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index ec91e4def6484..0eba57e5af771 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -242,3 +242,36 @@ func.func @non_index_loop_bounds(%init: i32, %lb: i32, %n: i32) -> i32 {
   }
   return %r : i32
 }
+
+// -----
+
+// Tests loops with constant bounds where the step and (upper bound - lower bound) are
+// very close, so `ceil(float(ub - lb) / step)` used to round to 1, and the front peeling
+// was skipped. See https://github.com/llvm/llvm-project/issues/231062
+
+// CHECK-LABEL:   func.func @peel_front_step_near_ub(
+// CHECK-SAME:                                         %[[MEMREF:[^:]*]]: memref<2xindex>) {
+// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[C16777216:.*]] = arith.constant 16777216 : index
+// CHECK-DAG:       %[[C16777217:.*]] = arith.constant 16777217 : index
+// CHECK:           scf.for %[[IV0:.*]] = %[[C0]] to %[[C16777216]] step %[[C16777216]] {
+// CHECK:             %[[SLOT0:.*]] = arith.divui %[[IV0]], %[[C16777216]] : index
+// CHECK:             memref.store %[[IV0]], %[[MEMREF]][%[[SLOT0]]] : memref<2xindex>
+// CHECK:           }
+// CHECK:           scf.for %[[IV1:.*]] = %[[C16777216]] to %[[C16777217]] step %[[C16777216]] {
+// CHECK:             %[[SLOT1:.*]] = arith.divui %[[IV1]], %[[C16777216]] : index
+// CHECK:             memref.store %[[IV1]], %[[MEMREF]][%[[SLOT1]]] : memref<2xindex>
+// CHECK:           }
+// CHECK:           return
+
+func.func @peel_front_step_near_ub(%out: memref<2xindex>) {
+  %c0 = arith.constant 0 : index
+  %step = arith.constant 16777216 : index
+  %ub = arith.constant 16777217 : index
+
+  scf.for %i = %c0 to %ub step %step {
+    %slot = arith.divui %i, %step : index
+    memref.store %i, %out[%slot] : memref<2xindex>
+  }
+  return
+}
\ No newline at end of file

>From 961d8daaa759a8c0a6886cd3a8e0d8e9602d7a75 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Mon, 12 Oct 2026 14:36:04 +0530
Subject: [PATCH 2/4] [mlir][scf] add newline at end of file in
 for-loop-peeling-front.mlir

---
 mlir/test/Dialect/SCF/for-loop-peeling-front.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index 0eba57e5af771..84b561a8571f6 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -274,4 +274,4 @@ func.func @peel_front_step_near_ub(%out: memref<2xindex>) {
     memref.store %i, %out[%slot] : memref<2xindex>
   }
   return
-}
\ No newline at end of file
+}

>From 12436849c6fa632341c6499f11e4472bba4b8b25 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Tue, 13 Oct 2026 10:52:44 +0530
Subject: [PATCH 3/4] [mlir][scf] add comments explaining how the new
 implementation figures whether there is more than one iteration

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index d0719643603c9..e46b588a2f124 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -238,7 +238,12 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
   auto ubInt = getConstantIntValue(forOp.getUpperBound());
   auto stepInt = getConstantIntValue(forOp.getStep());
 
-  // Peeling is not needed if there is one or less iteration.
+  // Peeling is not needed if there is at most one iteration. 
+  // Per the scf.for semantics, the trip count is
+  // max(0, ceil((UB - LB) / Step)), so for a positive step there is more than
+  // one iteration exactly when UB - LB > Step. 
+  // Compare the integers directly: going through float is lossy and wrongly skipped
+  // peeling of loops with two iterations.
   if (lbInt && ubInt && stepInt && *ubInt - *lbInt <= *stepInt)
     return failure();
 

>From 40d93499a6a1717cf14883e79fd0a0a053b945e3 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Tue, 13 Oct 2026 10:56:46 +0530
Subject: [PATCH 4/4] [mlir][scf] format comments explaining how the new
 implementation figures whether there is more than one iteration

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index e46b588a2f124..70e37a3868548 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -238,12 +238,12 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
   auto ubInt = getConstantIntValue(forOp.getUpperBound());
   auto stepInt = getConstantIntValue(forOp.getStep());
 
-  // Peeling is not needed if there is at most one iteration. 
+  // Peeling is not needed if there is at most one iteration.
   // Per the scf.for semantics, the trip count is
   // max(0, ceil((UB - LB) / Step)), so for a positive step there is more than
-  // one iteration exactly when UB - LB > Step. 
-  // Compare the integers directly: going through float is lossy and wrongly skipped
-  // peeling of loops with two iterations.
+  // one iteration exactly when UB - LB > Step.
+  // Compare the integers directly: going through float is lossy and wrongly
+  // skipped peeling of loops with two iterations.
   if (lbInt && ubInt && stepInt && *ubInt - *lbInt <= *stepInt)
     return failure();