Skip to content

[Relax] Preserve out_dtype in AdjustMatmulOrder - #20296

Merged
tlopex merged 1 commit into
apache:mainfrom
StrongbodyStrongmind:fix-20202
Sep 10, 2026
Merged

[Relax] Preserve out_dtype in AdjustMatmulOrder#20296
tlopex merged 1 commit into
apache:mainfrom
StrongbodyStrongmind:fix-20202

Conversation

@StrongbodyStrongmind

@StrongbodyStrongmind StrongbodyStrongmind commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Fixes #20202.

AdjustMatmulOrder rebuilt reassociated outer matmuls with
out_dtype=void, which could change the result dtype of the original
program. For example, a float16 matmul chain whose outer matmul explicitly
returns float32 was rewritten to return float16.

This patch preserves the original outer MatmulAttrs::out_dtype when
constructing the replacement outer matmul. Newly introduced inner matmuls
continue to infer their output dtype.

Tests:

  • Added regression coverage for both reassociation directions.
  • Covered explicit and inferred output dtypes.
  • 12 new regression cases passed.
  • Full AdjustMatmulOrder test file: 41 passed.
  • pre-commit and git diff --check: passed.

Preserve the original outer matmul output dtype in both cost-based and compile-time grouping rewrites. Add regression coverage for both reassociation directions, transposed inner matmuls, inferred output dtype, and idempotence.

Fixes apache#20202.

Validation: 41 AdjustMatmulOrder tests passed; LLVM VM float32 output checks passed in both directions; pre-commit passed on changed files. Before the fix, all six explicit float32 regression cases failed.

@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.

Thanks for the fix! LGTM

@tlopex
tlopex merged commit 6b34c75 into apache:main Sep 10, 2026
8 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 drops the outer matmul out_dtype

2 participants