zupengwang opened a new pull request, #20169:
URL: https://github.com/apache/tvm/pull/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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]