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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Prashanth (StarOne01)

<details>
<summary>Changes</summary>

fixes #<!-- -->231062 which affects loops with constant bounds where the `step` and `upperbound - lowerbound` are very close and close enogh to exceeds step by less than about half a float32 ulp at that magnitude.

eg: 

`ub - lb`: 20000001, `step`: 20000000
`ub - lb`: 1099511693312, `step`:  1099511627776

removing the float conversation and hence the ceil operation and comparing `*ubInt - *lbInt <= *stepInt` with repro as a test.


---
Full diff: https://github.com/llvm/llvm-project/pull/231100.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp (+1-1) 
- (modified) mlir/test/Dialect/SCF/for-loop-peeling-front.mlir (+33) 


``````````diff
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

``````````

</details>


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