[Mlir-commits] [mlir] [mlir][affine] Implement ValueBoundsOpInterface for affine.for (PR #214614)
Dhairyashil R G
llvmlistbot at llvm.org
Thu Aug 6 20:19:59 PDT 2026
https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/214614
`scf.for` has provided induction variable bounds through `ValueBoundsOpInterface` for a long time, but `affine.for` has no model at all. `ValueBoundsConstraintSet` therefore cannot derive any bound for an affine induction variable, not even `iv >= lowerBound`. Queries just return "unknown", which is easy to miss because it looks the same as a bound that genuinely cannot be proven.
Add a model for the induction variable. The lower bound of an `affine.for` is the maximum over the results of its lower bound map and the upper bound is the minimum over the results of its upper bound map, so the induction variable is constrained by every individual result. This follows how `AffineMinOpInterface` and `AffineMaxOpInterface` in the same file align their maps with the constraint set; the shared logic is factored into `alignBoundExpr`.
When both maps have exactly one result the step is taken into account as well, following `scf::ForOpInterface`: the induction variable is always a multiple of `step` away from the lower bound, so it never exceeds `lb + (tripCount - 1) * step`. That is tighter than `ub - 1` whenever the trip count is not a multiple of the step, for `affine.for %i = 0 to 300 step 128` the induction variable only ever takes {0, 128, 256}, so the bound is 256 rather than 299. As in `scf::ForOpInterface` this does not replace the `iv < ub` constraint, since constraints multiplying two constraint set dimensions aren't supported. The step is not applied when either map has multiple results (`max`/`min` bounds), since no single result can drive the arithmetic.
No bounds are inferred for `iter_args`.
Beyond the lit tests, the bound arithmetic was checked against an oracle sweep: 6268 loops over an exhaustive `(lb, ub, step)` grid plus 4000 random cases, with the true induction variable range computed by simulating each loop in Python rather than by asking the compiler. Across 12536 queries there were 0 unsound bounds and 0 loose ones — every bound is exactly `max(iv)` and `min(iv)`, so the model captures every provable case rather than a conservative subset. Two controls: reverting the change makes all 12536 queries unprovable, and mutating the step arithmetic to be one step too tight makes the sweep report 5643 unsound bounds with counterexamples. The harness discriminates in both directions.
No existing test needed updating: on its own this causes no `in_bounds` folding in the vector dialect, because `isInBounds` still bails on any non-constant index. Using these bounds there is a follow-up I'd like to send once this lands.
Assisted-by: Claude
>From 812fe9930f232f6d007a73cdbe83666eec2dea00 Mon Sep 17 00:00:00 2001
From: Dhairyashil Ghatage <dhairyashil25 at gmail.com>
Date: Fri, 7 Aug 2026 00:28:05 +0530
Subject: [PATCH] [mlir][affine] Implement ValueBoundsOpInterface for
affine.for
`scf.for` has provided induction variable bounds through
`ValueBoundsOpInterface` for a long time, but `affine.for` has no model
at all, so `ValueBoundsConstraintSet` cannot derive any bound for an
affine induction variable -- not even `iv >= lowerBound`. Queries simply
return "unknown".
Add a model for the induction variable. The lower bound of an
`affine.for` is the maximum over the results of its lower bound map and
the upper bound is the minimum over the results of its upper bound map,
so the induction variable is constrained by each individual result. This
mirrors how `AffineMinOpInterface` and `AffineMaxOpInterface` in the same
file align their maps with the constraint set.
When both maps have a single result the step is taken into account as
well, following `scf::ForOpInterface`: the induction variable is always a
multiple of `step` away from the lower bound, so it never exceeds
`lb + (tripCount - 1) * step`. That is tighter than `ub - 1` whenever the
trip count is not a multiple of the step -- for
`affine.for %i = 0 to 300 step 128` the induction variable only ever
takes {0, 128, 256}, so the upper bound is 256 rather than 299.
No bounds are inferred for `iter_args`.
Assisted-by: Claude
---
.../Affine/IR/ValueBoundsOpInterfaceImpl.cpp | 59 ++++++++++++
.../value-bounds-op-interface-impl.mlir | 91 +++++++++++++++++++
2 files changed, 150 insertions(+)
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 40475d37d2fd7..2244eb45a813a 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -50,6 +50,64 @@ struct AffineApplyOpInterface
}
};
+/// Express `expr`, a result of `map`, in terms of the constraint set by
+/// replacing the dims and symbols of `map` with the expressions for the
+/// corresponding `operands`.
+static AffineExpr alignBoundExpr(AffineExpr expr, AffineMap map,
+ ValueRange operands,
+ ValueBoundsConstraintSet &cstr) {
+ SmallVector<AffineExpr> dimReplacements =
+ llvm::map_to_vector(operands.take_front(map.getNumDims()),
+ [&](Value v) { return cstr.getExpr(v); });
+ SmallVector<AffineExpr> symReplacements =
+ llvm::map_to_vector(operands.drop_front(map.getNumDims()),
+ [&](Value v) { return cstr.getExpr(v); });
+ return expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
+}
+
+struct AffineForOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<AffineForOpInterface,
+ AffineForOp> {
+ void populateBoundsForIndexValue(Operation *op, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto forOp = cast<AffineForOp>(op);
+
+ // Only the induction variable is handled. Bounds for iter_args are not
+ // inferred.
+ if (value != forOp.getInductionVar())
+ return;
+
+ AffineMap lbMap = forOp.getLowerBoundMap();
+ AffineMap ubMap = forOp.getUpperBoundMap();
+ ValueRange lbOperands = forOp.getLowerBoundOperands();
+ ValueRange ubOperands = forOp.getUpperBoundOperands();
+
+ // The lower bound is the maximum over the results of `lbMap` and the upper
+ // bound is the minimum over the results of `ubMap`, so the induction
+ // variable is bounded by every individual result.
+ for (AffineExpr expr : lbMap.getResults())
+ cstr.bound(value) >= alignBoundExpr(expr, lbMap, lbOperands, cstr);
+ for (AffineExpr expr : ubMap.getResults())
+ cstr.bound(value) < alignBoundExpr(expr, ubMap, ubOperands, cstr);
+
+ // With a single lower and a single upper bound the step can be taken into
+ // account as well: the induction variable is always a multiple of `step`
+ // away from the lower bound, so it never exceeds
+ // `lb + (tripCount - 1) * step`. That is tighter than `ub - 1` whenever the
+ // trip count is not a multiple of the step, e.g. `affine.for %i = 0 to 300
+ // step 128` only ever yields {0, 128, 256}. This does not replace the
+ // `iv < ub` bound above, since multiplying two constraint set dimensions is
+ // not supported.
+ int64_t step = forOp.getStepAsInt();
+ if (step == 1 || lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
+ return;
+ AffineExpr lb = alignBoundExpr(lbMap.getResult(0), lbMap, lbOperands, cstr);
+ AffineExpr ub = alignBoundExpr(ubMap.getResult(0), ubMap, ubOperands, cstr);
+ AffineExpr tripCount = (ub - lb).ceilDiv(step);
+ cstr.bound(value) <= lb + (tripCount - 1) * step;
+ }
+};
+
struct AffineMinOpInterface
: public ValueBoundsOpInterface::ExternalModel<AffineMinOpInterface,
AffineMinOp> {
@@ -157,6 +215,7 @@ void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+ AffineForOp::attachInterface<AffineForOpInterface>(*ctx);
AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
AffineDelinearizeIndexOp::attachInterface<
diff --git a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
index a4310b91a37b3..beb80229bf857 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -236,3 +236,94 @@ func.func @linearize_static_no_outer_bound(%arg0: index, %arg1: index) -> index
"test.compare"(%0, %c6) {cmp = "LT"} : (index, index) -> ()
return %1 : index
}
+
+// -----
+
+// The induction variable of an affine.for with constant bounds is bounded below
+// by the lower bound and above by `lb + (tripCount - 1) * step`.
+
+func.func @affine_for_iv_constant_bounds() {
+ %c0 = arith.constant 0 : index
+ %c384 = arith.constant 384 : index
+ %c385 = arith.constant 385 : index
+ affine.for %i = 0 to 385 step 128 {
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c0) {cmp = "GE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c384) {cmp = "LE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c385) {cmp = "LT"} : (index, index) -> ()
+ }
+ return
+}
+
+// -----
+
+// Step alignment: `0 to 300 step 128` yields only {0, 128, 256}, so the
+// induction variable never exceeds 256 even though the upper bound is 300.
+
+func.func @affine_for_iv_step_alignment() {
+ %c256 = arith.constant 256 : index
+ affine.for %i = 0 to 300 step 128 {
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c256) {cmp = "LE"} : (index, index) -> ()
+ }
+ return
+}
+
+// -----
+
+// Non-zero lower bound: `5 to 300 step 128` yields {5, 133, 261}.
+
+func.func @affine_for_iv_nonzero_lb() {
+ %c5 = arith.constant 5 : index
+ %c261 = arith.constant 261 : index
+ affine.for %i = 5 to 300 step 128 {
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c5) {cmp = "GE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c261) {cmp = "LE"} : (index, index) -> ()
+ }
+ return
+}
+
+// -----
+
+// Bounds given as affine maps over loop-invariant values are handled as well,
+// including the step-aligned upper bound: `%n to %n + 300 step 128` yields
+// {%n, %n + 128, %n + 256}.
+
+func.func @affine_for_iv_symbolic_bounds(%n: index) {
+ %ub = affine.apply affine_map<()[s0] -> (s0 + 256)>()[%n]
+ affine.for %i = affine_map<()[s0] -> (s0)>()[%n]
+ to affine_map<()[s0] -> (s0 + 300)>()[%n] step 128 {
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %n) {cmp = "GE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %ub) {cmp = "LE"} : (index, index) -> ()
+ }
+ return
+}
+
+// -----
+
+// A `max`/`min` bound constrains the induction variable by every result of the
+// map. The step is not taken into account for such loops, so the tighter bound
+// `%i <= 256` implied by `step 128` cannot be proven here.
+
+func.func @affine_for_iv_multi_result_bounds(%n: index) {
+ %c0 = arith.constant 0 : index
+ %c256 = arith.constant 256 : index
+ %c300 = arith.constant 300 : index
+ affine.for %i = max affine_map<()[s0] -> (0, s0)>()[%n] to 300 step 128 {
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c0) {cmp = "GE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %n) {cmp = "GE"} : (index, index) -> ()
+ // expected-remark @below{{true}}
+ "test.compare"(%i, %c300) {cmp = "LT"} : (index, index) -> ()
+ // expected-error @below{{unknown}}
+ "test.compare"(%i, %c256) {cmp = "LE"} : (index, index) -> ()
+ }
+ return
+}
More information about the Mlir-commits
mailing list