[Op]: Added SearchSorted op def.

docs/articles_en/documentation/openvino-ir-format/operation-sets/operation-specs.rst

@@ -197,6 +197,7 @@ Operation Specifications
    ScatterElementsUpdate-12 <operation-specs/movement/scatter-elements-update-12>
    ScatterNDUpdate-3 <operation-specs/movement/scatter-nd-update-3>
    ScatterUpdate-3 <operation-specs/movement/scatter-update-3>
+   SearchSorted-15 <operation-specs/sort/search-sorted-15>
    Select-1 <operation-specs/condition/select-1>
    Selu-1 <operation-specs/activation/selu-1>
    ShapeOf-1 <operation-specs/shape/shape-of-1>

docs/articles_en/documentation/openvino-ir-format/operation-sets/operation-specs/sort/search-sorted-15.rst

@@ -0,0 +1,72 @@
+SearchSorted
+===============
+
+
+.. meta::
+  :description: Learn about SearchSorted - a sorting and maximization
+                operation, which requires two input tensors.
+
+
+**Versioned name**: *SearchSorted-15*
+
+**Category**: *Sorting and maximization*
+
+**Short description**: Find the indices from the innermost dimension of sorted_sequence such that, if the corresponding values in values were inserted before the indices, when sorted, the order of the corresponding innermost dimension within sorted_sequence would be preserved.
+
+**Detailed description**: `Reference <https://pytorch.org/docs/stable/generated/torch.searchsorted.html>`__.
+
+**Attributes**
+
+* *right*
+  * **Description**: *right*  if False, return the first suitable location that is found. If True, return the last such index. If no suitable index found, return 0 for non-numerical value (eg. nan, inf) or the size of innermost dimension within sorted_sequence (one pass the last index of the innermost dimension). In other words, if False, gets the lower bound index for each value in values on the corresponding innermost dimension of the sorted_sequence. If True, gets the upper bound index instead. Default value is False. side does the same and is preferred. It will error if side is set to “left” while this is True.
+  * **Range of values**: true or false
+  * **Type**: boolean
+  * **Default value**: false
+  * **Required**: *no*
+
+**Inputs**:
+
+* **1**: ND input tensor of type *T*, containing monotonically increasing sequence on the innermost dimension **Required.**
+
+* **2**: ND input tensor of type *T*, containing the search values. **Required.**
+
+**Outputs**:
+
+* **1**: Tensor of type *TOut*, with the same shape as second input tensor, containing the indices.
+
+**Types**
+
+* *T*: any supported floating-point and integer type.
+
+* *TOut*: int64.
+
+**Example**
+
+.. code-block:: xml
+   :force:
+
+   <layer ... type="SearchSorted" ... >
+       <right="True"/>
+       <input>
+           <port id="0">
+               <dim>7</dim>
+               <dim>256</dim>
+               <dim>200</dim>
+               <dim>200</dim>
+           </port>
+           <port id="1">
+               <dim>7</dim>
+               <dim>256</dim>
+               <dim>200</dim>
+               <dim>10</dim>
+           </port>
+       </input>
+       <output>
+           <port id="0" precision="INT64">
+               <dim>7</dim>
+               <dim>256</dim>
+               <dim>200</dim>
+               <dim>10</dim>
+           </port>
+       </output>
+   </layer>

src/core/include/openvino/op/ops.hpp

@@ -169,6 +169,7 @@
 #include "openvino/op/scatter_elements_update.hpp"
 #include "openvino/op/scatter_nd_update.hpp"
 #include "openvino/op/scatter_update.hpp"
+#include "openvino/op/search_sorted.hpp"
 #include "openvino/op/select.hpp"
 #include "openvino/op/selu.hpp"
 #include "openvino/op/shape_of.hpp"

src/core/include/openvino/op/search_sorted.hpp

