diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc index c76729c4ef26..519a05412b73 100644 --- a/src/arith/iter_affine_map.cc +++ b/src/arith/iter_affine_map.cc @@ -2175,10 +2175,16 @@ ffi::Array IterMapSimplify(const ffi::Array& indices, if (rewrite.empty() && !is_one(input_pred) && check_level != IterMapLevel::Bijective) { // The input predicate may cause detect iter map to fail // but we can still detect the iter map without the input predicate - // in which case the resulting iter map is valid and can be used for simplification. - rewrite = DetectIterMap(indices, input_iters, IntImm::Bool(true), check_level, ana, - /*simplify_trivial_iterators=*/simplify_trivial_iterators) - ->indices; + // in which case an unpadded iter map is valid and can be used for + // simplification. + auto fallback = DetectIterMap(indices, input_iters, IntImm::Bool(true), check_level, ana, + /*simplify_trivial_iterators=*/simplify_trivial_iterators); + // A padded fallback is not equivalent over the original iterator domain unless its + // padding predicate is also preserved. IterMapSimplify only returns expressions, so it + // cannot carry that predicate to callers. + if (!fallback->indices.empty() && is_zero(fallback->padding_predicate)) { + rewrite = fallback->indices; + } } if (rewrite.empty()) { diff --git a/tests/python/arith/test_arith_iter_affine_map.py b/tests/python/arith/test_arith_iter_affine_map.py index 375760118b5f..e3cb75749ec6 100644 --- a/tests/python/arith/test_arith_iter_affine_map.py +++ b/tests/python/arith/test_arith_iter_affine_map.py @@ -1234,6 +1234,33 @@ def fsymbolic_fuse2(i): ) +def test_iter_map_simplify_predicate_fallback_requires_no_padding(): + fused = tvm.tirx.Var("fused", "int64") + predicate = fused % 2 == 0 + unpadded_index = fused // 4 * 4 + fused % 4 + simplified = tvm.arith.iter_map_simplify( + [unpadded_index], + var_dom([(fused, 1024)]), + predicate=predicate, + ) + tvm.ir.assert_structural_equal(simplified, [fused]) + + kernel = tvm.tirx.Var("kernel", "int64") + value = fused % 14 + kernel + index = (value - 1) // 2 + predicate = (value + 1) % 2 == 0 + + # The parity predicate is not a bound constraint, so IterMapSimplify falls back to + # detecting the map without it. That fallback requires left-padding the iterator; + # discarding the corresponding padding predicate would change the index expression. + simplified = tvm.arith.iter_map_simplify( + [index], + var_dom([(fused, 1024), (kernel, 3)]), + predicate=predicate, + ) + tvm.ir.assert_structural_equal(simplified, [index]) + + def test_iter_map_simplify_symbolic_reshape(): n = tvm.tirx.Var("n", "int64") fused = tvm.tirx.Var("fused", "int64")