diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc index 6c91ab808010..326dd1b2ae79 100644 --- a/src/arith/int_set.cc +++ b/src/arith/int_set.cc @@ -426,11 +426,16 @@ class IntervalSetEvaluator : public ExprFunctor { 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()) { + // one-sided constraints loses it to infinity, keep the original bound, + // unless it references a dom_map_ var that must be relaxed out. + auto f_depends_on_relaxed = [this](const PrimExpr& bound) { + return UsesVar(bound, + [this](const VarNode* v) { return dom_map_.count(ffi::GetRef(v)) != 0; }); + }; + if (is_neg_inf(min_value) && val->HasLowerBound() && !f_depends_on_relaxed(val->min_value)) { min_value = val->min_value; } - if (is_pos_inf(max_value) && val->HasUpperBound()) { + if (is_pos_inf(max_value) && val->HasUpperBound() && !f_depends_on_relaxed(val->max_value)) { max_value = val->max_value; } return IntervalSet(min_value, max_value); diff --git a/tests/python/arith/test_arith_intset.py b/tests/python/arith/test_arith_intset.py index d526bff2411f..8a023c17d474 100644 --- a/tests/python/arith/test_arith_intset.py +++ b/tests/python/arith/test_arith_intset.py @@ -442,6 +442,19 @@ def test_constraint_scope_preserves_parametric_bounds(): assert analyzer.can_prove_equal(res.max_value, 14) +def test_relax_does_not_leak_dom_map_vars(): + analyzer = tvm.arith.Analyzer() + i = tvm.tirx.Var("i", "int64") + j = tvm.tirx.Var("j", "int64") + zero = tvm.tirx.const(0, "int64") + dom = { + j: tvm.arith.IntervalSet(zero, i), + i: tvm.arith.IntervalSet(zero, tvm.arith.int_set.pos_inf()), + } + res = analyzer.int_set(j, dom) + assert res.max_value.same_as(tvm.arith.int_set.pos_inf()) + + def test_estimate_region_accepts_external_analyzer(): i = tvm.tirx.Var("i", "int32") tile = tvm.tirx.Var("tile", "int32")