@@ -0,0 +1,46 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/op.hpp"
+
+namespace ov {
+namespace op {
+namespace v15 {
+/// \brief SearchSorted operation.
+///
+/// \ingroup ov_ops_cpp_api
+class OPENVINO_API SearchSorted : public Op {
+public:
+    OPENVINO_OP("SearchSorted", "opset15", Op);
+
+    SearchSorted() = default;
+    /// \brief Constructs a SearchSorted operation.
+    /// \param sorted_sequence Sorted sequence to search in.
+    /// \param values          Values to search indexs for.
+    /// \param right_mode      If False, return the first suitable index that is found for given value. If True, return
+    /// the last such index.
+    SearchSorted(const Output<Node>& sorted_sequence, const Output<Node>& values, bool right_mode = false);
+
+    void validate_and_infer_types() override;
+    bool visit_attributes(AttributeVisitor& visitor) override;
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+
+    bool get_right_mode() const {
+        return m_right_mode;
+    }
+
+    void set_right_mode(bool right_mode) {
+        m_right_mode = right_mode;
+    }
+
+private:
+    void validate();
+    void infer_type();
+    bool m_right_mode;
+};
+}  // namespace v15
+}  // namespace op
+}  // namespace ov

src/core/shape_inference/include/search_sorted_shape_inference.hpp

@@ -0,0 +1,33 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/search_sorted.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace op {
+namespace v15 {
+template <class TShape, class TRShape = result_shape_t<TShape>>
+std::vector<TRShape> shape_infer(const SearchSorted* op, const std::vector<TShape>& input_shapes) {
+    const auto& sorted_shape = input_shapes[0];
+    const auto& values_shape = input_shapes[1];
+    auto output_shape = values_shape;
+    TShape::merge_into(output_shape, sorted_shape);
+
+    if (output_shape.rank().is_static()) {
+        auto last_it = output_shape.end() - 1;
+        if (values_shape.rank().is_static()) {
+            *last_it = *(input_shapes[1].end() - 1);
+        } else {
+            *last_it = Dimension::dynamic();
+        }
+    }
+
+    return {output_shape};
+}
+}  // namespace v15
+}  // namespace op
+}  // namespace ov

src/core/src/op/search_sorted.cpp

@@ -0,0 +1,68 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <openvino/op/search_sorted.hpp>
+
+#include "itt.hpp"
+#include "openvino/core/validation_util.hpp"
+#include "search_sorted_shape_inference.hpp"
+
+namespace ov {
+namespace op {
+namespace v15 {
+
+SearchSorted::SearchSorted(const Output<Node>& sorted_sequence, const Output<Node>& values, bool right_mode)
+    : Op({sorted_sequence, values}),
+      m_right_mode(right_mode) {
+    constructor_validate_and_infer_types();
+}
+
+void SearchSorted::infer_type() {
+    const auto& output_shapes = shape_infer(this, ov::util::get_node_input_partial_shapes(*this));
+    set_output_type(0, ov::element::i64, output_shapes[0]);
+}
+
+void SearchSorted::validate() {
+    NODE_VALIDATION_CHECK(this,
+                          get_input_element_type(0) == get_input_element_type(1),
+                          "Sorted sequence and values must have the same element type.");
+
+    const auto& sorted_shape = get_input_partial_shape(0);
+    const auto& values_shape = get_input_partial_shape(1);
+
+    if (sorted_shape.rank().is_static() && values_shape.rank().is_static()) {
+        NODE_VALIDATION_CHECK(this,
+                              sorted_shape.rank().get_length() == values_shape.rank().get_length(),
+                              "Sorted sequence and values have different ranks.");
+
+        for (int64_t i = 0; i < sorted_shape.rank().get_length() - 1; ++i) {
+            NODE_VALIDATION_CHECK(this,
+                                  sorted_shape[i].compatible(values_shape[i]),
+                                  "Sorted sequence and values has different ",
+                                  i,
+                                  " dimension.");
+        }
+    }
+}
+
+void SearchSorted::validate_and_infer_types() {
+    OV_OP_SCOPE(v15_SearchSorted_validate_and_infer_types);
+    validate();
+    infer_type();
+}
+
+bool SearchSorted::visit_attributes(AttributeVisitor& visitor) {
+    OV_OP_SCOPE(v15_SearchSorted_visit_attributes);
+    visitor.on_attribute("right_mode", m_right_mode);
+    return true;
+}
+
+std::shared_ptr<Node> SearchSorted::clone_with_new_inputs(const OutputVector& new_args) const {
+    OV_OP_SCOPE(v15_SearchSorted_clone_with_new_inputs);
+    check_new_args_count(this, new_args);
+    return std::make_shared<SearchSorted>(new_args.at(0), new_args.at(1), get_right_mode());
+}
+}  // namespace v15
+}  // namespace op
+}  // namespace ov

src/core/tests/type_prop/search_sorted.cpp

