Fix MatMul SmartReshape in case of 1D input (#24701)

### Details:
Handled a case when "Other" input to MatMul is 1D


### Tickets:
 - *CVS-141638*
This commit is contained in:
Ivan Tikhonov 2024-05-27 18:24:22 +04:00
parent 65c3b1783c
commit c5bd13a63c
2 changed files with 66 additions and 4 deletions

View File

@ -36,7 +36,12 @@ bool relax_hc_reshape_followed_by_matmul(const ov::pass::pattern::PatternValueMa
// avoiding loop creation
return false;
const auto idx = reshape_is_A_input ? (matmul->get_transpose_b() ? -1 : -2) : (matmul->get_transpose_a() ? -2 : -1);
bool is_1d = ov::pass::pattern::rank_equals(1)(shape_source);
int64_t idx = -1;
if (!is_1d) {
idx = reshape_is_A_input ? (matmul->get_transpose_b() ? -1 : -2) : (matmul->get_transpose_a() ? -2 : -1);
}
const auto in_C_0 = std::make_shared<ov::op::v3::ShapeOf>(shape_source);
const auto in_C_1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {idx});
const auto in_C_2 = ov::op::v0::Constant::create(ov::element::i64, {}, {0});
@ -53,7 +58,6 @@ bool relax_hc_reshape_followed_by_matmul(const ov::pass::pattern::PatternValueMa
auto reshape_input = pattern_to_output.at(reshape_label).get_node_shared_ptr()->input(1);
reshape_input.replace_source_output(new_reshape_pattern);
return true;
}
@ -61,7 +65,7 @@ bool relax_hc_reshape_followed_by_matmul(const ov::pass::pattern::PatternValueMa
ov::pass::ReshapeAMatMul::ReshapeAMatMul() {
MATCHER_SCOPE(ReshapeAMatMul);
auto other_input_label = pattern::any_input();
auto other_input_label = pattern::any_input(ov::pass::pattern::has_static_rank());
auto reshape_input_label = pattern::any_input();
auto reshape_pattern_label = pattern::any_input();
auto reshape_predicate = [](ov::Output<ov::Node> output) -> bool {
@ -86,7 +90,7 @@ ov::pass::ReshapeAMatMul::ReshapeAMatMul() {
ov::pass::ReshapeBMatMul::ReshapeBMatMul() {
MATCHER_SCOPE(ReshapeBMatMul);
auto other_input_label = pattern::any_input();
auto other_input_label = pattern::any_input(ov::pass::pattern::has_static_rank());
auto reshape_input_label = pattern::any_input();
auto reshape_pattern_label = pattern::any_input();
auto reshape_predicate = [](ov::Output<ov::Node> output) -> bool {

View File

@ -378,6 +378,64 @@ TEST_F(TransformationTestsF, SmartReshapeReshapeAMatMulSeveralConsumers) {
manager.register_pass<ov::pass::ReshapeAMatMul>();
}
TEST_F(TransformationTestsF, SmartReshapeReshapeA_1DOtherInput) {
{
auto input_to_reshape = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3, 2, 3});
auto reshape_const = ov::op::v0::Constant::create(ov::element::i32, {2}, {3, 6});
auto reshape = std::make_shared<ov::op::v1::Reshape>(input_to_reshape, reshape_const, false);
auto other_input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{6});
auto matmul = std::make_shared<ov::op::v0::MatMul>(reshape, other_input);
model = std::make_shared<ov::Model>(ov::NodeVector{matmul}, ov::ParameterVector{input_to_reshape, other_input});
manager.register_pass<ov::pass::ReshapeAMatMul>();
}
{
auto input_to_reshape = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3, 2, 3});
auto other_input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{6});
const auto in_C_0 = std::make_shared<ov::op::v3::ShapeOf>(other_input);
const auto in_C_1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
const auto in_C_2 = ov::op::v0::Constant::create(ov::element::i64, {}, {0});
const auto C = std::make_shared<ov::op::v8::Gather>(in_C_0, in_C_1, in_C_2);
const auto N = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
const auto new_reshape_pattern = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{N, C}, 0);
auto reshape = std::make_shared<ov::op::v1::Reshape>(input_to_reshape, new_reshape_pattern, false);
auto matmul = std::make_shared<ov::op::v0::MatMul>(reshape, other_input);
model_ref =
std::make_shared<ov::Model>(ov::NodeVector{matmul}, ov::ParameterVector{input_to_reshape, other_input});
}
}
TEST_F(TransformationTestsF, SmartReshapeReshapeB_1DOtherInput) {
{
auto input_to_reshape = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3, 2, 3});
auto reshape_const = ov::op::v0::Constant::create(ov::element::i32, {2}, {3, 6});
auto reshape = std::make_shared<ov::op::v1::Reshape>(input_to_reshape, reshape_const, false);
auto other_input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3});
auto matmul = std::make_shared<ov::op::v0::MatMul>(other_input, reshape);
model = std::make_shared<ov::Model>(ov::NodeVector{matmul}, ov::ParameterVector{input_to_reshape, other_input});
manager.register_pass<ov::pass::ReshapeBMatMul>();
}
{
auto input_to_reshape = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3, 2, 3});
auto other_input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{3});
const auto in_C_0 = std::make_shared<ov::op::v3::ShapeOf>(other_input);
const auto in_C_1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
const auto in_C_2 = ov::op::v0::Constant::create(ov::element::i64, {}, {0});
const auto C = std::make_shared<ov::op::v8::Gather>(in_C_0, in_C_1, in_C_2);
const auto N = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
const auto new_reshape_pattern = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{C, N}, 0);
auto reshape = std::make_shared<ov::op::v1::Reshape>(input_to_reshape, new_reshape_pattern, false);
auto matmul = std::make_shared<ov::op::v0::MatMul>(other_input, reshape);
model_ref =
std::make_shared<ov::Model>(ov::NodeVector{matmul}, ov::ParameterVector{input_to_reshape, other_input});
}
}
TEST_F(TransformationTestsF, SmartReshapeReshapeBMatMulSeveralConsumers) {
// Reshape has 2 consumers: matmul and reduce.
// Since reshape movement leads to loop creation (circular dependencies), the transformation can't be applied