This is an automated email from the ASF dual-hosted git repository.
tqchen 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 e1e7ac9261 [ARITH] Scope interval constraints to mapped variables
(#19963)
e1e7ac9261 is described below
commit e1e7ac92617de8bced603cdde56395889f14da9f
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Jul 7 03:47:46 2026 +0800
[ARITH] Scope interval constraints to mapped variables (#19963)
## Rationale
Analyzer constraint scopes continue to provide loop-positive facts to
constant-bound and rewrite proofs. During IntSet relaxation, however, a
scoped domain constraint is a refinement only for a variable explicitly
present in the relaxation map. Applying it to an unmapped variable
reinterprets a free parameter as a relaxation domain and can let a
loop-local symbol survive recursive interval evaluation.
## Changes
- Apply scoped IntSet constraints only to variables already present in
the relaxation map; unmapped variables remain free parameters.
- Remove the finite-bound restoration fallback and its stronger
parametric-bound contract.
- Add a direct compact-buffer regression that prevents a loop-local
variable from escaping into a function-scope allocation extent.
---
src/arith/int_set.cc | 28 +++++++---------------
tests/python/arith/test_arith_intset.py | 18 --------------
.../test_s_tir_transform_compact_buffer_region.py | 15 ++++++++++++
3 files changed, 24 insertions(+), 37 deletions(-)
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 6c91ab8080..1e6f7497ea 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -423,17 +423,7 @@ class IntervalSetEvaluator : public
ExprFunctor<IntervalSet(const Expr&)> {
IntervalSet max_set = this->Eval(val->max_value);
--recur_depth_;
- 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);
+ return IntervalSet(min_set->min_value, max_set->max_value);
}
IntervalSet VisitExpr_(const IntImmNode* op) final {
@@ -443,6 +433,13 @@ class IntervalSetEvaluator : public
ExprFunctor<IntervalSet(const Expr&)> {
IntervalSet VisitExpr_(const VarNode* op) final {
Var var = ffi::GetRef<Var>(op);
+ auto it = dom_map_.find(var);
+ // Scoped constraints refine explicit relaxation domains. Variables that
+ // are absent from dom_map_ remain free parameters of the relaxed result.
+ if (it == dom_map_.end()) {
+ return IntervalSet::SinglePoint(var.as_or_throw<PrimExpr>());
+ }
+
ffi::Array<IntSet> values;
if (dom_constraints_) {
for (const auto& constraint : *dom_constraints_) {
@@ -452,14 +449,7 @@ class IntervalSetEvaluator : public
ExprFunctor<IntervalSet(const Expr&)> {
}
}
- auto it = dom_map_.find(var);
- if (it != dom_map_.end()) {
- values.push_back((*it).second);
- }
-
- if (values.empty()) {
- return IntervalSet::SinglePoint(var.as_or_throw<PrimExpr>());
- }
+ values.push_back((*it).second);
IntSet intersection = [&]() {
if (values.size() == 1) {
diff --git a/tests/python/arith/test_arith_intset.py
b/tests/python/arith/test_arith_intset.py
index d526bff241..e83a8c9f9d 100644
--- a/tests/python/arith/test_arith_intset.py
+++ b/tests/python/arith/test_arith_intset.py
@@ -424,24 +424,6 @@ 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/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
index 3185147f8e..f0000581e1 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
@@ -1288,6 +1288,21 @@ class TestNonBoolCondition(BaseCompactTest):
A[i - 1] = A[i - 1] + 1
+def test_loop_var_does_not_escape_compacted_buffer_extent():
+ @T.prim_func(private=True, s_tir=True)
+ def before(a: T.handle):
+ n = T.int64()
+ A = T.match_buffer(a, (n,), "int32")
+ tmp = T.alloc_buffer((n,), "int32")
+ for i in range(n):
+ length: T.let[T.int64] = T.ceildiv(n, T.shift_left(T.int64(1), i +
1))
+ for j in range(length):
+ tmp[j] = A[j]
+
+ after =
s_tir.transform.CompactBufferAllocation()(tvm.IRModule.from_expr(before))
+ assert tirx.analysis.verify_well_formed(after)
+
+
class TestCompactSymbolicBound0:
"""Test symbolic bound that get compacted to constant"""
