Skip to content

[Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder - #20287

Merged
tlopex merged 2 commits into
apache:mainfrom
emecii:fix/identity-permute-dims
Sep 8, 2026
Merged

[Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder#20287
tlopex merged 2 commits into
apache:mainfrom
emecii:fix/identity-permute-dims

Conversation

@emecii

@emecii emecii commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Summary

AdjustMatmulOrder matched any relax.permute_dims around an inner matmul as though it were a transpose. An explicit identity axis order, such as [0, 1], preserves the inner result and must not trigger that transpose-specific reassociation.

This change checks explicit axes before the transpose rewrite and leaves identity permutations unchanged. It preserves the existing behavior for genuine transposes.

Fixes #20277.

Validation

  • python -m pytest tests/python/relax/test_transform_adjust_matmul_order.py -q — 27 passed with a local CPU/LLVM build.
  • The new integer LLVM regression returns [-14, 42] after AdjustMatmulOrder; the previous implementation returns [-20, 40].
  • pre-commit run --files src/relax/transform/adjust_matmul_order.cc tests/python/relax/test_transform_adjust_matmul_order.py

Generated-by: OpenAI Codex (GPT-6)

An explicit identity permute_dims does not transpose the inner matmul. Skip the transpose-specific reassociation when its axes preserve every dimension.

The focused Relax regression covers the IR distinction from a real transpose and the CPU/LLVM integer result reported in apache#20277.

Fixes apache#20277

Generated-by: OpenAI Codex (GPT-6)

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The identity fix looks good. One follow-up: a 3D permutation like axes=[1,0,2] still reaches the transpose rewrite and produces incorrect results. Could we make both branches check that only the last two axes are swapped, with batch axes unchanged, and add a regression? This also happens on the base branch.

@emecii

emecii commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Updated both rewrite branches in 27ca382 to accept only a last-two-axis swap and preserve batch axes. Added 3D LHS/RHS regressions; the focused suite is 29 passed.

@tlopex
tlopex merged commit 6583b2b into apache:main Sep 8, 2026
5 checks passed
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] AdjustMatmulOrder treats an identity permute_dims (axes=[0,1]) as a transpose and produces a wrong result

2 participants