diff --git a/src/common/transformations/tests/op_conversions/batch_norm_decompositoin.cpp b/src/common/transformations/tests/op_conversions/batch_norm_decomposition_test.cpp similarity index 99% rename from src/common/transformations/tests/op_conversions/batch_norm_decompositoin.cpp rename to src/common/transformations/tests/op_conversions/batch_norm_decomposition_test.cpp index 432647f3ef4..595975082e6 100644 --- a/src/common/transformations/tests/op_conversions/batch_norm_decompositoin.cpp +++ b/src/common/transformations/tests/op_conversions/batch_norm_decomposition_test.cpp @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 // +#include "transformations/op_conversions/batch_norm_decomposition.hpp" + #include #include @@ -12,7 +14,6 @@ #include "openvino/opsets/opset1.hpp" #include "openvino/opsets/opset5.hpp" #include "transformations/init_node_info.hpp" -#include "transformations/op_conversions/batch_norm_decomposition.hpp" #include "transformations/utils/utils.hpp" using namespace ov; using namespace testing; diff --git a/src/common/transformations/tests/op_conversions/scaled_dot_product_decomposition_test.cpp b/src/common/transformations/tests/op_conversions/scaled_dot_product_decomposition_test.cpp new file mode 100644 index 00000000000..fec2fa3472b --- /dev/null +++ b/src/common/transformations/tests/op_conversions/scaled_dot_product_decomposition_test.cpp @@ -0,0 +1,227 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include "common_test_utils/ov_test_utils.hpp" +#include "openvino/core/model.hpp" +#include "openvino/op/add.hpp" +#include "openvino/op/broadcast.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert_like.hpp" +#include "openvino/op/divide.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/greater_eq.hpp" +#include "openvino/op/logical_not.hpp" +#include "openvino/op/matmul.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/range.hpp" +#include "openvino/op/scaled_dot_product_attention.hpp" +#include "openvino/op/select.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/softmax.hpp" +#include "openvino/op/sqrt.hpp" +#include "openvino/op/squeeze.hpp" +#include "openvino/op/transpose.hpp" +#include "openvino/op/unsqueeze.hpp" +#include "transformations/op_conversions/scaled_dot_product_attention_decomposition.hpp" + +using namespace ov; +using namespace testing; + +const std::shared_ptr scaled_dot_product_attention_decomposition( + const std::shared_ptr query, + const std::shared_ptr key, + const std::shared_ptr value, + const std::shared_ptr attention_mask, + const std::shared_ptr scale, + const bool casual); + +TEST_F(TransformationTestsF, ScaledDotProductAttentionDecompositionStaticBasic) { + const PartialShape query_shape{1, 32, 32}; + const PartialShape key_shape{1, 32, 32}; + const PartialShape value_shape{1, 32, 32}; + const PartialShape attention_mask_shape{1, 32, 32}; + const PartialShape scale_shape{1}; + + const auto query = std::make_shared(element::f32, query_shape); + const auto key = std::make_shared(element::f32, key_shape); + const auto value = std::make_shared(element::f32, value_shape); + const auto attention_mask = std::make_shared(element::f32, attention_mask_shape); + const auto scale = std::make_shared(element::f32, scale_shape); + const auto casual = false; + { + const auto scaled_dot_product_attention = + std::make_shared(query, key, value, attention_mask, scale, casual); + + model = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + manager.register_pass(); + } + + { + const auto scaled_dot_product_attention = + scaled_dot_product_attention_decomposition(query, key, value, attention_mask, scale, casual); + model_ref = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + } +} + +TEST_F(TransformationTestsF, ScaledDotProductAttentionDecompositionStaticBroadcastKey) { + const PartialShape query_shape{2, 1, 3, 1}; + const PartialShape key_shape{1, 7, 4, 1}; + const PartialShape value_shape{2, 7, 4, 1}; + const PartialShape attention_mask_shape{1, 1, 1, 4}; + const PartialShape scale_shape{1}; + + const auto query = std::make_shared(element::f32, query_shape); + const auto key = std::make_shared(element::f32, key_shape); + const auto value = std::make_shared(element::f32, value_shape); + const auto attention_mask = std::make_shared(element::f32, attention_mask_shape); + const auto scale = std::make_shared(element::f32, scale_shape); + const auto casual = false; + { + const auto scaled_dot_product_attention = + std::make_shared(query, key, value, attention_mask, scale, casual); + + model = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + manager.register_pass(); + } + + { + const auto scaled_dot_product_attention = + scaled_dot_product_attention_decomposition(query, key, value, attention_mask, scale, casual); + model_ref = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + } +} + +TEST_F(TransformationTestsF, ScaledDotProductAttentionDecompositionStaticBroadcastQuery) { + const PartialShape query_shape{1, 7, 1, 4}; + const PartialShape key_shape{2, 1, 1, 4}; + const PartialShape value_shape{1, 1, 1, 6}; + const PartialShape attention_mask_shape{1, 1, 3, 1}; + const PartialShape scale_shape{}; + + const auto query = std::make_shared(element::f32, query_shape); + const auto key = std::make_shared(element::f32, key_shape); + const auto value = std::make_shared(element::f32, value_shape); + const auto attention_mask = std::make_shared(element::f32, attention_mask_shape); + const auto scale = std::make_shared(element::f32, scale_shape); + const auto casual = false; + { + const auto scaled_dot_product_attention = + std::make_shared(query, key, value, attention_mask, scale, casual); + + model = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + manager.register_pass(); + } + + { + const auto scaled_dot_product_attention = + scaled_dot_product_attention_decomposition(query, key, value, attention_mask, scale, casual); + model_ref = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + } +} + +TEST_F(TransformationTestsF, ScaledDotProductAttentionDecompositionDynamic) { + const PartialShape query_shape{-1, -1, -1}; + const PartialShape key_shape{-1, -1, -1}; + const PartialShape value_shape{-1, -1, -1}; + const PartialShape attention_mask_shape{-1, -1, -1}; + const PartialShape scale_shape{}; + + const auto query = std::make_shared(element::f32, query_shape); + const auto key = std::make_shared(element::f32, key_shape); + const auto value = std::make_shared(element::f32, value_shape); + const auto attention_mask = std::make_shared(element::f32, attention_mask_shape); + const auto scale = std::make_shared(element::f32, scale_shape); + const auto casual = false; + { + const auto scaled_dot_product_attention = + std::make_shared(query, key, value, attention_mask, scale, casual); + + model = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + manager.register_pass(); + } + + { + const auto scaled_dot_product_attention = + scaled_dot_product_attention_decomposition(query, key, value, attention_mask, scale, casual); + model_ref = std::make_shared(NodeVector{scaled_dot_product_attention}, + ParameterVector{query, key, value, attention_mask, scale}); + } +} + +const std::shared_ptr scaled_dot_product_attention_decomposition( + const std::shared_ptr query, + const std::shared_ptr key, + const std::shared_ptr value, + const std::shared_ptr attention_mask, + const std::shared_ptr scale, + const bool casual) { + const auto q_shape = std::make_shared(query, element::i32); + const auto k_shape = std::make_shared(key, element::i32); + const auto minus_one = ov::op::v0::Constant::create(element::i32, Shape{}, {-1}); + const auto minus_two = ov::op::v0::Constant::create(element::i32, Shape{}, {-2}); + const auto zero_i = ov::op::v0::Constant::create(element::i32, Shape{}, {0}); + const auto one_i = ov::op::v0::Constant::create(element::i32, Shape{}, {1}); + const auto one_f = std::make_shared(one_i, query); + const auto zero_f = std::make_shared(zero_i, query); + + const auto q_scaled = std::make_shared(query, scale); + auto k_rank = std::make_shared(k_shape, element::i32)->output(0); + const auto k_last_dim = std::make_shared(k_rank, minus_one); + const auto k_next_dim = std::make_shared(k_rank, minus_two)->output(0); + k_rank = std::make_shared(k_rank, zero_i); + auto minus_inf = + ov::op::v0::Constant::create(element::f32, Shape{}, {-std::numeric_limits::infinity()})->output(0); + const auto keep_dim_last = std::make_shared(k_next_dim, zero_i); + const auto k_dims_before_transpose = + std::make_shared(zero_i, keep_dim_last, one_i, element::i32); + + const auto transpose_dims = + std::make_shared(OutputVector{k_dims_before_transpose, k_last_dim, k_next_dim}, 0); + const auto k_transposed = std::make_shared(key, transpose_dims); + auto scaled_atten = std::make_shared(q_scaled, k_transposed)->output(0); + minus_inf = std::make_shared(minus_inf, scaled_atten); + + Output mask; + Output atten_mask; + if (!casual) { + mask = attention_mask; + if (mask.get_element_type() == element::boolean) { + atten_mask = std::make_shared(mask, scaled_atten); + const auto inv_mask = std::make_shared(mask); + atten_mask = std::make_shared(inv_mask, atten_mask, minus_inf); + } else { + atten_mask = mask; + } + } else { + const auto target_s_len = std::make_shared(q_shape, minus_two, zero_i); + const auto source_s_len = std::make_shared(k_shape, minus_two, zero_i); + const auto ssl = std::make_shared(source_s_len, zero_i); + const auto tsl = std::make_shared(target_s_len, zero_i); + const auto mask_shape = std::make_shared(OutputVector{tsl, ssl}, 0); + mask = std::make_shared(minus_inf, mask_shape); + auto horizontal_range = + std::make_shared(zero_i, source_s_len, one_i, element::i32)->output(0); + horizontal_range = std::make_shared(horizontal_range, zero_i); + const auto stop = std::make_shared(target_s_len, one_i); + auto vertical_range = std::make_shared(one_i, stop, one_i, element::i32)->output(0); + vertical_range = std::make_shared(vertical_range, one_i); + const auto triu = std::make_shared(horizontal_range, vertical_range); + atten_mask = std::make_shared(triu, mask, zero_f); + } + scaled_atten = std::make_shared(scaled_atten, atten_mask); + scaled_atten = std::make_shared(scaled_atten, -1); + + const std::shared_ptr result = std::make_shared(scaled_atten, value); + return result; +} diff --git a/src/core/tests/visitors/op/scaled_dot_product_attention.cpp b/src/core/tests/visitors/op/scaled_dot_product_attention.cpp new file mode 100644 index 00000000000..b4cfa1e224f --- /dev/null +++ b/src/core/tests/visitors/op/scaled_dot_product_attention.cpp @@ -0,0 +1,29 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/scaled_dot_product_attention.hpp" + +#include + +#include "openvino/op/unique.hpp" +#include "visitors/visitors.hpp" + +using namespace ov; +using ov::test::NodeBuilder; + +TEST(attributes, scaled_dot_product_attention) { + NodeBuilder::get_ops().register_factory(); + const auto query = std::make_shared(element::f32, Shape{1, 32, 32}); + const auto key = std::make_shared(element::f32, Shape{1, 32, 32}); + const auto value = std::make_shared(element::f32, Shape{1, 32, 32}); + const auto casual = false; + + const auto op = std::make_shared(query, key, value, casual); + NodeBuilder builder(op, {query, key, value}); + auto g_sdpa = ov::as_type_ptr(builder.create()); + + constexpr auto expected_attr_count = 1; + EXPECT_EQ(builder.get_value_map_size(), expected_attr_count); + EXPECT_EQ(op->get_causal(), g_sdpa->get_causal()); +}