Skip to content

[Relax] Skip 1-D inputs in CombineParallelMatmul - #20389

Open
Nanmur wants to merge 1 commit into
apache:mainfrom
Nanmur:codex/fix-combine-parallel-matmul-1d
Open

Nanmur wants to merge 1 commit into
apache:mainfrom
Nanmur:codex/fix-combine-parallel-matmul-1d

Conversation

@Nanmur

@Nanmur Nanmur commented Sep 19, 2026

Copy link
Copy Markdown
Contributor

Motivation

CombineParallelMatmul assumes both operands have matrix dimensions when it concatenates RHS tensors and splits the combined output. Valid matmul expressions with a 1-D operand therefore crash the pass:

  • A vector LHS produces a 1-D output, but the rewrite constructs relax.split(axis=1).
  • A vector RHS makes the batch-dimension loop use rhs_dim - 2, which underflows and leads to an out-of-bounds shape access.

Changes

  • Skip rewrite groups whose RHS rank is less than two.
  • Skip rewrite groups whose shared LHS rank is less than two.
  • Add regression tests for both vector-LHS and vector-RHS matmul branches, verifying that the pass leaves valid unsupported patterns unchanged.

Testing

  • python -m pytest tests/python/relax/test_transform_combine_parallel_matmul.py -q (14 passed)
  • Release/LLVM source build on Windows
  • Ruff check and format check
  • clang-format dry run
  • ASF header and file-type checks

Fixes #20197

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] CombineParallelMatmul crashes on 1-D matmul boundary cases

1 participant