diff --git a/src/common/transformations/include/transformations/symbolic_transformations/reshape_optimizations.hpp b/src/common/transformations/include/transformations/symbolic_transformations/reshape_optimizations.hpp index 8cdcaee05e7..6ad409b4d96 100644 --- a/src/common/transformations/include/transformations/symbolic_transformations/reshape_optimizations.hpp +++ b/src/common/transformations/include/transformations/symbolic_transformations/reshape_optimizations.hpp @@ -16,7 +16,39 @@ class TRANSFORMATIONS_API ReshapeOptimizations; /** * @ingroup ov_transformation_common_api * @brief Searches for Flatten-like Reshape operations and simplifies 2nd input of such Reshape using special zero - * feature + * feature. + * The transformation works in 2 cases: + * 1. all in/out dims are static, or we can match them via the symbols. + 2. only one out dim doesn't have the corresponding input static dim, + and we can't match it using symbols. Besides that the output shape must not contain zero dims, + because then value -1 in 2nd input to Reshape op can't guarantee an unambiguous determination of the remaining dim + value. + + for example: + Before: + +-------------+ +----------+ + |data | | Concat | + |shape: (0, 0)| | shape (2)|<- the values might be determined in runtime + +-----+-------+ +-----+----+ the empty data tensor can be reshaped to + | | any other empty shape, e.g. (0, 800) + +-----v---------------+ | + | Reshape | | + | shape (0,-1) <--+ + | special zero = False| + +---------------------+ + + After: + +---------------+ + +--------------+ | Constant | + | data | | shape (2) | + | shape: (0, 0)| | values (0, -1)|<- -1 means copy the corresponding input dim + +-----+--------+ +-------+-------+ + | | + +----v----------------+ | + | Reshape <----+ + | shape (0,0) | <- it might cause inconsistency + | special zero = True | + +---------------------+ */ class ov::pass::ReshapeOptimizations : public ov::pass::MatcherPass { public: diff --git a/src/common/transformations/src/transformations/symbolic_transformations/reshape_optimizations.cpp b/src/common/transformations/src/transformations/symbolic_transformations/reshape_optimizations.cpp index eee0c5d4db4..482f56367be 100644 --- a/src/common/transformations/src/transformations/symbolic_transformations/reshape_optimizations.cpp +++ b/src/common/transformations/src/transformations/symbolic_transformations/reshape_optimizations.cpp @@ -34,16 +34,23 @@ ov::pass::ReshapeOptimizations::ReshapeOptimizations() { const auto& out_shape = reshape->get_output_partial_shape(0); const auto& out_rank = out_shape.size(); + int64_t cnt_static_zeros = 0; std::vector output_pattern(out_rank, -1); for (size_t i = 0; i < out_rank; ++i) { - if (out_shape[i].is_static()) + if (out_shape[i].is_static()) { output_pattern[i] = out_shape[i].get_length(); - else if (i >= in_rank) + if (output_pattern[i] == 0) { + ++cnt_static_zeros; + } + } else if (i >= in_rank) { break; - else if (dims_are_equal(in_shape[i], out_shape[i])) + } else if (dims_are_equal(in_shape[i], out_shape[i])) { output_pattern[i] = 0; + } } - if (std::count(output_pattern.begin(), output_pattern.end(), -1) <= 1) { + + int64_t cnt_neg_ones = std::count(output_pattern.begin(), output_pattern.end(), -1); + if (cnt_neg_ones == 0 || (cnt_neg_ones == 1 && cnt_static_zeros == 0)) { auto new_pattern = ov::op::v0::Constant::create(element::i64, Shape{output_pattern.size()}, output_pattern); ov::copy_runtime_info(reshape->get_input_node_shared_ptr(1), new_pattern); reshape->set_special_zero(true); diff --git a/src/common/transformations/tests/symbolic_transformations/reshape_optimizations.cpp b/src/common/transformations/tests/symbolic_transformations/reshape_optimizations.cpp index 2f1fdba8dc2..8535abfe644 100644 --- a/src/common/transformations/tests/symbolic_transformations/reshape_optimizations.cpp +++ b/src/common/transformations/tests/symbolic_transformations/reshape_optimizations.cpp @@ -149,3 +149,26 @@ TEST_F(TransformationTestsF, NegativeTest) { manager.register_pass(); } } + +TEST_F(TransformationTestsF, ZeroDimsInOutputShape) { + // [A, B] + auto shape = PartialShape{0, 0}; + { + auto data = make_shared(element::f32, shape); + auto b = make_shared(element::f32, PartialShape{-1}); + + auto a = ov::op::v0::Constant::create(element::i64, Shape{1}, {0}); + + auto shape_of = make_shared(b); + auto indices = ov::op::v0::Constant::create(element::i64, {1}, {0}); + auto axis = ov::op::v0::Constant::create(element::i64, {}, {0}); + auto b_dim = make_shared(shape_of, indices, axis); + + auto pattern = make_shared(OutputVector{a, b_dim}, 0); + + auto reshape = make_shared(data, pattern, false); + + model = make_shared(NodeVector{reshape}, ParameterVector{data, b}); + manager.register_pass(); + } +}