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

[Groverkss]

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

Same as #126426 but for scf.forall loop

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: Kunwar Grover (Groverkss)

Changes

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

Same as #126426 but for scf.forall loop

---

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

2 Files Affected:

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

diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 410a6bffd345e..9f87e27861564 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -130,6 +130,17 @@ struct ForallOpInterface
     : public ValueBoundsOpInterface::ExternalModel<ForallOpInterface,
                                                    ForallOp> {
 
+  static AffineExpr getTripCountExpr(scf::ForallOp forAllOp,
+                                     ValueBoundsConstraintSet &cstr,
+                                     int64_t idx) {
+    AffineExpr lbExpr = cstr.getExpr(forAllOp.getMixedLowerBound()[idx]);
+    AffineExpr ubExpr = cstr.getExpr(forAllOp.getMixedUpperBound()[idx]);
+    AffineExpr stepExpr = cstr.getExpr(forAllOp.getMixedStep()[idx]);
+    AffineExpr tripCountExpr =
+        AffineExpr(ubExpr - lbExpr).ceilDiv(stepExpr); // (ub - lb) / step
+    return tripCountExpr;
+  }
+
   void populateBoundsForIndexValue(Operation *op, Value value,
                                    ValueBoundsConstraintSet &cstr) const {
     auto forallOp = cast<ForallOp>(op);
@@ -141,11 +152,20 @@ struct ForallOpInterface
     assert(blockArg.getArgNumber() < forallOp.getInductionVars().size() &&
            "expected index value to be an induction var");
     int64_t idx = blockArg.getArgNumber();
-    // TODO: Take into account step size.
     AffineExpr lb = cstr.getExpr(forallOp.getMixedLowerBound()[idx]);
     AffineExpr ub = cstr.getExpr(forallOp.getMixedUpperBound()[idx]);
     cstr.bound(value) >= lb;
     cstr.bound(value) < ub;
+    // 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(forallOp, cstr, idx) - cstr.getExpr(1);
+    AffineExpr computedUpperBound =
+        lb + AffineExpr(tripCountMinusOne *
+                        cstr.getExpr(forallOp.getMixedStep()[idx]));
+    cstr.bound(value) <= computedUpperBound;
   }
 
   void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
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 339d97df001c5..60fe96d52d20b 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -379,3 +379,12 @@ func.func @scf_for_result_infer_dynamic_init_big_step(%i : index) {
   "test.compare"(%0, %7) {cmp = "LE"} : (index, index) -> ()
   return
 }
+
+func.func @scf_forall_computed_upper_bound(%x: index) {
+  %c6 = arith.constant 6 : index
+  scf.forall (%iv) = (0) to (8) step (3) {
+    // expected-remark @below{{true}}
+    "test.compare"(%iv, %c6) {cmp = "LE"} : (index, index) -> ()
+  }
+  return
+}

[AviadCo]

LGTM, Please see small comments