[Fix][Relax] Validate non-concat axes before fusing parallel matmul biases

src/relax/transform/combine_parallel_matmul.cc

@@ -117,19 +117,26 @@ Patterns CreatePatterns(const BranchInfo& branch_info) {
 /*! \brief Create a rewriter for the given parallel matmul branches. */
 ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, Expr>)> GetRewriter(
     const Patterns& patterns, const BranchInfo& branch_info, FCheck check) {
-  auto batch_dims_compatible = [](size_t rhs_dim, const std::vector<size_t>& indices,
-                                  const std::vector<ffi::Array<PrimExpr>>& rhs_shapes) {
-    arith::Analyzer ana;
-    for (auto ind : indices) {
-      TVM_FFI_ICHECK_EQ(static_cast<int>(rhs_shapes[ind].size()), rhs_dim);
-      // -2 for reduction and concat axes
-      for (size_t i = 0; i < rhs_dim - 2; ++i) {
-        if (!ana->CanProve(rhs_shapes[indices[0]][i] == rhs_shapes[ind][i])) {
-          return false;
+  auto shapes_compatible_excluding_trailing_axes =
+      [](const std::vector<ffi::Array<PrimExpr>>& shapes, size_t num_trailing_axes_excluded) {
+        arith::Analyzer ana;
+        size_t ndim = shapes[0].size();
+        for (const auto& shape : shapes) {
+          TVM_FFI_ICHECK_EQ(shape.size(), ndim);
+          for (size_t i = 0; i < ndim - num_trailing_axes_excluded; ++i) {
+            if (!ana->CanProve(shapes[0][i] == shape[i])) {
+              return false;
+            }
+          }
         }
-      }
-    }
-    return true;
+        return true;
+      };
+  auto batch_dims_compatible = [&](const std::vector<size_t>& indices,
+                                   const std::vector<ffi::Array<PrimExpr>>& rhs_shapes) {
+    std::vector<ffi::Array<PrimExpr>> selected;
+    selected.reserve(indices.size());
+    for (size_t ind : indices) selected.push_back(rhs_shapes[ind]);
+    return shapes_compatible_excluding_trailing_axes(selected, 2);
   };
 
   return [=](ffi::Map<DFPattern, Var> matchings, ffi::Map<Var, Expr> bindings) {
@@ -145,7 +152,7 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, E
     ffi::Map<Var, Expr> replacements;
 
     for (const auto& [rhs_dim, indices] : GroupShapes(rhs_shapes)) {
-      if (indices.size() == 1 || !batch_dims_compatible(rhs_dim, indices, rhs_shapes)) continue;
+      if (indices.size() == 1 || !batch_dims_compatible(indices, rhs_shapes)) continue;
 
       auto lhs = matchings[patterns.input];
 
@@ -210,6 +217,37 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var, E
         continue;
       }
 
+      if (branch_info.bias_dim) {
+        std::vector<ffi::Array<PrimExpr>> bias_shapes;
+        bool bias_shape_unknown = false;
+        for (const auto& bias_var : bias) {
+          auto bias_shape_opt = GetTensorType(bias_var)->GetShape();
+          if (!bias_shape_opt) {
+            bias_shape_unknown = true;
+            break;
+          }
+          bias_shapes.push_back(bias_shape_opt.value());
+        }
+        if (bias_shape_unknown) {
+          return ffi::Map<Var, Expr>{};
+        }
+        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);
       auto matmul_combined = matmul(lhs, concat_rhs, splits[0].out_dtype);
 

tests/python/relax/test_transform_combine_parallel_matmul.py

@@ -696,6 +696,120 @@ def expected(
     tvm.ir.assert_structural_equal(after, expected)
 
 
+def test_skip_bias_fusion_with_incompatible_bias_shapes():
+    """Do not fuse biases whose shapes are broadcast-compatible with their
+    own matmul but not concat-compatible with each other
+
+    Regression test for https://github.com/apache/tvm/issues/20205.  Both
+    branches are individually valid: `b0` broadcasts against `lv0`'s [2, 4]
+    output, and `b1` already matches `lv1`'s [2, 5] output exactly.  But
+    `b0` and `b1` disagree on dimension 0 (1 vs 2).
+    """
+
+    @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((1, 4), "float32"),
+        b1: R.Tensor((2, 5), "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)
+
+
+def test_fuse_bias_with_compatible_non_concat_axes():
+    """Biases with the same rank and matching non-concat axes still fuse
+
+    Companion to `test_skip_bias_fusion_with_incompatible_bias_shapes`:
+    here `b0` and `b1` both have a leading dimension of 1, so they agree
+    on every axis except the one being concatenated, and the bias-fusing
+    shortcut remains valid.
+    """
+
+    @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((1, 4), "float32"),
+        b1: R.Tensor((1, 5), "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
+
+    @R.function(private=True)
+    def expected(
+        x: R.Tensor((2, 3), dtype="float32"),
+        w0: R.Tensor((3, 4), dtype="float32"),
+        w1: R.Tensor((3, 5), dtype="float32"),
+        b0: R.Tensor((1, 4), dtype="float32"),
+        b1: R.Tensor((1, 5), dtype="float32"),
+    ):
+        with R.dataflow():
+            lv = R.concat((w0, w1), axis=1)
+            lv1 = R.matmul(x, lv, out_dtype="float32")
+            lv2 = R.concat((b0, b1), axis=1)
+            lv3 = R.add(lv1, lv2)
+            lv4 = R.split(lv3, indices_or_sections=[4], axis=1)
+            lv0 = lv4[0]
+            lv1_1 = lv4[1]
+            out = (lv0, lv1_1)
+            R.output(out)
+        return out
+
+    after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
+
+    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: