[Fix][Relax] Require matching bias output width when fusing parallel matmul biases

Author: OmarAzizi

src/relax/transform/combine_parallel_matmul.cc

@@ -234,6 +234,18 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, E
         if (!shapes_compatible_excluding_trailing_axes(bias_shapes, 1)) {
           continue;
         }
+        arith::Analyzer ana;
+        bool bias_widths_match = true;
+        for (size_t i = 0; i < splits.size(); ++i) {
+          const auto& shape = bias_shapes[i];
+          if (!ana->CanProve(shape[shape.size() - 1] == splits[i].split_size)) {
+            bias_widths_match = false;
+            break;
+          }
+        }
+        if (!bias_widths_match) {
+          continue;
+        }
       }
 
       auto concat_rhs = concat(Tuple(rhs), rhs_dim - 1);

tests/python/relax/test_transform_combine_parallel_matmul.py

@@ -779,6 +779,37 @@ def expected(
     tvm.ir.assert_structural_equal(after, expected)
 
 
+def test_skip_bias_fusion_when_bias_relies_on_its_own_broadcast():
+    """Do not fuse biases whose last dimension only matches via broadcast
+
+    Biases can agree with each other on every non-concat axis and still be
+    unsafe to fuse if a bias's own last dimension doesn't equal its
+    branch's actual output width and it was only valid by broadcasting
+    against that branch's own matmul.
+    """
+
+    @R.function(private=True)
+    def before(
+        x: R.Tensor((2, 3), "float32"),
+        w0: R.Tensor((3, 4), "float32"),
+        w1: R.Tensor((3, 5), "float32"),
+        b0: R.Tensor((2, 1), "float32"),
+        b1: R.Tensor((2, 1), "float32"),
+    ):
+        with R.dataflow():
+            lv0 = R.matmul(x, w0)
+            lv1 = R.matmul(x, w1)
+            y0 = R.add(lv0, b0)
+            y1 = R.add(lv1, b1)
+            out = (y0, y1)
+            R.output(out)
+        return out
+
+    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
+
+    tvm.ir.assert_structural_equal(after, before)
+
+
 @pytest.mark.parametrize("float32_branch", [0, 1])
 def test_skip_matmuls_with_different_output_dtypes(float32_branch):
     if float32_branch == 0: