[Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder - #20287
Merged
Conversation
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
requested changes
Sep 8, 2026
tlopex
left a comment
Member
There was a problem hiding this comment.
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.
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
approved these changes
Sep 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AdjustMatmulOrdermatched anyrelax.permute_dimsaround 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.[-14, 42]afterAdjustMatmulOrder; 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.pyGenerated-by: OpenAI Codex (GPT-6)