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