[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod#19978
[Arith] Fix const-int-bound modular-set tightening for Mod/FloorMod#19978sbinabdullah wants to merge 2 commits into
Conversation
The "tighter bounds using modular set" fast path in ConstIntBound's
Mod/FloorMod visitors normalized the modular-set base with
`base % modulus` instead of `base % gcd(coeff, modulus)`. When the base
was not already smaller than the gcd, this produced invalid bounds with
min_value > max_value.
Example: for free int64 n,
const_int_bound((n * 320 + 255) % 256)
returned [255, 191], while the true value set is {63, 127, 191, 255},
i.e. [63, 255].
This is more than a precision issue: the invalid bound made
Analyzer::CanProve(..., kSymbolicBound) "prove" the bounds predicate of
imperfect dynamic loop splits (e.g. fuse(n, 5, 64) then split by 256:
bx*256 + tx < n*320), so tir schedule split() silently dropped the
T.where predicate. Every dlight gpu.Fallback-scheduled kernel with a
dynamic fused extent then ran without an out-of-bounds guard. In
practice this corrupted the paged-KV cache on WebGPU: the unguarded
tir_kv_cache_transpose_append kernel's excess threads read
position_map past ntoken (getting 0, not the -1 sentinel) and
overwrote page 0 / slot 0 with stale data on every decode.
Fix: normalize the residue modulo the gcd. For FloorMod (divisor > 0)
the result is in {r, r + g, ..., modulus - g + r} with
r = base % g >= 0. For truncated Mod, mirror the residue set on the
negative side when the dividend can be negative.
Adds const-int-bound regression tests covering the floormod case, the
truncmod mixed-sign / non-negative cases, and the original base=0
example from the comment.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes a bug in ConstIntBoundAnalyzer where the residue for modular-set-based bounds was normalized modulo the divisor instead of the GCD of the coefficient and divisor, which could lead to invalid bounds and incorrect optimization of out-of-bounds guards. It also adds support for negative dividends in truncmod and introduces regression tests. The reviewer feedback suggests that immediately returning the modular-set-based bounds for both truncmod and floormod discards the interval-based bounds of the dividend. To prevent losing precision when the dividend has a tight range, the reviewer recommends intersecting the modular-set-based bounds with the fallback interval-based bounds.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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); | ||
| if (a.min_value >= 0) { | ||
| return 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); | ||
| if (a.max_value < 0) { | ||
| return MakeBound(tight_min, -neg_base); | ||
| } | ||
| return MakeBound(tight_min, tight_max); | ||
| } |
There was a problem hiding this comment.
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));
}There was a problem hiding this comment.
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 (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); | ||
| } |
There was a problem hiding this comment.
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));
}There was a problem hiding this comment.
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]).
Address review feedback: returning the modular-set-based bound directly
discarded the interval-based bound of the dividend, so a tight dividend
range (e.g. a in [63, 127] with a == 63 (mod 64), where floormod(a, 256)
is exactly {63, 127}) produced an unnecessarily loose [63, 255].
Restructure the Mod/FloorMod visitors to always compute the
interval-based bound first and intersect it with the modular-set-based
bound when the fast path applies. Adds tight-range regression cases for
floormod, truncmod with a negative dividend range, and floormod of a
negative dividend range.
|
CI note: the failing
I reproduced this on a clean detached worktree at the PR target SHA pre-commit run ruff-format --all-files --show-diff-on-failureIt produces the exact same two-file diff and exits 1. Upstream |
tlopex
left a comment
tlopex
left a comment
There was a problem hiding this comment.
Thanks for catching this. it looks good to me
2 participants
Summary
The "tighter bounds using modular set" fast path in
ConstIntBoundAnalyzer'sMod/FloorModvisitors normalizes the modular-set base withbase % modulusinstead ofbase % gcd(coeff, modulus). When the base is not already smaller than the gcd, this produces an invalid bound withmin_value > max_value.Example (free
int64n):The true value set is
{63, 127, 191, 255}: sincegcd(320, 256) = 64, the dividend satisfiesa == 255 == 63 (mod 64), sofloormod(a, 256)can only take values63 + 64k.Why it matters (miscompilation, not just precision)
The invalid bound makes
Analyzer::CanProve(..., kSymbolicBound)"prove" the bounds predicate of imperfect dynamic loop splits. For example, aftersch.fuse()ofT.grid(n, 5, 64)andsch.split(..., factors=[None, 256]), the predicatewas deemed provable (the symbolic-bound path reduces it to
const_int_bound((n*320 + 255) % 256 - 254).min_value >= 1, which the broken[255, 191]bound satisfies), sosplit()silently dropped theT.wherepredicate. Any kernel scheduled through dlightgpu.Fallbackwith a dynamic fused extent then ran without an out-of-bounds guard.Downstream, this corrupted the paged-KV cache on WebGPU: the unguarded
tir_kv_cache_transpose_appendkernel's excess threads readposition_mappastntoken(zero-initialized memory returns0, not the-1sentinel) and overwrote page 0 / slot 0 of the KV cache with stale data on every decode, breaking multi-sequence decoding.Fix
Normalize the residue modulo the gcd:
FloorMod(divisor > 0): result is in{r, r + g, ..., modulus - g + r}withr = base % gnormalized to[0, g).Mod: same set on the non-negative side; when the dividend can be negative, mirror the residue set (r' = (g - r) % g) on the negative side.Testing
tests/python/arith/test_arith_const_int_bound.py(TestModBoundWithModularSet) covering the floormod case above, the truncmod mixed-sign / non-negative-dividend cases, and the originalbase = 0example from the code comment.tests/python/arithpasses (964 passed);tests/python/s_tir/schedule -k "split or fuse"andtests/python/s_tir/dlightpass.sch.fuse+ imperfectsch.spliton a dynamic extent emits theT.wherebounds predicate again, and that the recompiled WebGPU kernel contains the bounds guard.