[Mlir-commits] [mlir] [mlir][scf] Add value bound for computed upper bound of forall loop (PR #171158)

[Kunwar Grover]

https://github.com/Groverkss created https://github.com/llvm/llvm-project/pull/171158

Same as https://github.com/llvm/llvm-project/pull/126426 but for scf.forall loop

>From cc45adb4020dc4feaf8df88158bda959415e978d Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Mon, 8 Dec 2025 17:03:22 +0000
Subject: [PATCH] [mlir][scf] Add value bound for computed upper bound of
 forall loop

---
 .../SCF/IR/ValueBoundsOpInterfaceImpl.cpp     | 22 ++++++++++++++++++-
 .../SCF/value-bounds-op-interface-impl.mlir   |  9 ++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

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
+}