diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc index 6c91ab808010..1e6f7497ea9e 100644 --- a/src/arith/int_set.cc +++ b/src/arith/int_set.cc @@ -423,17 +423,7 @@ class IntervalSetEvaluator : public ExprFunctor { 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 VisitExpr_(const VarNode* op) final { Var var = ffi::GetRef(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()); + } + ffi::Array values; if (dom_constraints_) { for (const auto& constraint : *dom_constraints_) { @@ -452,14 +449,7 @@ class IntervalSetEvaluator : public ExprFunctor { } } - 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()); - } + 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 d526bff2411f..e83a8c9f9d41 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 3185147f8e14..f0000581e1f1 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 @@ def expected(): 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"""