This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 52441cdf2bb26eefc565d557f99b2c22d89b1047 Author: alumkal <[email protected]> AuthorDate: Fri Oct 16 22:41:51 2026 +0800 [FIX][Sym] Reject non-positive divisors in iter map floordiv/floormod (#20756) Fixes #20736. ### Problem `SplitFloorDivConst` and `SplitFloorModConst` cancel a non-unit iterator scale against the divisor. With a negative scale this produces a negative divisor: `floordiv(63 - x, 2)` (scale `-1`, base `63`) becomes `floordiv(x - 63, -2)`. `PadDividendToDivisor` and the split construction assume a positive divisor: - `left_pad = floormod(-63, -2) = -1`. In the first pass the mark's padding is `max(info.left_pad, -1) = 0`, so the split is not shifted, while `new_base = floordiv(base - left_pad, rhs) = 31` still assumes the shift. - The new split gets `lower_factor = -2`. As a result, over `x in [0, 64)`: | expression | `iter_map_simplify` on main | wrong values | |---|---|---| | `floordiv(63 - x, 2)` | `x // -2 + 31` | 32 / 64 | | `floormod(63 - x, 2)` | `x % -2 * -1` | 64 / 64 | | `floordiv(63 - x, 4)` | `x // -4 + 15` | 48 / 64 | | `floormod(65 - x, 4)` | `x % -4 * -1` | 64 / 64 | | `floordiv(x + 1, -2)` | `x // -2 + -1` | 32 / 64 | | `floormod(x + 1, -2)` | `x % -2` | 64 / 64 | Downstream, a buffer flattening pass that uses the iter-map simplifier turns a reversed block index `A[(63 - bx) // 2, tx]` into an index that reads the previous row for odd `bx` and row -1 (out of bounds) for `bx = 63`. ### Fix After the scale handling, report a provably non-positive divisor as not representable (via `ErrorLogger`), so these expressions are left unchanged. Negative scales that divide out exactly (e.g. `floordiv(x * -2 + 62, 2)` -> `x * -1 + 31`) take the earlier exact-division branches and are still simplified. Divisors of unknown sign are not affected, so symbolic divisors such as `n * 4096` keep simplifying as before. `DetectIterMap` callers already handle this failure: `IterMapSimplify` keeps the original expressions, `SubspaceDivide` returns empty, and region evaluation falls back to `EvalSet`. ### Testing - New `test_non_positive_divisor_is_not_simplified`: fails on main, passes with this change. - `tirx/tests/python/sym/test_iter_affine_map.py`: 46 passed. - `tirx/tests/python/sym`, `tirx/base/test_index_map.py`, `tirx/transform`, `tirx/analysis`, `tirx/test_layout.py`: 1466 passed. The 7 failures are unrelated: 6 need LLVM (not built here) and `test_vectorize_nested_predicates_preserve_both_masks` fails identically on main. - A probe over 15 negative-scale and negative-divisor expressions gives 0 wrong values after the change. --- tirx/src/sym/iter_affine_map.cc | 17 +++++++++++++++++ tirx/tests/python/sym/test_iter_affine_map.py | 22 ++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/tirx/src/sym/iter_affine_map.cc b/tirx/src/sym/iter_affine_map.cc index 8627bbf609..def26c2903 100644 --- a/tirx/src/sym/iter_affine_map.cc +++ b/tirx/src/sym/iter_affine_map.cc @@ -1993,6 +1993,16 @@ ffi::Optional<PrimExpr> IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P // We handle scale!=1 in above code, hence we only consider floordiv(x, rhs) below // where x=floormod(floordiv(iter, lower_factor), extent) + base + if (analyzer_->CanProveLess(rhs, 1)) { + // The padding and the split constructed below assume a positive divisor. + // A non-positive one can be written directly or come from dividing the + // divisor by a negative scale above, e.g. floordiv(63 - x, 2) becomes + // floordiv(x - 63, -2). + ErrorLogger(this) << "Cannot represent as IterMap: the divisor " << rhs + << " of FloorDiv is not positive."; + return std::nullopt; + } + auto pair = PadDividendToDivisor(lhs, base, rhs); if (!pair) return std::nullopt; IterSplitExpr padded = pair->first; @@ -2112,6 +2122,13 @@ ffi::Optional<PrimExpr> IterMapRewriter::SplitFloorModConst(IterSplitExpr lhs, 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 + if (analyzer_->CanProveLess(rhs, 1)) { + // See SplitFloorDivConst: the padding and the split below assume a + // positive divisor. + ErrorLogger(this) << "Cannot represent as IterMap: the divisor " << rhs + << " of FloorMod is not positive."; + return std::nullopt; + } PrimExpr source_upper_bound = lhs->source->extent; auto origin_it = padded_origin_map_.find(lhs->source); if (origin_it != padded_origin_map_.end()) { diff --git a/tirx/tests/python/sym/test_iter_affine_map.py b/tirx/tests/python/sym/test_iter_affine_map.py index 44ccc72d40..03b675bd44 100644 --- a/tirx/tests/python/sym/test_iter_affine_map.py +++ b/tirx/tests/python/sym/test_iter_affine_map.py @@ -250,6 +250,28 @@ def test_compound_floormod_two_regression(): ) +def test_non_positive_divisor_is_not_simplified(): + # Dividing out a negative scale turns floordiv(63 - x, 2) into + # floordiv(x - 63, -2); padding with a negative divisor used to drop the + # remainder (e.g. x // -2 + 31, wrong for every odd x). Such expressions + # are left unchanged, as are directly written negative divisors. + x = tvm_tirx.tirx.Var("x", "int32") + fld = tvm_tirx.tirx.floordiv + flm = tvm_tirx.tirx.floormod + for index in [ + fld(63 - x, 2), + flm(63 - x, 2), + fld(63 - x, 4), + flm(65 - x, 4), + fld(x + 1, -2), + flm(x + 1, -2), + ]: + assert_iter_map_simplify({index: index}, var_dom([(x, 64)])) + + # A negative scale that divides out exactly is still simplified. + assert_iter_map_simplify({fld(x * -2 + 62, 2): x * -1 + 31}, var_dom([(x, 64)])) + + def test_nested_floormod_requires_divisible_extents(): x = tvm_tirx.tirx.Var("x", "int32") flm = tvm_tirx.tirx.floormod
