Skip to content

[BugFix][Arith] Reject padded IterMapSimplify fallback - #20169

Open
zupengwang wants to merge 1 commit into
apache:mainfrom
zupengwang:fix/conv2d-transpose-output-padding-19524
Open

[BugFix][Arith] Reject padded IterMapSimplify fallback#20169
zupengwang wants to merge 1 commit into
apache:mainfrom
zupengwang:fix/conv2d-transpose-output-padding-19524

Conversation

@zupengwang

Copy link
Copy Markdown

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.

@zupengwang
zupengwang marked this pull request as ready for review August 24, 2026 02:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug][Relax] conv2d_transpose produces wrong results on CUDA when output_padding > 0

1 participant