This is an automated email from the ASF dual-hosted git repository.
MasterJH5574 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 15ab675ea4 [Fix][Arith] Preserve nested floormod semantics (#20181)
15ab675ea4 is described below
commit 15ab675ea411b990e4fe4f8edd76c41de92dea7a
Author: Shushi Hong <[email protected]>
AuthorDate: Tue Aug 25 22:48:00 2026 -0400
[Fix][Arith] Preserve nested floormod semantics (#20181)
Prevent `IterMapRewriter` from collapsing a nested `floormod` when the
inner modulus can wrap over the source domain and is not divisible by
the outer modulus. Preserve valid no-wrap simplification after
right-only padding by consulting the original source extent, with
regression coverage for boundary, padding, and divisible cases.
Related downstream TVM change:
[tile-ai/tvm#65](https://github.com/tile-ai/tvm/pull/65)
Related TileLang-side change:
[tile-ai/tilelang#3065](https://github.com/tile-ai/tilelang/pull/3065)
Co-authored-by: Lei Wang <[email protected]>
---
src/arith/iter_affine_map.cc | 17 ++++++++++++++++
tests/python/arith/test_arith_iter_affine_map.py | 26 ++++++++++++++++++++++++
2 files changed, 43 insertions(+)
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index c76729c4ef..4b60cb5a1d 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -2021,6 +2021,23 @@ PrimExpr
IterMapRewriter::SplitFloorModConst(IterSplitExpr lhs, PrimExpr base, P
// We handle scale!=1 in above code, hence we only consider floormod(x, rhs)
below
// where x=floormod(floordiv(iter, lower_factor), extent) + base
+ PrimExpr source_upper_bound = lhs->source->extent;
+ auto origin_it = padded_origin_map_.find(lhs->source);
+ if (origin_it != padded_origin_map_.end()) {
+ auto padding_it = padded_iter_map_.find(origin_it->second);
+ TVM_FFI_ICHECK(padding_it != padded_iter_map_.end());
+ // Right padding only contains values excluded by padding_predicate_, so it
+ // cannot make the inner floormod wrap over the original iterator domain.
+ // Keep left-padded marks conservative because padding shifts their values.
+ if (is_zero(padding_it->second.left_pad)) {
+ source_upper_bound = origin_it->second->extent;
+ }
+ }
+ bool inner_mod_can_wrap =
+ !analyzer_->CanProve(source_upper_bound <= lhs->lower_factor *
lhs->extent);
+ if (inner_mod_can_wrap && !CanProveDivisible(lhs->extent, rhs)) {
+ return PrimExpr();
+ }
auto pair = PadDividendToDivisor(lhs, base, rhs);
IterSplitExpr padded = pair.first;
if (!padded.defined()) {
diff --git a/tests/python/arith/test_arith_iter_affine_map.py
b/tests/python/arith/test_arith_iter_affine_map.py
index 375760118b..c684117ca9 100644
--- a/tests/python/arith/test_arith_iter_affine_map.py
+++ b/tests/python/arith/test_arith_iter_affine_map.py
@@ -250,6 +250,32 @@ def test_compound_floormod_two_regression():
)
+def test_nested_floormod_requires_divisible_extents():
+ x = tvm.tirx.Var("x", "int32")
+ flm = tvm.tirx.floormod
+ non_divisible = flm(flm(x, 64), 7)
+
+ # The inner floormod does not wrap at or below its exact domain boundary.
+ assert_iter_map_simplify({non_divisible: flm(x, 7)}, var_dom([(x, 63)]))
+ assert_iter_map_simplify({non_divisible: flm(x, 7)}, var_dom([(x, 64)]))
+
+ # One value beyond the boundary makes the non-divisible inner floormod
observable.
+ assert_iter_map_simplify({non_divisible: non_divisible}, var_dom([(x,
65)]))
+ assert_iter_map_simplify({non_divisible: non_divisible}, var_dom([(x,
128)]))
+
+ # A non-zero domain minimum becomes left padding. Keep these cases
+ # conservative because padding shifts the iterator values.
+ assert_iter_map_simplify(
+ {non_divisible: non_divisible}, {x: tvm.ir.Range.from_min_extent(1,
63)}
+ )
+ assert_iter_map_simplify(
+ {non_divisible: non_divisible}, {x: tvm.ir.Range.from_min_extent(1,
64)}
+ )
+
+ divisible = flm(flm(x, 64), 8)
+ assert_iter_map_simplify({divisible: flm(x, 8)}, var_dom([(x, 128)]))
+
+
def test_predicate():
x = tvm.tirx.Var("x", "int32")
y = tvm.tirx.Var("y", "int32")