@@ -0,0 +1,84 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/op/search_sorted.hpp"
+
+#include "common_test_utils/type_prop.hpp"
+
+using namespace std;
+using namespace ov;
+
+#define EXPECT_THROW_SUBSTRING(STATEMENT, SUBSTRING)              \
+    try {                                                         \
+        STATEMENT;                                                \
+        FAIL() << "Exception not thrown";                         \
+    } catch (const NodeValidationFailure& error) {                \
+        EXPECT_THAT(error.what(), testing::HasSubstr(SUBSTRING)); \
+    } catch (...) {                                               \
+        FAIL() << "Unexpected exception thrown";                  \
+    }
+
+static void PerformShapeTest(const PartialShape& sorted_shape,
+                             const PartialShape& values_shape,
+                             const PartialShape& expected_output_shape) {
+    auto sorted = make_shared<op::v0::Parameter>(element::i32, sorted_shape);
+    auto values = make_shared<op::v0::Parameter>(element::i32, values_shape);
+    auto search_sorted_op = make_shared<op::v15::SearchSorted>(sorted, values);
+    EXPECT_EQ(search_sorted_op->get_element_type(), element::i64);
+    EXPECT_EQ(search_sorted_op->get_output_partial_shape(0), expected_output_shape);
+}
+
+TEST(type_prop, search_sorted_shape_infer_equal_inputs) {
+    PerformShapeTest({1, 3, 6}, {1, 3, 6}, {1, 3, 6});
+}
+
+TEST(type_prop, search_sorted_shape_infer_sorted_dynamic) {
+    PerformShapeTest(PartialShape::dynamic(), {1, 3, 6}, {1, 3, 6});
+}
+
+TEST(type_prop, search_sorted_shape_infer_values_dynamic) {
+    PerformShapeTest({1, 3, 7, 5}, PartialShape::dynamic(), {1, 3, 7, -1});
+}
+
+TEST(type_prop, search_sorted_shape_infer_different_last_dim) {
+    PerformShapeTest({1, 3, 7, 100}, {1, 3, 7, 10}, {1, 3, 7, 10});
+}
+
+TEST(type_prop, search_sorted_shape_infer_both_dynamic_1) {
+    PerformShapeTest({1, -1, 7, -1}, {-1, 3, -1, 10}, {1, 3, 7, 10});
+}
+
+TEST(type_prop, search_sorted_shape_infer_both_dynamic_2) {
+    PerformShapeTest({1, -1, 7, 50}, {-1, 3, -1, -1}, {1, 3, 7, -1});
+}
+
+TEST(type_prop, search_sorted_shape_infer_both_dynamic_3) {
+    PerformShapeTest(PartialShape::dynamic(), PartialShape::dynamic(), PartialShape::dynamic());
+}
+
+TEST(type_prop, search_sorted_shape_infer_both_dynamic_4) {
+    PerformShapeTest({-1, -1, 50}, {-1, -1, 20}, {-1, -1, 20});
+}
+
+TEST(type_prop, search_sorted_shape_infer_different_types) {
+    auto sorted = make_shared<ov::op::v0::Parameter>(element::f32, Shape{1, 3, 6});
+    auto values = make_shared<ov::op::v0::Parameter>(element::i32, Shape{1, 3, 6});
+    EXPECT_THROW_SUBSTRING(make_shared<op::v15::SearchSorted>(values, sorted),
+                           std::string("must have the same element type"));
+}
+
+TEST(type_prop, search_sorted_shape_infer_wrong_rank) {
+    auto sorted = make_shared<ov::op::v0::Parameter>(element::i32, Shape{1, 1, 3, 6});
+    auto values = make_shared<ov::op::v0::Parameter>(element::i32, Shape{1, 3, 6});
+    EXPECT_THROW_SUBSTRING(make_shared<op::v15::SearchSorted>(sorted, values),
+                           std::string("Sorted sequence and values have different ranks"));
+}
+
+TEST(type_prop, search_sorted_shape_infer_wrong_dim) {
+    auto sorted = make_shared<ov::op::v0::Parameter>(element::i32, Shape{1, 1, 3, 6});
+    auto values = make_shared<ov::op::v0::Parameter>(element::i32, Shape{1, 1, 5, 6});
+    EXPECT_THROW_SUBSTRING(make_shared<op::v15::SearchSorted>(sorted, values), std::string(" different 2 dimension."));
+}
+
+#undef EXPECT_THROW_SUBSTRING