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

[lrcyyds1]

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: 5a8dae4

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