This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 434dc7636e [ARITH][TIR] Track positive loop extents in analyzer
visitors (#19927)
434dc7636e is described below
commit 434dc7636e06c6c4a25eb1845869b596199c51bb
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Jul 5 21:23:12 2026 -0700
[ARITH][TIR] Track positive loop extents in analyzer visitors (#19927)
This PR updates `IRMutatorWithAnalyzer` and `IRVisitorWithAnalyzer` to
add the constraint `loop_extent > 0` while visiting a `ForNode` body.
If execution reaches the loop body, the loop must have at least one
iteration, so the body context can safely assume the extent is positive.
The constraint is scoped only to the loop body, leaving the loop header
expressions (`min`, `extent`, `step`) outside of this assumption.
While validating this change, it exposed an issue in `IntSetAnalyzer`:
one-sided constraints such as `m > 0` could cause existing parametric
bounds like `m - 1` to be recursively relaxed to `+inf`. This caused
`DomainTouched` to lose finite symbolic bounds. The PR fixes this by
preserving existing parametric bounds when recursive interval relaxation
would otherwise replace them with infinity.
---
src/arith/int_set.cc | 12 ++++++++++-
src/arith/ir_mutator_with_analyzer.cc | 24 +++++++++++++++++++++-
src/arith/ir_visitor_with_analyzer.cc | 10 ++++++++-
tests/python/arith/test_arith_intset.py | 18 ++++++++++++++++
.../tirx-transform/test_tir_transform_simplify.py | 16 +++++++++++++++
5 files changed, 77 insertions(+), 3 deletions(-)
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 64f46a605f..8a6920d787 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -423,7 +423,17 @@ class IntervalSetEvaluator : public
ExprFunctor<IntervalSet(const Expr&)> {
IntervalSet max_set = this->Eval(val->max_value);
--recur_depth_;
- return IntervalSet(min_set->min_value, max_set->max_value);
+ PrimExpr min_value = min_set->min_value;
+ PrimExpr max_value = max_set->max_value;
+ // IntSet keeps symbolic bounds parametric. If relaxing a bound under
+ // one-sided constraints loses it to infinity, keep the original bound.
+ if (is_neg_inf(min_value) && val->HasLowerBound()) {
+ min_value = val->min_value;
+ }
+ if (is_pos_inf(max_value) && val->HasUpperBound()) {
+ max_value = val->max_value;
+ }
+ return IntervalSet(min_value, max_value);
}
IntervalSet VisitExpr_(const IntImmNode* op) final {
diff --git a/src/arith/ir_mutator_with_analyzer.cc
b/src/arith/ir_mutator_with_analyzer.cc
index 440ce85ed7..98576efa2c 100644
--- a/src/arith/ir_mutator_with_analyzer.cc
+++ b/src/arith/ir_mutator_with_analyzer.cc
@@ -188,7 +188,29 @@ Stmt IRMutatorWithAnalyzer::VisitStmt_(const ForNode* op) {
Range dom = Range::FromMinExtent(op->min, op->extent);
analyzer_->Bind(op->loop_var, dom);
iter_vars_.Set(op->loop_var, dom);
- return StmtExprMutator::VisitStmt_(op);
+
+ PrimExpr min = this->VisitPrimExpr(op->min);
+ PrimExpr extent = this->VisitPrimExpr(op->extent);
+ ffi::Optional<PrimExpr> step{std::nullopt};
+ if (op->step.has_value()) {
+ step = this->VisitPrimExpr(*op->step);
+ }
+ Stmt body = constraint_scope_.WithNewScope([&]() -> Stmt {
+ EnterConstraintFacts(&constraint_scope_.Current(), analyzer_,
+ extent > IntImm(extent.ty(), 0));
+ return this->VisitStmt(op->body);
+ });
+ if (min.same_as(op->min) && extent.same_as(op->extent) &&
body.same_as(op->body) &&
+ step.same_as(op->step)) {
+ return ffi::GetRef<Stmt>(op);
+ } else {
+ auto n = this->CopyOnWrite(op);
+ n->min = std::move(min);
+ n->extent = std::move(extent);
+ n->step = std::move(step);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
});
}
diff --git a/src/arith/ir_visitor_with_analyzer.cc
b/src/arith/ir_visitor_with_analyzer.cc
index 11ba198115..89949b82f0 100644
--- a/src/arith/ir_visitor_with_analyzer.cc
+++ b/src/arith/ir_visitor_with_analyzer.cc
@@ -36,7 +36,15 @@ using namespace tirx;
void IRVisitorWithAnalyzer::VisitStmt_(const ForNode* op) {
constraint_scope_.WithNewScope([&]() {
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
- StmtExprVisitor::VisitStmt_(op);
+ this->VisitExpr(op->min);
+ this->VisitExpr(op->extent);
+ if (op->step.has_value()) {
+ this->VisitExpr(*op->step);
+ }
+ constraint_scope_.WithNewScope([&]() {
+ constraint_scope_.Current().Emplace(analyzer_, op->extent >
IntImm(op->extent.ty(), 0));
+ this->VisitStmt(op->body);
+ });
});
}
diff --git a/tests/python/arith/test_arith_intset.py
b/tests/python/arith/test_arith_intset.py
index e83a8c9f9d..d526bff241 100644
--- a/tests/python/arith/test_arith_intset.py
+++ b/tests/python/arith/test_arith_intset.py
@@ -424,6 +424,24 @@ def test_relax_cyclic_variable_dependency():
assert res is not None
+def test_constraint_scope_preserves_parametric_bounds():
+ """One-sided constraints should not erase symbolic bounds."""
+ analyzer = tvm.arith.Analyzer()
+ i = tvm.tirx.Var("i", "int32")
+ m = tvm.tirx.Var("m", "int32")
+ analyzer.bind(i, tvm.ir.Range.from_min_extent(0, m))
+
+ with analyzer.constraint_scope(m > 0):
+ res = analyzer.int_set(i - 1)
+ assert analyzer.can_prove_equal(res.min_value, -1)
+ assert analyzer.can_prove_equal(res.max_value, m - 2)
+
+ with analyzer.constraint_scope(m < 16):
+ res = analyzer.int_set(i)
+ assert analyzer.can_prove_equal(res.min_value, 0)
+ assert analyzer.can_prove_equal(res.max_value, 14)
+
+
def test_estimate_region_accepts_external_analyzer():
i = tvm.tirx.Var("i", "int32")
tile = tvm.tirx.Var("tile", "int32")
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py
b/tests/python/tirx-transform/test_tir_transform_simplify.py
index c2121ebeca..47da902d51 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -96,6 +96,22 @@ def test_if_likely():
assert not isinstance(body.body.body.then_case, tvm.tirx.IfThenElse)
+def test_loop_body_knows_dynamic_extent_is_positive():
+ @T.prim_func(private=True, s_tir=True)
+ def before(A: T.Buffer((1,), "float32"), m: T.int32, n: T.int32):
+ for i in T.serial(m, n // 4):
+ if n // 4 - m > 0:
+ A[0] = 1.0
+
+ @T.prim_func(private=True, s_tir=True)
+ def expected(A: T.Buffer((1,), "float32"), m: T.int32, n: T.int32):
+ for i in T.serial(m, n // 4):
+ A[0] = 1.0
+
+ after =
tvm.tirx.transform.StmtSimplify()(tvm.IRModule.from_expr(before))["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
def _apply_simplify(
func,
transitively_prove_inequalities=False,
