[mlir][scf]: Add value bound for the computed upper bound of for loop

[amirBish]

Add additional bound for the induction variable of the scf.for such that:
%iv <= %lower_bound + (%trip_count - 1) * step

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: Amir Bishara (amirBish)

Changes

Add additional bound for the induction variable of the scf.for such that:
%iv <= %lower_bound + (%trip_count - 1) * step

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp (+17-6)
• (modified) mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir (+21)

diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 8a27bf186d1c2a4..a6f0ec74ad52fdb 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -20,6 +20,16 @@ namespace {
 struct ForOpInterface
     : public ValueBoundsOpInterface::ExternalModel<ForOpInterface, ForOp> {
 
+  static AffineExpr getTripCountExpr(scf::ForOp forOp,
+                                     ValueBoundsConstraintSet &cstr) {
+    AffineExpr lbExpr = cstr.getExpr(forOp.getLowerBound());
+    AffineExpr ubExpr = cstr.getExpr(forOp.getUpperBound());
+    AffineExpr stepExpr = cstr.getExpr(forOp.getStep());
+    AffineExpr tripCountExpr =
+        AffineExpr(ubExpr - lbExpr).ceilDiv(stepExpr); // (ub - lb) / step
+    return tripCountExpr;
+  }
+
   /// Populate bounds of values/dimensions for iter_args/OpResults. If the
   /// value/dimension size does not change in an iteration, we can deduce that
   /// it the same as the initial value/dimension.
@@ -77,11 +87,7 @@ struct ForOpInterface
     // `value` is result of `forOp`, we can prove that:
     // %result == %init_arg + trip_count * (%yielded_value - %iter_arg).
     // Where trip_count is (ub - lb) / step.
-    AffineExpr lbExpr = cstr.getExpr(forOp.getLowerBound());
-    AffineExpr ubExpr = cstr.getExpr(forOp.getUpperBound());
-    AffineExpr stepExpr = cstr.getExpr(forOp.getStep());
-    AffineExpr tripCountExpr =
-        AffineExpr(ubExpr - lbExpr).ceilDiv(stepExpr); // (ub - lb) / step
+    AffineExpr tripCountExpr = getTripCountExpr(forOp, cstr);
     AffineExpr oneIterAdvanceExpr =
         cstr.getExpr(yieldedValue) - cstr.getExpr(iterArg);
     cstr.bound(value) ==
@@ -93,9 +99,14 @@ struct ForOpInterface
     auto forOp = cast<ForOp>(op);
 
     if (value == forOp.getInductionVar()) {
-      // TODO: Take into account step size.
       cstr.bound(value) >= forOp.getLowerBound();
       cstr.bound(value) < forOp.getUpperBound();
+      AffineExpr tripCountMinusOne =
+          getTripCountExpr(forOp, cstr) - cstr.getExpr(1);
+      AffineExpr computedUpperBound =
+          cstr.getExpr(forOp.getLowerBound()) +
+          AffineExpr(tripCountMinusOne * cstr.getExpr(forOp.getStep()));
+      cstr.bound(value) <= computedUpperBound;
       return;
     }
 
diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
index b48f38f592dc922..53dbb4ceff4f051 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -270,6 +270,27 @@ func.func @compare_scf_for(%a: index, %b: index, %c: index) {
 
 // -----
 
+func.func @scf_for_induction_var_upper_bound(%i : index) {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %c5 = arith.constant 5 : index
+  %c8 = arith.constant 8 : index
+  %c10 = arith.constant 10 : index
+  scf.for %iv = %c0 to %c10 step %c4 {
+    // expected-remark @below{{true}}
+    "test.compare"(%iv, %c8) {cmp = "LE"} : (index, index) -> ()
+  }
+  scf.for %iv = %c2 to %c8 step %c3 {
+    // expected-remark @below{{true}}
+    "test.compare"(%iv, %c5) {cmp = "LE"} : (index, index) -> ()
+  }
+  return
+}
+
+// -----
+
 func.func @scf_for_result_infer() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index

[amirBish]

@matthias-springer Thanks for the review, feel free to check it again.

[matthias-springer]

Looks good!