apache / apache/tvm

[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases

Open
#20,205 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

[Bug][Relax] CombineParallelMatmul fails valid branches with broadcast-compatible biases

### Expected behavior

`relax.transform.CombineParallelMatmul()` should combine bias-add branches only when their biases can be concatenated along the output-channel axis. Otherwise it should leave the original valid branches
unchanged.

### Actual behavior

The pass checks that bias tensors have the same rank, but does not check non-concat dimensions. It rewrites two valid broadcast additions into an invalid bias concat.

```text
matmul(x, w1): [2, 4] + b1: [1, 4] -> [2, 4]
matmul(x, w2): [2, 5] + b2: [2, 5] -> [2, 5]

Both original additions are valid. The pass constructs concat((b1, b2), axis=1), which is invalid because the row dimensions are 1 and 2.

### Environment

OS: Linux x86_64
Target: llvm
Relax VM exec_mode: bytecode
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05

### Steps to reproduce

from tvm import relax

builder = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType((2, 3), "float32"))
w1 = relax.Var("w1", relax.TensorType((3, 4), "float32"))
w2 = relax.Var("w2", relax.TensorType((3, 5), "float32"))
b1 = relax.Var("b1", relax.TensorType((1, 4), "float32"))
b2 = relax.Var("b2", relax.TensorType((2, 5), "float32"))

with builder.function("main", [x, w1, w2, b1, b2]):
with builder.dataflow():
mm1 = builder.emit(relax.op.matmul(x, w1))
mm2 = builder.emit(relax.op.matmul(x, w2))
y1 = builder.emit(relax.op.add(mm1, b1))
y2 = builder.emit(relax.op.add(mm2, b2))
output = builder.emit_output((y1, y2))
builder.emit_func_output(output)

mod = relax.transform.Normalize()(builder.finalize())
relax.transform.CombineParallelMatmul()(mod)

Observed error:

ValueError: Concat expects the input tensors to have the same shape on every
dimension except the one indicated by the input axis. However, the input
contains tensors whose shapes on dimension 0 is T.int64(1) and T.int64(2)

Changing b2 to shape [1, 5] makes the concat valid; the original and transformed LLVM bytecode VM outputs then match exactly.

### Suspected cause

src/relax/transform/combine_parallel_matmul.cc records only a bias rank and later unconditionally builds:

auto concat_bias = concat(Tuple(bias), bias_dim - 1);

### Suggested fix

Validate compatibility of every bias shape on non-concat axes before applying the fused-bias rewrite. If that cannot be proven, combine only the matmuls and leave the individual bias additions in place.

### Triage

- needs-triage
- type: bug
- relax

```markdown

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with src/relax/transform/combine_parallel_matmul.cc and reproduce the issue using the Relax BlockBuilder example in the report. Trace how bias rank and shapes are checked before concat is built. Done means incompatible non-concat dimensions no longer produce an invalid concat: matmuls may combine, while the original bias additions remain valid and transformed outputs match.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
75/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.