From 9995074b526339a57ff7b9f034f216f8223db767 Mon Sep 17 00:00:00 2001
From: Amir Bishara <amir.bishara@mobileye.com>
Date: Sun, 9 Feb 2025 16:10:19 +0200
Subject: [PATCH] [mlir][scf]: Add value bound for the computed upper bound of
 for loop

Add additional bound for the induction variable of the `scf.for`
such that:
`%iv <= %lower_bound + (%trip_count - 1) * step`
---
 .../SCF/IR/ValueBoundsOpInterfaceImpl.cpp     | 27 +++++++++---
 .../SCF/value-bounds-op-interface-impl.mlir   | 41 +++++++++++++++++++
 2 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 8a27bf186d1c2..410a6bffd345e 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,18 @@ 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();
+      // iv <= lb + ((ub-lb)/step - 1) * step
+      // This bound does not replace the `iv < ub` constraint mentioned above,
+      // since constraints involving the multiplication of two constraint set
+      // dimensions are not supported.
+      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 b48f38f592dc9..339d97df001c5 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,47 @@ func.func @compare_scf_for(%a: index, %b: index, %c: index) {
 
 // -----
 
+func.func @scf_for_induction_var_upper_bound() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : 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
+}
+
+// -----
+
+#map_ceildiv_dynamic_divisor = affine_map<(i)[s] -> (i ceildiv s)>
+func.func @scf_for_induction_var_computed_upper_bound(%upperBound: index, %step: index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %tripCount = affine.apply #map_ceildiv_dynamic_divisor (%upperBound)[%step]
+  %tripCountMinusOne = arith.subi %tripCount, %c1 : index
+  %computedUpperBound = arith.muli %tripCountMinusOne, %step : index
+  scf.for %iv = %c0 to %upperBound step %step {
+    // TODO: Value bounds analysis will fail to compute upper bound
+    // because multiplication/division of unknown block arguments is
+    // not supported.
+    // expected-error @below{{unknown}}
+    "test.compare"(%iv, %computedUpperBound) {cmp = "LE"} : (index, index) -> ()
+  }
+  return
+}
+
+// -----
+
 func.func @scf_for_result_infer() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index