apache / apache/tvm

[Bug][Relax] CombineParallelMatmul crashes on 1-D matmul boundary cases

Open
#20,197 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Expected behavior

relax.transform.CombineParallelMatmul should either correctly handle parallel R.matmul branches involving 1-D operands, or safely leave them unchanged if those rank combinations are outside the pass's
supported scope.

It should not construct an invalid relax.split or throw an internal IndexError on well-formed Relax IR.

### Actual behavior

There are two related failures.

For vector-LHS matmul:

x: Tensor((16,), float32)
A/B: Tensor((16, 32), float32)
matmul(x, A) + matmul(x, B)

CombineParallelMatmul tries to combine the branches, but the combined matmul output is 1-D. It then creates relax.split(axis=1), which is invalid:

ValueError:
In relax.split, the input axis 1 is out of range.
The input tensor has 1 dimensions, so axis should be in range [-1, 1).

For RHS rank 1:

x: Tensor((8, 16), float32)
A/B: Tensor((16,), float32)
matmul(x, A) + matmul(x, B)

the pass throws:

IndexError: Index 1 out of bounds 1

### Environment

OS: Linux x86_64
Python: 3.10.12
TVM version: 0.26.dev1
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05

### Steps to reproduce

import traceback

from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class VectorLHSFromExpand:
@R.function
def main(
x: R.Tensor((16,), dtype="float32"),
A: R.Tensor((16, 32), dtype="float32"),
B: R.Tensor((16, 32), dtype="float32"),
) -> R.Tensor((32,), dtype="float32"):
weight = R.add(A, B)
out = R.matmul(x, weight, out_dtype="float32")
return out

@I.ir_module
class RHSVectorBranches:
@R.function
def main(
x: R.Tensor((8, 16), dtype="float32"),
A: R.Tensor((16,), dtype="float32"),
B: R.Tensor((16,), dtype="float32"),
) -> R.Tensor((8,), dtype="float32"):
y0 = R.matmul(x, A, out_dtype="float32")
y1 = R.matmul(x, B, out_dtype="float32")
out = R.add(y0, y1)
return out

def run_case(name, mod, passes):
print("\n##", name)
current = mod

for transform_pass in passes:
pass_name = transform_pass.info.name
try:
current = transform_pass(current)
print("OK:", pass_name)
except Exception as err:
print("ERROR in", pass_name + ":", type(err).__name__ + ":", err)
traceback.print_exc(limit=2)
return

relax.analysis.well_formed(current)
print("OK: final module is well formed")

def main():
# Control: after ExpandMatmulOfSum, the IR is still well-formed and
# LegalizeOps succeeds if CombineParallelMatmul is not applied.
expanded = relax.transform.ExpandMatmulOfSum()(
relax.transform.Normalize()(VectorLHSFromExpand)
)
relax.analysis.well_formed(expanded)
relax.transform.LegalizeOps()(expanded)
print("OK: VectorLHSFromExpand is well formed after ExpandMatmulOfSum")
print("OK: LegalizeOps succeeds without CombineParallelMatmul")

run_case(
"Vector LHS through ExpandMatmulOfSum",
VectorLHSFromExpand,
[
relax.transform.Normalize(),
relax.transform.ExpandMatmulOfSum(),
relax.transform.CombineParallelMatmul(),
],
)

run_case(
"RHS vector branches",
RHSVectorBranches,
[
relax.transform.Normalize(),
relax.transform.CombineParallelMatmul(),
],
)

if __name__ == "__main__":
main()

Observed output:

OK: VectorLHSFromExpand is well formed after ExpandMatmulOfSum
OK: LegalizeOps succeeds without CombineParallelMatmul

## Vector LHS through ExpandMatmulOfSum
OK: Normalize
OK: ExpandMatmulOfSum
ERROR in CombineParallelMatmul: ValueError: In relax.split, the input axis 1 is out of range. The input tensor has 1 dimensions, so axis should be in range [-1, 1).

## RHS vector branches
OK: Normalize
ERROR in CombineParallelMatmul: IndexError: Index 1 out of bounds 1

This is triggered by explicitly running CombineParallelMatmul, or by a pass sequence that first creates parallel matmul branches with ExpandMatmulOfSum. I have not observed this under the CPU default build
pipeline.

### Triage

- needs-triage
- type: bug
- relax
- transform

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by running the supplied VectorLHSFromExpand and RHSVectorBranches reproductions through Normalize, ExpandMatmulOfSum, and CombineParallelMatmul. Trace CombineParallelMatmul at the failing rank-1 cases; done means the pass no longer constructs an invalid relax.split or raises IndexError, and the resulting module remains well formed or the cases are safely left unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.