[Arith] Fix dom_map var leak in IntervalSet relaxation

src/arith/int_set.cc

@@ -426,11 +426,16 @@ class IntervalSetEvaluator : public ExprFunctor<IntervalSet(const Expr&)> {
     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<Var>(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);

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")