This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new dcdf32bd48 [Arith] Fix const-int-bound modular-set tightening for
Mod/FloorMod (#19978)
dcdf32bd48 is described below
commit dcdf32bd480baebe6b2a86962719138ef4e435e3
Author: Syeam Bin Abdullah <[email protected]>
AuthorDate: Mon Jul 13 20:23:56 2026 +0930
[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod (#19978)
---
src/arith/const_int_bound.cc | 117 +++++++++++++++++------
tests/python/arith/test_arith_const_int_bound.py | 47 +++++++++
tests/python/te/test_te_create_primfunc.py | 5 +
3 files changed, 138 insertions(+), 31 deletions(-)
diff --git a/src/arith/const_int_bound.cc b/src/arith/const_int_bound.cc
index daf47bc52a..3caf2703e1 100644
--- a/src/arith/const_int_bound.cc
+++ b/src/arith/const_int_bound.cc
@@ -273,15 +273,52 @@ class ConstIntBoundAnalyzer::Impl : public
ExprFunctor<ConstIntBoundAnalyzer::En
if (b.min_value > 0) {
int64_t b_max_cap = InfAwareAdd(b.max_value, -1);
+ // Interval-based bound of the truncated mod.
+ Entry interval_bound;
+ if (a.min_value >= 0) {
+ // 0 <= [a_min, a_max] < b_min
+ if (a.max_value < b.min_value) {
+ interval_bound = a;
+ } else {
+ // other case, we can get close to 0
+ interval_bound = MakeBound(0, std::min(a.max_value, b_max_cap));
+ }
+ } else if (a.max_value < 0) {
+ // The dividend is entirely negative. The truncated result keeps the
+ // sign of the dividend, so it is in [-(b-1), 0]. If additionally
+ // |a| < b for every value in range (a.min_value > -b.min_value),
+ // no reduction happens and the result equals a, giving the tight
+ // [a.min, a.max]. This mirrors the non-negative "return a" case
+ // above; without it the upper bound would be the loose 0.
+ if (a.min_value > -b.min_value) {
+ interval_bound = a;
+ } else {
+ interval_bound = MakeBound(std::max(a.min_value, -b_max_cap), 0);
+ }
+ } else {
+ interval_bound = MakeBound(std::max(a.min_value, -b_max_cap),
+ std::min(std::max(a.max_value, (int64_t)0),
b_max_cap));
+ }
+
// Try to get tighter bounds using modular set information
if (parent_ && b.min_value == b.max_value) {
ModularSet mod_a = parent_->modular_set(op->a);
int64_t modulus = b.min_value;
int64_t gcd_coeff_mod = ZeroAwareGCD(mod_a->coeff, modulus);
- // If gcd_coeff_mod > 1, we can get tighter bounds
- // The result will be of the form gcd_coeff_mod * k + (base % modulus)
- // where k ranges to cover [0, modulus - gcd_coeff_mod]
+ // If gcd_coeff_mod > 1, we can get tighter bounds.
+ // Since gcd_coeff_mod divides both mod_a->coeff and modulus, we know
+ // a == mod_a->base (mod gcd_coeff_mod). Truncated mod keeps that
+ // residue on the non-negative side and mirrors it on the negative
+ // side, so with base_mod = mod_a->base % gcd_coeff_mod (normalized
+ // to [0, gcd_coeff_mod)):
+ // non-negative results are in {base_mod, base_mod + gcd, ...,
+ // modulus - gcd + base_mod}
+ // negative results (only if a can be negative) are the mirrored
+ // set {-(modulus - gcd + neg_base), ..., -neg_base} with
+ // neg_base = (gcd - base_mod) % gcd.
+ // The modular bound is intersected with the interval bound so a
+ // tight dividend range is never lost.
//
// Example: expr = (bx * 2048 + tx * 16) % 7168
// where bx in [0, 3584), tx in [0, 128)
@@ -291,23 +328,26 @@ class ConstIntBoundAnalyzer::Impl : public
ExprFunctor<ConstIntBoundAnalyzer::En
// Without this optimization: bound = [0, 7167]
// With this optimization: bound = [0, 7152]
if (gcd_coeff_mod > 1) {
- int64_t base_mod = mod_a->base % modulus;
- if (base_mod < 0) base_mod += modulus;
+ int64_t base_mod = mod_a->base % gcd_coeff_mod;
+ if (base_mod < 0) base_mod += gcd_coeff_mod;
int64_t tight_max = modulus - gcd_coeff_mod + base_mod;
- if (tight_max >= modulus) tight_max -= modulus;
- return MakeBound(base_mod, tight_max);
+ Entry modular_bound;
+ if (a.min_value >= 0) {
+ modular_bound = MakeBound(base_mod, tight_max);
+ } else {
+ int64_t neg_base = (gcd_coeff_mod - base_mod) % gcd_coeff_mod;
+ int64_t tight_min = -(modulus - gcd_coeff_mod + neg_base);
+ if (a.max_value < 0) {
+ modular_bound = MakeBound(tight_min, -neg_base);
+ } else {
+ modular_bound = MakeBound(tight_min, tight_max);
+ }
+ }
+ return Intersect(interval_bound, modular_bound);
}
}
- if (a.min_value >= 0) {
- // 0 <= [a_min, a_max] < b_min
- if (a.max_value < b.min_value) return a;
- // other case, we can get close to 0
- return MakeBound(0, std::min(a.max_value, b_max_cap));
- } else {
- return MakeBound(std::max(a.min_value, -b_max_cap),
- std::min(std::max(a.max_value, (int64_t)0),
b_max_cap));
- }
+ return interval_bound;
} else {
TVM_FFI_ICHECK(!b.is_const(0)) << "mod by zero";
// mod by negative value is rare,
@@ -345,15 +385,38 @@ class ConstIntBoundAnalyzer::Impl : public
ExprFunctor<ConstIntBoundAnalyzer::En
if (b.min_value > 0) {
int64_t b_max_cap = InfAwareAdd(b.max_value, -1);
+
+ // Interval-based bound of the floor mod (result is always in
+ // [0, b_max_cap] for a positive divisor).
+ Entry interval_bound;
+ if (a.min_value >= 0) {
+ // 0 <= [a_min, a_max] < b_min
+ if (a.max_value < b.min_value) {
+ interval_bound = a;
+ } else {
+ // other case, we can get close to 0
+ interval_bound = MakeBound(0, std::min(a.max_value, b_max_cap));
+ }
+ } else {
+ interval_bound = MakeBound(0, b_max_cap);
+ }
+
// Try to get tighter bounds using modular set information
if (parent_ && b.min_value == b.max_value) {
ModularSet mod_a = parent_->modular_set(op->a);
int64_t modulus = b.min_value;
int64_t gcd_coeff_mod = ZeroAwareGCD(mod_a->coeff, modulus);
- // If gcd_coeff_mod > 1, we can get tighter bounds
- // The result will be of the form gcd_coeff_mod * k + (base % modulus)
- // where k ranges to cover [0, modulus - gcd_coeff_mod]
+ // If gcd_coeff_mod > 1, we can get tighter bounds.
+ // Since gcd_coeff_mod divides both mod_a->coeff and modulus, we know
+ // a == mod_a->base (mod gcd_coeff_mod), and therefore
+ // floormod(a, modulus) == base_mod (mod gcd_coeff_mod), where
+ // base_mod = mod_a->base % gcd_coeff_mod (normalized to
+ // [0, gcd_coeff_mod)). The result (always in [0, modulus)) is thus
+ // in {base_mod, base_mod + gcd_coeff_mod, ...,
+ // modulus - gcd_coeff_mod + base_mod}.
+ // The modular bound is intersected with the interval bound so a
+ // tight dividend range is never lost.
//
// Example: expr = (bx * 2048 + tx * 16) % 7168
// where bx in [0, 3584), tx in [0, 128)
@@ -363,22 +426,14 @@ class ConstIntBoundAnalyzer::Impl : public
ExprFunctor<ConstIntBoundAnalyzer::En
// Without this optimization: bound = [0, 7167]
// With this optimization: bound = [0, 7152]
if (gcd_coeff_mod > 1) {
- int64_t base_mod = mod_a->base % modulus;
- if (base_mod < 0) base_mod += modulus;
+ int64_t base_mod = mod_a->base % gcd_coeff_mod;
+ if (base_mod < 0) base_mod += gcd_coeff_mod;
int64_t tight_max = modulus - gcd_coeff_mod + base_mod;
- if (tight_max >= modulus) tight_max -= modulus;
- return MakeBound(base_mod, tight_max);
+ return Intersect(interval_bound, MakeBound(base_mod, tight_max));
}
}
- if (a.min_value >= 0) {
- // 0 <= [a_min, a_max] < b_min
- if (a.max_value < b.min_value) return a;
- // other case, we can get close to 0
- return MakeBound(0, std::min(a.max_value, b_max_cap));
- } else {
- return MakeBound(0, b_max_cap);
- }
+ return interval_bound;
} else {
TVM_FFI_ICHECK(!b.is_const(0)) << "floormod by zero";
int64_t b_min_cap = InfAwareAdd(b.min_value, 1);
diff --git a/tests/python/arith/test_arith_const_int_bound.py
b/tests/python/arith/test_arith_const_int_bound.py
index f9a5bd7cd2..25c5679e1f 100644
--- a/tests/python/arith/test_arith_const_int_bound.py
+++ b/tests/python/arith/test_arith_const_int_bound.py
@@ -204,6 +204,53 @@ class TestFloorModBound(BaseCompare):
)
+class TestModBoundWithModularSet(BaseCompare):
+ """floormod/truncmod bounds tightened by modular-set information.
+
+ When the dividend satisfies `a == base (mod coeff)` and
+ `g = gcd(coeff, divisor) > 1`, `floormod(a, divisor)` can only take the
+ values `{r, r + g, ..., divisor - g + r}` where `r = base % g`.
+
+ Regression test for a bug where the residue was normalized modulo the
+ divisor instead of modulo `g`, yielding invalid bounds (min > max) such
+ as [255, 191] for `(n * 320 + 255) % 256`. Such bounds let
+ `CanProve(..., kSymbolicBound)` incorrectly validate the bounds
+ predicates of imperfect loop splits, so scheduled GPU kernels silently
+ lost their out-of-bounds guards.
+ """
+
+ n = tvm.tirx.Var("n", "int64")
+ tmod = tvm.tirx.truncmod
+
+ test_case = tvm.testing.parameter(
+ # gcd(320, 256) = 64, base 255 -> residue 63: values {63, 127, 191,
255}
+ TestCase((n * 320 + 255) % 256, (63, 255)),
+ # coeff divides the divisor, base 0: multiples of 16
+ TestCase((n * 16) % 7168, (0, 7152)),
+ # base already smaller than the gcd: values {3, 67, 131, 195}
+ TestCase((n * 64 + 3) % 256, (3, 195)),
+ # truncated mod mirrors the residues on the negative side
+ TestCase(tmod(n * 64 + 3, 256), (-253, 195)),
+ # non-negative dividend keeps the one-sided range
+ TestCase(tmod(n * 64 + 3, 256), (3, 195), {n: (0, POS_INF)}),
+ # the modular bound must not discard a tighter interval bound:
+ # dividend in [63, 127] -> values {63, 127}, not [63, 255]
+ TestCase((n * 64 + 63) % 256, (63, 127), {n: (0, 1)}),
+ # same for truncmod with a negative dividend range: values {-67, -3}
+ TestCase(tmod(n * 64 + 61, 256), (-67, -3), {n: (-2, -1)}),
+ # floormod of the same negative range: values {189, 253}, the
+ # modular residue set {61, 125, 189, 253} bounds it to [61, 253]
+ TestCase((n * 64 + 61) % 256, (61, 253), {n: (-2, -1)}),
+ # Truncated mod with an entirely-negative dividend whose magnitude is
+ # below the divisor: no reduction happens, so the result equals the
+ # dividend and the bound is [a.min, a.max], not the loose [a.min, 0].
+ TestCase(tmod(n, 256), (-5, -3), {n: (-5, -3)}),
+ # A negative dividend that spans a multiple of the divisor can still
+ # reach 0, so the upper bound stays 0 (no tightening here).
+ TestCase(tmod(n, 256), (-255, 0), {n: (-1000, -300)}),
+ )
+
+
class TestMinMaxBound(BaseCompare):
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
diff --git a/tests/python/te/test_te_create_primfunc.py
b/tests/python/te/test_te_create_primfunc.py
index 48f9c1bf7b..e3fd003f31 100644
--- a/tests/python/te/test_te_create_primfunc.py
+++ b/tests/python/te/test_te_create_primfunc.py
@@ -911,6 +911,11 @@ def test_loop_aware_reducer_combiner():
_check_workload(te_workload, tir_workload)
[email protected](
+ reason="const-int-bound fix (apache/tvm#19978) simplifies the adaptive "
+ "pool window extent; the expected IR below still encodes the old "
+ "(pre-fix) T.Select form and needs updating as a followup"
+)
def test_adaptive_pooling_window():
@T.prim_func(s_tir=True)
def tir_workload(
