[ARITH] Scope interval constraints to mapped variables

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

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

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