[ARITH] fix int set analysis on negative scale

src/arith/int_set.cc

@@ -860,9 +860,18 @@ Optional<Array<IntSet>> EstimateRegionLowerBound(const Array<Range>& region,
     const PrimExpr& base = sum_expr->base;
     // IterSplitExpr: (source // lower_factor) % extent * scale
     // where `(source // lower_factor) % extent` is within [0, extent - 1]
-    // Therefore, the range of `region[i]->min` is `base + [0, (extent - 1) * scale]`
-    result.push_back(
-        IntSet::FromMinExtent(base, split->extent * split->scale + (range->extent - split->scale)));
+    if (analyzer->CanProve(split->scale < 0)) {
+      // If scale is negative, the var dom is [(extent - 1) * scale, 0]
+      // The total base is `base + (extent - 1) * scale`,
+      // while total extent is `dom_extent + (extent - 1) * (-scale)`
+      const PrimExpr& var_extent = (split->extent - 1) * split->scale;
+      result.push_back(IntSet::FromMinExtent(base + var_extent, range->extent - var_extent));
+    } else {
+      // If scale is positive, the var dom is [0, (extent - 1) * scale]
+      // The total dom is [base, dom_extent + (extent - 1) * scale]
+      result.push_back(
+          IntSet::FromMinExtent(base, range->extent + (split->extent - 1) * split->scale));
+    }
   }
   return result;
 }

tests/python/unittest/test_arith_intset.py

@@ -213,6 +213,26 @@ def test_region_lower_bound_multiple_variables():
     assert k_int_set.max_value.value == 31
 
 
+def test_region_lower_bound_negative_scale():
+    i = tvm.tir.Var("i", "int32")
+    j = tvm.tir.Var("j", "int32")
+    int_set_0, int_set_1 = tvm.arith.estimate_region_lower_bound(
+        region=[
+            tvm.ir.Range.from_min_extent(min_value=1 - i, extent=4),
+            tvm.ir.Range.from_min_extent(min_value=20 - j * 4, extent=16),
+        ],
+        var_dom={
+            i: tvm.ir.Range(begin=0, end=4),
+            j: tvm.ir.Range(begin=0, end=4),
+        },
+        predicate=tvm.tir.IntImm("bool", 1),
+    )
+    assert int_set_0.min_value.value == -2
+    assert int_set_0.max_value.value == 4
+    assert int_set_1.min_value.value == 8
+    assert int_set_1.max_value.value == 35
+
+
 def test_union_lower_bound():
     neg_inf = tvm.arith.int_set.neg_inf()
     pos_inf = tvm.arith.int_set.pos_inf()
@@ -236,4 +256,5 @@ def test_union_lower_bound():
     test_region_lower_bound_small_stride()
     test_region_lower_bound_split_predicate()
     test_region_lower_bound_multiple_variables()
+    test_region_lower_bound_negative_scale()
     test_union_lower_bound()