diff --git a/src/relax/transform/combine_parallel_matmul.cc b/src/relax/transform/combine_parallel_matmul.cc index 952b7f4b6209..546e2539c8e0 100644 --- a/src/relax/transform/combine_parallel_matmul.cc +++ b/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, ffi::Map)> GetRewriter( const Patterns& patterns, const BranchInfo& branch_info, FCheck check) { - auto batch_dims_compatible = [](size_t rhs_dim, const std::vector& indices, - const std::vector>& rhs_shapes) { - arith::Analyzer ana; - for (auto ind : indices) { - TVM_FFI_ICHECK_EQ(static_cast(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>& 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& indices, + const std::vector>& rhs_shapes) { + std::vector> 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 matchings, ffi::Map bindings) { @@ -145,7 +152,7 @@ ffi::TypedFunction(ffi::Map, ffi::Map 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, ffi::Map> 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{}; + } + 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); diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py b/tests/python/relax/test_transform_combine_parallel_matmul.py index 57be71ae8806..fce2e756af62 100644 --- a/tests/python/relax/test_transform_combine_parallel_matmul.py +++ b/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: