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 149f7515e4 [FIX][Sym] Reject non-positive divisors in iter map
floordiv/floormod (#20756)
149f7515e4 is described below
commit 149f7515e45958a4344818a645ea542e558a2619
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
