diff --git a/src/core/include/openvino/op/ops.hpp b/src/core/include/openvino/op/ops.hpp index 1df29258cb8928..adeb9c25611960 100644 --- a/src/core/include/openvino/op/ops.hpp +++ b/src/core/include/openvino/op/ops.hpp @@ -172,6 +172,7 @@ #include "openvino/op/scatter_nd_update.hpp" #include "openvino/op/scatter_update.hpp" #include "openvino/op/search_sorted.hpp" +#include "openvino/op/segment_max.hpp" #include "openvino/op/select.hpp" #include "openvino/op/selu.hpp" #include "openvino/op/shape_of.hpp" diff --git a/src/core/include/openvino/op/segment_max.hpp b/src/core/include/openvino/op/segment_max.hpp new file mode 100644 index 00000000000000..944cdb4ded97fe --- /dev/null +++ b/src/core/include/openvino/op/segment_max.hpp @@ -0,0 +1,46 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "openvino/op/op.hpp" + +namespace ov::op::v16 { +/// \brief An operation which computes the maximum values along segments of a tensor. +/// \ingroup ov_ops_cpp_api +class OPENVINO_API SegmentMax : public ov::op::Op { +public: + OPENVINO_OP("SegmentMax", "opset16", ov::op::Op); + + SegmentMax() = default; + + /// \brief Constructs a SegmentMax operation. + /// + /// \param data Input tensor with data + /// \param segment_ids Indices of segments in the data input tensor + /// \param fill_mode The value assigned to segments which are empty + SegmentMax(const Output& data, const Output& segment_ids, const op::FillMode fill_mode); + + /// \brief Constructs a SegmentMax operation. + /// + /// \param data Input tensor with data + /// \param segment_ids Indices of segments in the data input tensor + /// \param num_segments The segments count + /// \param fill_mode The value assigned to segments which are empty + SegmentMax(const Output& data, + const Output& segment_ids, + const Output& num_segments, + const op::FillMode fill_mode); + + bool visit_attributes(AttributeVisitor& visitor) override; + void validate_and_infer_types() override; + std::shared_ptr clone_with_new_inputs(const OutputVector& new_args) const override; + + const op::FillMode get_fill_mode() const; + +private: + op::FillMode m_fill_mode{}; +}; + +} // namespace ov::op::v16 diff --git a/src/core/include/openvino/op/util/attr_types.hpp b/src/core/include/openvino/op/util/attr_types.hpp index 5f8895d62fe9b1..5039cf25411dd7 100644 --- a/src/core/include/openvino/op/util/attr_types.hpp +++ b/src/core/include/openvino/op/util/attr_types.hpp @@ -20,6 +20,12 @@ enum class PadMode { CONSTANT = 0, EDGE, REFLECT, SYMMETRIC }; OPENVINO_API std::ostream& operator<<(std::ostream& s, const PadMode& type); +/// \brief Fill modes for the `SegmentMax` operator. +enum class FillMode { ZERO = 0, LOWEST }; + +OPENVINO_API +std::ostream& operator<<(std::ostream& s, const FillMode& type); + /// \brief Padding Type used for `Convolution` and `Pooling` /// /// Follows ONNX padding type definitions @@ -219,6 +225,14 @@ class OPENVINO_API AttributeAdapter : public EnumAttributeAdapterBa OPENVINO_RTTI("AttributeAdapter"); }; +template <> +class OPENVINO_API AttributeAdapter : public EnumAttributeAdapterBase { +public: + AttributeAdapter(op::FillMode& value) : EnumAttributeAdapterBase(value) {} + + OPENVINO_RTTI("AttributeAdapter"); +}; + template <> class OPENVINO_API AttributeAdapter : public EnumAttributeAdapterBase { public: diff --git a/src/core/shape_inference/include/segment_max_shape_inference.hpp b/src/core/shape_inference/include/segment_max_shape_inference.hpp new file mode 100644 index 00000000000000..fac0c4b44599f9 --- /dev/null +++ b/src/core/shape_inference/include/segment_max_shape_inference.hpp @@ -0,0 +1,82 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include + +#include "openvino/op/segment_max.hpp" +#include "utils.hpp" + +namespace ov { +namespace op { +namespace v16 { +template > +std::vector shape_infer(const SegmentMax* op, + const std::vector& input_shapes, + const ITensorAccessor& tensor_accessor = make_tensor_accessor()) { + NODE_VALIDATION_CHECK(op, input_shapes.size() == 2 || input_shapes.size() == 3); + + // validate shape of data input + const auto& data_shape = input_shapes[0]; + const auto is_data_shape_rank_static = data_shape.rank().is_static(); + if (is_data_shape_rank_static) { + NODE_SHAPE_INFER_CHECK(op, input_shapes, data_shape.size() > 0, "The data input cannot be a scalar."); + } + + // validate segment_ids input + const auto& segment_ids_shape = input_shapes[1]; + const auto is_segment_ids_rank_static = segment_ids_shape.rank().is_static(); + if (is_segment_ids_rank_static) { + NODE_SHAPE_INFER_CHECK(op, + input_shapes, + segment_ids_shape.size() == 1, + "segment_ids must be a 1D input. Got: ", + segment_ids_shape); + if (is_data_shape_rank_static) { + NODE_SHAPE_INFER_CHECK(op, + input_shapes, + data_shape[0].compatible(segment_ids_shape[0]), + "The number of elements in segment_ids must match the first dimension of data."); + } + } + const auto segment_ids = ov::op::get_input_const_data_as(op, 1, tensor_accessor); + if (segment_ids) { + NODE_VALIDATION_CHECK(op, + std::is_sorted(segment_ids->begin(), segment_ids->end()), + "segment_ids must be sorted."); + } + + // validate num_segments input + const auto num_segments_available = op->inputs().size() == 3; + const auto num_segments = num_segments_available ? get_input_const_data_as_shape(op, 2, tensor_accessor) + : ov::optional{}; + if (num_segments_available) { + const auto& num_segments_shape = input_shapes[2]; + NODE_SHAPE_INFER_CHECK(op, + input_shapes, + num_segments_shape.rank().compatible(0), + "num_segments must be a scalar input. Got: ", + num_segments_shape); + } + + if (!is_data_shape_rank_static) { + return {PartialShape::dynamic()}; + } + using TDim = typename TShape::value_type; + auto output_shapes = std::vector{data_shape}; + auto& output_shape = output_shapes[0]; + if (num_segments) { + output_shape[0] = TDim((*num_segments)[0]); + } else if (segment_ids && !num_segments_available) { + output_shape[0] = TDim(segment_ids->back() + 1); + } else { + // if num_segments input is provided but the value is unknown, the first dimension should be dynamic + output_shape[0] = Dimension::dynamic(); + } + return output_shapes; +} +} // namespace v16 +} // namespace op +} // namespace ov diff --git a/src/core/src/op/segment_max.cpp b/src/core/src/op/segment_max.cpp new file mode 100644 index 00000000000000..7b7e2276ba019a --- /dev/null +++ b/src/core/src/op/segment_max.cpp @@ -0,0 +1,72 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/segment_max.hpp" + +#include "itt.hpp" +#include "openvino/core/validation_util.hpp" +#include "openvino/op/op.hpp" +#include "segment_max_shape_inference.hpp" + +namespace ov { +namespace op { +namespace v16 { + +SegmentMax::SegmentMax(const Output& data, const Output& segment_ids, const op::FillMode fill_mode) + : Op({data, segment_ids}), + m_fill_mode(fill_mode) { + constructor_validate_and_infer_types(); +} + +SegmentMax::SegmentMax(const Output& data, + const Output& segment_ids, + const Output& num_segments, + const op::FillMode fill_mode) + : Op({data, segment_ids, num_segments}), + m_fill_mode(fill_mode) { + constructor_validate_and_infer_types(); +} + +bool SegmentMax::visit_attributes(ov::AttributeVisitor& visitor) { + OV_OP_SCOPE(v16_SegmentMax_visit_attributes); + visitor.on_attribute("fill_mode", m_fill_mode); + return true; +} + +void SegmentMax::validate_and_infer_types() { + OV_OP_SCOPE(v16_SegmentMax_validate_and_infer_types); + const auto& segment_ids_element_type = get_input_element_type(1); + NODE_VALIDATION_CHECK(this, + segment_ids_element_type == element::i32 || segment_ids_element_type == element::i64, + "The element type of the segment_ids input be i32 or i64. Got: ", + segment_ids_element_type); + if (inputs().size() == 3) { + const auto& num_segments_element_type = get_input_element_type(2); + NODE_VALIDATION_CHECK(this, + num_segments_element_type == element::i32 || num_segments_element_type == element::i64, + "The element type of the num_segments input be i32 or i64. Got: ", + num_segments_element_type); + } + + const auto output_shapes = shape_infer(this, ov::util::get_node_input_partial_shapes(*this)); + set_output_type(0, get_input_element_type(0), output_shapes[0]); +} + +std::shared_ptr SegmentMax::clone_with_new_inputs(const ov::OutputVector& new_args) const { + OV_OP_SCOPE(v16_SegmentMax_clone_with_new_inputs); + check_new_args_count(this, new_args); + if (new_args.size() == 3) { + return std::make_shared(new_args.at(0), new_args.at(1), new_args.at(2), m_fill_mode); + } else { + return std::make_shared(new_args.at(0), new_args.at(1), m_fill_mode); + } +} + +const op::FillMode SegmentMax::get_fill_mode() const { + return m_fill_mode; +} + +} // namespace v16 +} // namespace op +} // namespace ov diff --git a/src/core/src/op/util/attr_types.cpp b/src/core/src/op/util/attr_types.cpp index 5795fe9992dcc0..b6c2b8927eb217 100644 --- a/src/core/src/op/util/attr_types.cpp +++ b/src/core/src/op/util/attr_types.cpp @@ -21,6 +21,14 @@ OPENVINO_API EnumNames& EnumNames::get() { return enum_names; } +template <> +OPENVINO_API EnumNames& EnumNames::get() { + static auto enum_names = + EnumNames("ov::op::FillMode", + {{"zero", ov::op::FillMode::ZERO}, {"lowest", ov::op::FillMode::LOWEST}}); + return enum_names; +} + template <> OPENVINO_API EnumNames& EnumNames::get() { static auto enum_names = EnumNames("ov::op::PadType", @@ -132,6 +140,10 @@ std::ostream& op::operator<<(std::ostream& s, const ov::op::PadMode& type) { return s << as_string(type); } +std::ostream& op::operator<<(std::ostream& s, const ov::op::FillMode& type) { + return s << as_string(type); +} + std::ostream& op::operator<<(std::ostream& s, const ov::op::PadType& type) { return s << as_string(type); } diff --git a/src/core/tests/type_prop/segment_max.cpp b/src/core/tests/type_prop/segment_max.cpp new file mode 100644 index 00000000000000..a51ee7161e8f0d --- /dev/null +++ b/src/core/tests/type_prop/segment_max.cpp @@ -0,0 +1,381 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/segment_max.hpp" + +#include + +#include "common_test_utils/test_assertions.hpp" +#include "common_test_utils/type_prop.hpp" +#include "openvino/op/add.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/reduce_max.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/strided_slice.hpp" + +namespace ov::test { +using op::v0::Constant, op::v0::Parameter, op::v1::Add, op::v1::ReduceMax, op::v1::StridedSlice, op::v3::ShapeOf; + +class TypePropSegmentMaxTest : public TypePropOpTest {}; + +TEST_F(TypePropSegmentMaxTest, default_ctor) { + const auto data = std::make_shared(element::i32, PartialShape{3, 12, 225}); + const auto segment_ids = std::make_shared(element::i64, PartialShape{3}); + + const auto op = make_op(data, segment_ids, op::FillMode::LOWEST); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_fill_mode(), op::FillMode::LOWEST); + EXPECT_EQ(op->get_output_element_type(0), element::i32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 225})); +} + +TEST_F(TypePropSegmentMaxTest, non_default_args_no_values) { + PartialShape data_shape{3, 12, 81}; + auto data_symbols = set_shape_symbols(data_shape); + const auto data = std::make_shared(element::f32, data_shape); + const auto segment_ids = std::make_shared(element::i64, PartialShape{3}); + const auto num_segments = std::make_shared(element::i64, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::LOWEST); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::LOWEST); + EXPECT_THAT(get_shape_symbols(op->get_output_partial_shape(0)), + testing::ElementsAre(nullptr, data_symbols[1], data_symbols[2])); +} + +TEST_F(TypePropSegmentMaxTest, num_segments_bigger_than_segment_ids) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 0, 1}); + const auto num_segments = std::make_shared(element::i32, Shape{}, std::vector{40}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::LOWEST); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{40, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::LOWEST); +} + +TEST_F(TypePropSegmentMaxTest, incorrect_inputs) { + const auto data = std::make_shared(element::i32, PartialShape{3, 12, 225}); + const auto segment_ids = std::make_shared(element::i32, PartialShape{3}); + const auto num_segments = std::make_shared(element::i32, PartialShape{}); + { + const auto num_segments_f32 = std::make_shared(element::f32, PartialShape{}); + OV_EXPECT_THROW(std::ignore = make_op(data, segment_ids, num_segments_f32, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("The element type of the num_segments input be i32 or i64.")); + } + { + const auto segment_ids_f32 = std::make_shared(element::f32, PartialShape{3}); + OV_EXPECT_THROW(std::ignore = make_op(data, segment_ids_f32, num_segments, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("The element type of the segment_ids input be i32 or i64.")); + } + { + const auto segment_ids_nd = std::make_shared(element::i32, PartialShape{2, 3}); + OV_EXPECT_THROW(std::ignore = make_op(data, segment_ids_nd, num_segments, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("segment_ids must be a 1D input.")); + } + { + const auto num_segments_nd = std::make_shared(element::i32, PartialShape{1}); + OV_EXPECT_THROW(std::ignore = make_op(data, segment_ids, num_segments_nd, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("num_segments must be a scalar input.")); + } + { + const auto segment_ids_unsorted = + std::make_shared(element::i32, Shape{3}, std::vector{1, 0, 1}); + OV_EXPECT_THROW(std::ignore = make_op(data, segment_ids_unsorted, num_segments, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("segment_ids must be sorted.")); + } + { + const auto data_scalar = std::make_shared(element::i32, PartialShape{}); + OV_EXPECT_THROW(std::ignore = make_op(data_scalar, segment_ids, num_segments, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("The data input cannot be a scalar.")); + } + { + const auto segment_ids_short = std::make_shared(element::i32, Shape{2}, std::vector{1, 0}); + OV_EXPECT_THROW( + std::ignore = make_op(data, segment_ids_short, num_segments, op::FillMode::LOWEST), + ov::NodeValidationFailure, + testing::HasSubstr("The number of elements in segment_ids must match the first dimension of data.")); + } +} + +TEST_F(TypePropSegmentMaxTest, num_segments_from_graph_shapeof) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + const auto some_subgraph_result = std::make_shared(element::i32, PartialShape{8}); + const auto shape_of = std::make_shared(some_subgraph_result); + const auto num_segments = std::make_shared(shape_of, Constant::create(element::i64, Shape{}, {0})); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::LOWEST); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{8, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::LOWEST); +} + +TEST_F(TypePropSegmentMaxTest, num_segments_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + const auto max_segment_id = std::make_shared(segment_ids, Constant::create(element::i32, Shape{}, {0})); + const auto num_segments = std::make_shared(max_segment_id, Constant::create(element::i32, Shape{}, {1})); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::LOWEST); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{3, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::LOWEST); +} + +TEST_F(TypePropSegmentMaxTest, dynamic_num_segments_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, PartialShape::dynamic()); + const auto max_segment_id = std::make_shared(segment_ids, Constant::create(element::i32, Shape{}, {0})); + const auto num_segments = std::make_shared(max_segment_id, Constant::create(element::i32, Shape{}, {1})); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, interval_num_segments_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, PartialShape{{2, 4}}); + const auto max_segment_id = std::make_shared(segment_ids, Constant::create(element::i32, Shape{}, {0})); + const auto num_segments = std::make_shared(max_segment_id, Constant::create(element::i32, Shape{}, {1})); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, segment_ids_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids_input = + std::make_shared(element::i32, Shape{5}, std::vector{0, 1, 2, 3, 4}); + const auto begin = Constant::create(element::i64, Shape{1}, {0}); + const auto end = Constant::create(element::i64, Shape{1}, {3}); + const auto stride = Constant::create(element::i64, Shape{1}, {1}); + const auto segment_ids = std::make_shared(segment_ids_input, + begin, + end, + stride, + std::vector{0}, + std::vector{0}); + + const auto op = make_op(data, segment_ids, op::FillMode::ZERO); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{3, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, dynamic_segment_ids_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids_input = std::make_shared(element::i32, PartialShape::dynamic()); + const auto begin = Constant::create(element::i64, Shape{1}, {0}); + const auto end = Constant::create(element::i64, Shape{1}, {3}); + const auto stride = Constant::create(element::i64, Shape{1}, {1}); + const auto segment_ids = std::make_shared(segment_ids_input, + begin, + end, + stride, + std::vector{0}, + std::vector{0}); + + const auto op = make_op(data, segment_ids, op::FillMode::ZERO); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, interval_segment_ids_from_graph) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids_input = std::make_shared(element::i32, PartialShape{{2, 4}}); + const auto begin = Constant::create(element::i64, Shape{1}, {0}); + const auto end = Constant::create(element::i64, Shape{1}, {3}); + const auto stride = Constant::create(element::i64, Shape{1}, {1}); + const auto segment_ids = std::make_shared(segment_ids_input, + begin, + end, + stride, + std::vector{0}, + std::vector{0}); + + const auto op = make_op(data, segment_ids, op::FillMode::ZERO); + op->validate_and_infer_types(); + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, all_inputs_dynamic) { + const auto data = std::make_shared(element::f32, PartialShape::dynamic()); + const auto segment_ids = std::make_shared(element::i64, PartialShape::dynamic()); + const auto num_segments = std::make_shared(element::i64, PartialShape::dynamic()); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape::dynamic())); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, all_inputs_interval) { + const auto data = std::make_shared(element::f32, PartialShape{{2, 4}, {12, 15}, {18, 300}}); + const auto segment_ids = std::make_shared(element::i64, PartialShape{{2, 4}}); + const auto num_segments = std::make_shared(element::i64, PartialShape::dynamic()); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), {12, 15}, {18, 300}})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, data_dynamic) { + const auto data = std::make_shared(element::f32, PartialShape::dynamic()); + const auto segment_ids = std::make_shared(element::i64, PartialShape{6}); + const auto num_segments = std::make_shared(element::i64, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape::dynamic())); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, segment_ids_dynamic) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i64, PartialShape::dynamic()); + const auto num_segments = std::make_shared(element::i64, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, segment_ids_const_num_segments_dynamic) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + const auto num_segments = std::make_shared(element::i64, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, segment_ids_dynamic_num_segments_const) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, PartialShape::dynamic()); + const auto num_segments = std::make_shared(element::i32, Shape{}, std::vector{7}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{7, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, const_num_segments_zero) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + const auto num_segments = std::make_shared(element::i32, Shape{}, std::vector{0}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{0, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, segment_ids_const_no_num_segments) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + + const auto op = make_op(data, segment_ids, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{3, 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, num_segments_dynamic) { + const auto data = std::make_shared(element::f32, PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(element::i64, PartialShape{3}); + const auto num_segments = std::make_shared(element::i64, PartialShape::dynamic()); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), 12, 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, dynamic_dimensions) { + const auto data = std::make_shared(element::f32, PartialShape{3, Dimension::dynamic(), 81}); + const auto segment_ids = std::make_shared(element::i64, PartialShape{Dimension::dynamic()}); + const auto num_segments = std::make_shared(element::i64, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), Dimension::dynamic(), 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, dynamic_num_segments_const_segment_ids) { + const auto data = std::make_shared(element::f32, PartialShape{3, Dimension::dynamic(), 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + const auto num_segments = std::make_shared(element::i32, PartialShape{}); + + const auto op = make_op(data, segment_ids, num_segments, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{Dimension::dynamic(), Dimension::dynamic(), 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} + +TEST_F(TypePropSegmentMaxTest, no_num_segments_const_segment_ids) { + const auto data = std::make_shared(element::f32, PartialShape{3, Dimension::dynamic(), 81}); + const auto segment_ids = std::make_shared(element::i32, Shape{3}, std::vector{0, 1, 2}); + + const auto op = make_op(data, segment_ids, op::FillMode::ZERO); + op->validate_and_infer_types(); + + EXPECT_EQ(op->get_output_element_type(0), element::f32); + EXPECT_EQ(op->get_output_partial_shape(0), (PartialShape{3, Dimension::dynamic(), 81})); + EXPECT_EQ(op->get_fill_mode(), op::FillMode::ZERO); +} +} // namespace ov::test diff --git a/src/core/tests/visitors/op/segment_max.cpp b/src/core/tests/visitors/op/segment_max.cpp new file mode 100644 index 00000000000000..8f844ee5c3d1cb --- /dev/null +++ b/src/core/tests/visitors/op/segment_max.cpp @@ -0,0 +1,38 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/segment_max.hpp" + +#include + +#include "visitors/visitors.hpp" + +namespace ov::test { +using ov::op::v0::Parameter, ov::test::NodeBuilder; + +TEST(attributes, segment_max_v16_with_num_segments) { + NodeBuilder::opset().insert(); + const auto data = std::make_shared(ov::element::i32, ov::PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(ov::element::i64, ov::Shape{3}); + const auto num_segments = std::make_shared(ov::element::i64, ov::Shape{}); + + const auto op = std::make_shared(data, segment_ids, num_segments, ov::op::FillMode::ZERO); + NodeBuilder builder(op, {data, segment_ids, num_segments}); + auto g_op = ov::as_type_ptr(builder.create()); + + EXPECT_EQ(g_op->get_fill_mode(), op->get_fill_mode()); +} + +TEST(attributes, segment_max_v16_without_num_segments) { + NodeBuilder::opset().insert(); + const auto data = std::make_shared(ov::element::i32, ov::PartialShape{3, 12, 81}); + const auto segment_ids = std::make_shared(ov::element::i64, ov::Shape{3}); + + const auto op = std::make_shared(data, segment_ids, ov::op::FillMode::LOWEST); + NodeBuilder builder(op, {data, segment_ids}); + auto g_op = ov::as_type_ptr(builder.create()); + + EXPECT_EQ(g_op->get_fill_mode(), op->get_fill_mode()); +} +} // namespace ov::test diff --git a/src/plugins/intel_cpu/src/shape_inference/shape_inference.cpp b/src/plugins/intel_cpu/src/shape_inference/shape_inference.cpp index 7c664ee92b8fef..0274b0e5d0be4e 100644 --- a/src/plugins/intel_cpu/src/shape_inference/shape_inference.cpp +++ b/src/plugins/intel_cpu/src/shape_inference/shape_inference.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -104,6 +105,7 @@ #include "scatter_elements_update_shape_inference.hpp" #include "scatter_nd_base_shape_inference.hpp" #include "search_sorted_shape_inference.hpp" +#include "segment_max_shape_inference.hpp" #include "select_shape_inference.hpp" #include "shape_nodes.hpp" #include "shuffle_channels_shape_inference.hpp" @@ -435,6 +437,8 @@ using IStaticShapeInferFactory = // To use other version of operators, explicitly specify operator with opset version namespace. template <> const IStaticShapeInferFactory::TRegistry IStaticShapeInferFactory::registry{ + // opset16 + _OV_OP_SHAPE_INFER_MASK_REG(op::v16::SegmentMax, ShapeInferTA, util::bit::mask(1, 2)), // opset15 _OV_OP_SHAPE_INFER_MASK_REG(op::v15::Squeeze, ShapeInferTA, util::bit::mask(1)), _OV_OP_SHAPE_INFER_MASK_REG(op::v15::SearchSorted, ShapeInferTA, util::bit::mask()), diff --git a/src/plugins/intel_cpu/tests/unit/shape_inference_test/segment_max_shape_inference_test.cpp b/src/plugins/intel_cpu/tests/unit/shape_inference_test/segment_max_shape_inference_test.cpp new file mode 100644 index 00000000000000..cf3fca5b32cba0 --- /dev/null +++ b/src/plugins/intel_cpu/tests/unit/shape_inference_test/segment_max_shape_inference_test.cpp @@ -0,0 +1,116 @@ +// Copyright (C) 2018-2025 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include "common_test_utils/test_assertions.hpp" +#include "utils.hpp" + +using namespace ov::intel_cpu; +using ov::op::v0::Constant, ov::op::v0::Parameter; +using testing::HasSubstr; + +struct SegmentMaxTestParams { + ov::Shape data_shape; + std::vector segment_ids_val; + int64_t num_segments_val; + ov::Shape expected_output_shape; +}; + +class SegmentMaxStaticShapeInferenceTest: public OpStaticShapeInferenceTest {}; + +class SegmentMaxStaticTestSuite : public SegmentMaxStaticShapeInferenceTest, public ::testing::WithParamInterface {}; + +TEST_F(SegmentMaxStaticShapeInferenceTest, segment_ids_from_tensor_accessor) { + const auto data = std::make_shared(ov::element::i64, ov::PartialShape::dynamic()); + const auto segment_ids = std::make_shared(ov::element::i64, ov::PartialShape::dynamic()); + const auto op = make_op(data, segment_ids, ov::op::FillMode::ZERO); + + int64_t segment_ids_val[] = {0, 0, 1, 1, 4, 4, 5, 5, 5, 5}; + auto const_inputs = std::unordered_map{{1, {ov::element::i64, ov::Shape{10}, segment_ids_val}}}; + + const auto input_shapes = StaticShapeVector{{10, 12, 289}, {10}, {}}; + auto shape_infer = make_shape_inference(op); + const auto input_shape_refs = make_static_shape_refs(input_shapes); + const auto output_shapes = *shape_infer->infer(input_shape_refs, ov::make_tensor_accessor(const_inputs)); + EXPECT_EQ(output_shapes.size(), 1); + EXPECT_EQ(output_shapes.front(), StaticShape({6, 12, 289})); +} + +TEST_F(SegmentMaxStaticShapeInferenceTest, segment_ids_from_tensor_accessor_with_num_segments) { + const auto data = std::make_shared(ov::element::i64, ov::PartialShape::dynamic()); + const auto segment_ids = std::make_shared(ov::element::i64, ov::PartialShape::dynamic()); + const auto num_segments = std::make_shared(ov::element::i64, ov::PartialShape::dynamic()); + const auto op = make_op(data, segment_ids, num_segments, ov::op::FillMode::ZERO); + + int64_t segment_ids_val[] = {0, 0, 1, 1, 4, 4, 5, 5, 5, 5}; + int64_t num_segments_val[] = {6}; + auto const_inputs = std::unordered_map{{1, {ov::element::i64, ov::Shape{10}, segment_ids_val}}, + {2, {ov::element::i64, ov::Shape{}, num_segments_val}}}; + + const auto input_shapes = StaticShapeVector{{5, 12, 289}, {5}, {}}; + auto shape_infer = make_shape_inference(op); + const auto input_shape_refs = make_static_shape_refs(input_shapes); + const auto output_shapes = *shape_infer->infer(input_shape_refs, ov::make_tensor_accessor(const_inputs)); + EXPECT_EQ(output_shapes.size(), 1); + EXPECT_EQ(output_shapes.front(), StaticShape({6, 12, 289})); +} + +TEST_P(SegmentMaxStaticTestSuite, segment_max_static_shape_inference_with_num_segments) { + const auto& [data_shape, segment_ids_val, num_segments_val, expected_output_shape] = GetParam(); + const auto segment_ids_shape = ov::Shape{segment_ids_val.size()}; + + const auto data = std::make_shared(ov::element::f32, data_shape); + const auto segment_ids = std::make_shared(ov::element::i32, segment_ids_shape, segment_ids_val); + + std::shared_ptr op; + const auto num_segments = std::make_shared(ov::element::i32, ov::Shape{}, num_segments_val); + op = make_op(data, segment_ids, num_segments, ov::op::FillMode::ZERO); + + const auto input_shapes = StaticShapeVector{data_shape, segment_ids_shape, {}}; + auto shape_infer = make_shape_inference(op); + const auto input_shape_refs = make_static_shape_refs(input_shapes); + const auto output_shapes = *shape_infer->infer(input_shape_refs, ov::make_tensor_accessor()); + + EXPECT_EQ(output_shapes.size(), 1); + EXPECT_EQ(output_shapes.front(), StaticShape(expected_output_shape)); +} + +TEST_P(SegmentMaxStaticTestSuite, segment_max_static_shape_inference_no_num_segments) { + const auto& [data_shape, segment_ids_val, num_segments_val, expected_output_shape] = GetParam(); + const auto segment_ids_shape = ov::Shape{segment_ids_val.size()}; + + const auto data = std::make_shared(ov::element::f32, data_shape); + const auto segment_ids = std::make_shared(ov::element::i32, segment_ids_shape, segment_ids_val); + + std::shared_ptr op; + op = make_op(data, segment_ids, ov::op::FillMode::ZERO); + + const auto input_shapes = StaticShapeVector{data_shape, segment_ids_shape}; + auto shape_infer = make_shape_inference(op); + const auto input_shape_refs = make_static_shape_refs(input_shapes); + const auto output_shapes = *shape_infer->infer(input_shape_refs, ov::make_tensor_accessor()); + + EXPECT_EQ(output_shapes.size(), 1); + EXPECT_EQ(output_shapes.front(), StaticShape(expected_output_shape)); +} + +INSTANTIATE_TEST_SUITE_P(SegmentMaxStaticShapeInferenceTests, + SegmentMaxStaticTestSuite, + ::testing::Values(SegmentMaxTestParams{ov::Shape{6}, // data shape + std::vector{0, 0, 1, 4, 4, 4}, // segment_ids values + 5, // num_segments value + ov::Shape{5}}, // expected output shape + SegmentMaxTestParams{ov::Shape{1, 23}, // data shape + std::vector{200}, // segment_ids values + 201, // num_segments value + ov::Shape{201, 23}}, // expected output shape + SegmentMaxTestParams{ov::Shape{4, 23}, // data shape + std::vector{0, 0, 0, 0}, // segment_ids values + 1, // num_segments value + ov::Shape{1, 23}}, // expected output shape + SegmentMaxTestParams{ov::Shape{2, 6, 24, 1}, // data shape + std::vector{2, 6}, // segment_ids values + 7, // num_segments value + ov::Shape{7, 6, 24, 1}})); // expected output shape