This is an automated email from the ASF dual-hosted git repository.
tlopex 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 c46018d370 [BugFix][Arith] Reject padded IterMapSimplify fallback
(#20169)
c46018d370 is described below
commit c46018d37018d547078fc1dcdfda5eebfa5916a8
Author: Zupeng Wang <[email protected]>
AuthorDate: Wed Sep 2 08:52:52 2026 +0800
[BugFix][Arith] Reject padded IterMapSimplify fallback (#20169)
Fixes #19524.
## Problem
When predicate-aware iter-map detection fails, `IterMapSimplify` retries
detection without the predicate. The retry can introduce iterator
padding, but the API returns only simplified expressions and drops the
fallback's `padding_predicate`.
`FlattenBuffer` can therefore consume a non-equivalent flattened index.
In the `conv2d_transpose` reproducer from #19524, scheduled TIR contains
`(index - 1) // 2`, while the affected path generated a shifted CUDA
address and silently read out of bounds.
## Change
- Accept the predicate-free fallback only when it requires no iterator
padding.
- Preserve the original expression when the fallback would need a
padding predicate.
- Cover both sides of the decision: an unpadded fallback still
simplifies, while a padded fallback is rejected.
## Validation
Base: `apache/tvm@4e9a099d154d7c4644a40a1a9c00b8873226468e`
Environment: NVIDIA RTX 3090 (SM86), CUDA 13.0 (`nvcc 13.0.48`), GCC
11.5.0, CMake 4.4.2, TVM `0.26.dev0`.
- CUDA-enabled build: `589/589` targets built.
- `python -m pytest tests/python/arith/test_arith_iter_affine_map.py
-q`: `44 passed`.
- `python -m pytest
tests/python/relax/test_transform_legalize_ops_nn.py::test_conv2d_transpose
tests/python/relax/test_op_nn_convolution.py::test_conv2d_transpose_wrong_output_padding
-q`: `2 passed`.
- Scoped pre-commit checks on both changed files: all passed, including
Ruff and clang-format 20.1.8.
- `git diff --check`: passed.
GPU correctness experiment used 14 `conv2d_transpose` cases (`H=4..10`,
`output_padding in {0, 1}`), with PyTorch `F.conv_transpose2d` as the
oracle:
| Revision | Failing cases | Worst max absolute error |
| --- | ---: | ---: |
| Base | 5 / 14 | 0.3841745257 |
| This change | 0 / 14 | 7.4505806e-08 |
## Risk
For predicates that iter-map detection cannot parse, padded fallback
mappings now retain their original expressions instead of being
simplified. This is deliberately conservative and may reduce
simplification in those cases; unpadded fallbacks keep the existing
behavior.
## Not run locally
The full upstream arm, cpu, docker, gpu, and wasm CI matrices were not
run locally.
Co-authored-by: Wang <Zupeng>
---
src/arith/iter_affine_map.cc | 14 ++++++++----
tests/python/arith/test_arith_iter_affine_map.py | 27 ++++++++++++++++++++++++
2 files changed, 37 insertions(+), 4 deletions(-)
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 6537598cb9..2a766dccf8 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -2192,10 +2192,16 @@ ffi::Array<PrimExpr> IterMapSimplify(const
ffi::Array<PrimExpr>& 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 c684117ca9..6a5b809ac4 100644
--- a/tests/python/arith/test_arith_iter_affine_map.py
+++ b/tests/python/arith/test_arith_iter_affine_map.py
@@ -1260,6 +1260,33 @@ def test_iter_map_simplify_symbolic_predicate():
)
+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")
