Skip to content

[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod #19978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tqchen merged 4 commits into from
Jul 13, 2026
+138 −31
Merged

[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod #19978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ed75bb9
[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod
sbinabdullah Jul 10, 2026
f068050
[Arith] Intersect modular-set mod bounds with interval bounds
sbinabdullah Jul 10, 2026
eccd4ae
[Arith] Tighten truncmod bound for entirely-negative dividends with |…
sbinabdullah Jul 11, 2026
cf7b5c7
[Tests] xfail test_adaptive_pooling_window after const-int-bound fix
sbinabdullah Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 86 additions & 31 deletions src/arith/const_int_bound.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),

@spectrometerHBH spectrometerHBH Jul 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe there can be tigher bound when amax < 0.

@sbinabdullah sbinabdullah Jul 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

really good catch I think you’re right. will add an additional check for this.

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)
Expand All @@ -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);
}
Comment on lines 330 to 347

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When gcd_coeff_mod > 1, the modular-set-based bound is immediately returned, which completely discards the interval-based bounds of the dividend a. If a has a tight range (e.g., a.max_value < b.min_value), this can lead to unnecessarily loose bounds. Intersecting the modular-set-based bound with the fallback interval-based bound preserves maximum precision.

        if (gcd_coeff_mod > 1) {
          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;
          Entry fallback_bound;
          if (a.min_value >= 0) {
            if (a.max_value < b.min_value) {
              fallback_bound = a;
            } else {
              fallback_bound = MakeBound(0, std::min(a.max_value, b_max_cap));
            }
            return Intersect(fallback_bound, MakeBound(base_mod, tight_max));
          }
          int64_t neg_base = (gcd_coeff_mod - base_mod) % gcd_coeff_mod;
          int64_t tight_min = -(modulus - gcd_coeff_mod + neg_base);
          fallback_bound = MakeBound(std::max(a.min_value, -b_max_cap),
                                     std::min(std::max(a.max_value, (int64_t)0), b_max_cap));
          if (a.max_value < 0) {
            return Intersect(fallback_bound, MakeBound(tight_min, -neg_base));
          }
          return Intersect(fallback_bound, MakeBound(tight_min, tight_max));
        }

@sbinabdullah sbinabdullah Jul 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in 6fc6bfc. Both visitors now compute the interval-based bound first and return Intersect(interval_bound, modular_bound) when the modular fast path applies, so a tight dividend range is never lost. Added regression cases: (n*64+63) % 256 with n in [0,1] now gives [63, 127] (was [63, 255]), plus truncmod/floormod variants with a negative dividend range.

}

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,
Expand Down Expand Up @@ -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)
Expand All @@ -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));
}
Comment on lines 428 to 433

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the Mod visitor, immediately returning the modular-set-based bound for FloorMod discards the interval-based bounds of a. Intersecting the modular-set-based bound with the fallback interval-based bound ensures we do not lose precision when a has a tight range.

        if (gcd_coeff_mod > 1) {
          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;
          Entry fallback_bound;
          if (a.min_value >= 0) {
            if (a.max_value < b.min_value) {
              fallback_bound = a;
            } else {
              fallback_bound = MakeBound(0, std::min(a.max_value, b_max_cap));
            }
          } else {
            fallback_bound = MakeBound(0, b_max_cap);
          }
          return Intersect(fallback_bound, MakeBound(base_mod, tight_max));
        }

@sbinabdullah sbinabdullah Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 6fc6bfc together with the Mod visitor — FloorMod now intersects the modular-set bound with the interval-based bound as well (regression case: (n*64+63) % 256 with n in [0,1] → [63, 127]).

}

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);
Expand Down
47 changes: 47 additions & 0 deletions tests/python/arith/test_arith_const_int_bound.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
5 changes: 5 additions & 0 deletions tests/python/te/test_te_create_primfunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -911,6 +911,11 @@ def te_workload():
_check_workload(te_workload, tir_workload)


@pytest.mark.xfail(
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(
Expand Down
Loading