[CPU]: Added impl of SearchSorted op

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -253,7 +253,8 @@ static const TypeToNameMap& get_type_to_name_tbl() {
         {"EmbeddingBagOffsets", Type::EmbeddingBagOffsets},
         {"LLMMLP", Type::LLMMLP},
         {"QKVProjection", Type::QKVProjection},
-        {"RMS", Type::RMS}
+        {"RMS", Type::RMS},
+        {"SearchSorted", Type::SearchSorted}
     };
     return type_to_name_tbl;
 }
@@ -384,6 +385,7 @@ std::string NameFromType(const Type type) {
         CASE(LLMMLP);
         CASE(QKVProjection);
         CASE(RMS);
+        CASE(SearchSorted);
         CASE(Unknown);
     }
 #undef CASE

src/plugins/intel_cpu/src/cpu_types.h

@@ -131,7 +131,8 @@ enum class Type {
     CausalMaskPreprocess,
     LLMMLP,
     QKVProjection,
-    RMS
+    RMS,
+    SearchSorted,
 };
 
 enum class Algorithm {

src/plugins/intel_cpu/src/nodes/search_sorted.cpp

@@ -0,0 +1,116 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "search_sorted.h"
+
+#include "openvino/op/search_sorted.hpp"
+#include "openvino/reference/search_sorted.hpp"
+
+namespace ov {
+namespace intel_cpu {
+namespace node {
+SearchSorted::SearchSorted(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
+    : Node(op, context, NgraphShapeInferFactory(op, EMPTY_PORT_MASK)) {
+    std::string errorMessage;
+    if (!isSupportedOperation(op, errorMessage)) {
+        OPENVINO_THROW_NOT_IMPLEMENTED(errorMessage);
+    }
+    const auto ss_op = ov::as_type_ptr<const ov::op::v15::SearchSorted>(op);
+    right_mode = ss_op->get_right_mode();
+}
+
+bool SearchSorted::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept {
+    try {
+        if (!ov::is_type<ov::op::v15::SearchSorted>(op)) {
+            errorMessage = "Only opset15 SearchSorted operation is supported";
+            return false;
+        }
+    } catch (...) {
+        return false;
+    }
+    return true;
+}
+
+void SearchSorted::getSupportedDescriptors() {
+    // Validation is already done in the ov::opset15::SearchSorted.
+}
+
+void SearchSorted::initSupportedPrimitiveDescriptors() {
+    if (!supportedPrimitiveDescriptors.empty())
+        return;
+
+    ov::element::Type inputPrec = getOriginalInputPrecisionAtPort(0);
+    ov::element::Type outputPrec = getOriginalOutputPrecisionAtPort(0);
+
+    addSupportedPrimDesc({{LayoutType::ncsp, inputPrec}, {LayoutType::ncsp, inputPrec}},
+                         {{LayoutType::ncsp, outputPrec}},
+                         impl_desc_type::ref);
+}
+
+bool SearchSorted::created() const {
+    return getType() == Type::SearchSorted;
+}
+
+bool SearchSorted::needPrepareParams() const {
+    return false;
+}
+
+void SearchSorted::executeDynamicImpl(dnnl::stream strm) {
+    execute(strm);
+}
+
+template <typename INPUT_TYPE, typename OUTPUT_TYPE>
+void SearchSorted::executeImpl() {
+    ov::reference::search_sorted<INPUT_TYPE, OUTPUT_TYPE>(getSrcDataAtPortAs<const INPUT_TYPE>(0),
+                                                          getSrcDataAtPortAs<const INPUT_TYPE>(1),
+                                                          getDstDataAtPortAs<OUTPUT_TYPE>(0),
+                                                          ov::Shape{getSrcMemoryAtPort(0)->getStaticDims()},
+                                                          ov::Shape{getSrcMemoryAtPort(1)->getStaticDims()},
+                                                          right_mode);
+}
+
+namespace {
+struct SearchSortedContext {
+    SearchSorted& node;
+};
+}  // namespace
+
+template <typename T>
+struct SearchSorted::SearchSortedExecute {
+    using TInputType = typename std::tuple_element<0, T>::type;
+    using TOutputType = typename std::tuple_element<1, T>::type;
+
+    void operator()(SearchSortedContext& ctx) {
+        ctx.node.executeImpl<TInputType, TOutputType>();
+    }
+};
+void SearchSorted::execute(dnnl::stream strm) {
+    auto inputPrecision = getParentEdgeAt(0)->getMemory().getDesc().getPrecision();
+    auto outputPrecision = getChildEdgeAt(0)->getMemory().getDesc().getPrecision();
+
+    SearchSortedContext ctx = {*this};
+
+#define CASE(OV_TYPE) OV_CASE2(OV_TYPE, ov::element::i64, ov::element_type_traits<OV_TYPE>::value_type, int64_t)
+
+    OV_SWITCH(intel_cpu,
+              SearchSortedExecute,
+              ctx,
+              std::tie(inputPrecision, outputPrecision),
+              CASE(ov::element::f64),
+              CASE(ov::element::f32),
+              CASE(ov::element::f16),
+              CASE(ov::element::i64),
+              CASE(ov::element::i32),
+              CASE(ov::element::i16),
+              CASE(ov::element::i8),
+              CASE(ov::element::u64),
+              CASE(ov::element::u32),
+              CASE(ov::element::u16),
+              CASE(ov::element::u8))
+
+#undef CASE
+}
+}  // namespace node
+}  // namespace intel_cpu
+}  // namespace ov

src/plugins/intel_cpu/src/nodes/search_sorted.h

@@ -0,0 +1,37 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "node.h"
+
+namespace ov {
+namespace intel_cpu {
+namespace node {
+
+class SearchSorted : public Node {
+public:
+    SearchSorted(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);
+
+    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
+    void getSupportedDescriptors() override;
+    void initSupportedPrimitiveDescriptors() override;
+    void execute(dnnl::stream strm) override;
+    bool created() const override;
+    bool needPrepareParams() const override;
+    void executeDynamicImpl(dnnl::stream strm) override;
+
+private:
+    template <typename INPUT_TYPE, typename OUTPUT_TYPE>
+    void executeImpl();
+
+    template <typename T>
+    struct SearchSortedExecute;
+
+    bool right_mode = false;
+};
+
+}  // namespace node
+}  // namespace intel_cpu
+}  // namespace ov

src/plugins/intel_cpu/src/nodes_factory.cpp

@@ -90,6 +90,7 @@
 #include "nodes/rope.h"
 #include "nodes/scaled_attn.h"
 #include "nodes/scatter_update.h"
+#include "nodes/search_sorted.h"
 #include "nodes/string_tensor_pack.h"
 #include "nodes/string_tensor_unpack.h"
 #include "nodes/shapeof.h"
@@ -217,6 +218,7 @@ Node::NodesFactory::NodesFactory() : Factory("NodesFactory") {
     INTEL_CPU_NODE(Subgraph, Type::Subgraph);
     INTEL_CPU_NODE(Composite, Type::SubModel);
     INTEL_CPU_NODE(ScaledDotProductAttention, Type::ScaledDotProductAttention);
+    INTEL_CPU_NODE(SearchSorted, Type::SearchSorted);
 #if defined(OPENVINO_ARCH_X86_64)
     INTEL_CPU_NODE(FakeQuantize, Type::FakeQuantize);
     INTEL_CPU_NODE(GridSample, Type::GridSample);

src/plugins/intel_cpu/tests/functional/shared_tests_instances/single_layer_tests/search_sorted.cpp

@@ -0,0 +1,18 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "single_op_tests/search_sorted.hpp"
+
+namespace ov {
+namespace test {
+
+INSTANTIATE_TEST_SUITE_P(
+    smoke_SearchSortedTest,
+    SearchSortedLayerTest,
+    ::testing::Combine(::testing::ValuesIn(SearchSortedLayerTest::GenerateParams()),
+                       testing::Values(ElementType::f32, ElementType::f16, ElementType::i64, ElementType::u32)),
+    SearchSortedLayerTest::getTestCaseName);
+
+}  // namespace test
+}  // namespace ov

src/tests/functional/plugin/shared/include/single_op_tests/search_sorted.hpp

@@ -0,0 +1,15 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "shared_test_classes/single_op/search_sorted.hpp"
+
+namespace ov {
+namespace test {
+TEST_P(SearchSortedLayerTest, Inference) {
+    run();
+}
+}  // namespace test
+}  // namespace ov

src/tests/functional/shared_test_classes/include/shared_test_classes/single_op/search_sorted.hpp

@@ -0,0 +1,32 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "shared_test_classes/base/ov_subgraph.hpp"
+
+namespace ov {
+namespace test {
+
+using SearchSortedSpecificParams = std::tuple<InputShape,  // sorted shape
+                                              InputShape,  // values shape
+                                              bool>;
+
+using SearchSortedLayerTestParams = std::tuple<SearchSortedSpecificParams, ElementType>;
+
+class SearchSortedLayerTest : public testing::WithParamInterface<SearchSortedLayerTestParams>,
+                              public ov::test::SubgraphBaseTest {
+public:
+    static std::string getTestCaseName(testing::TestParamInfo<SearchSortedLayerTestParams> obj);
+    static const std::vector<SearchSortedSpecificParams> GenerateParams();
+
+protected:
+    void SetUp() override;
+    void generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) override;
+};
+
+extern const std::vector<SearchSortedSpecificParams> SearchSortedParamsVector;
+
+}  // namespace test
+}  // namespace ov

src/tests/functional/shared_test_classes/src/single_op/search_sorted.cpp

@@ -0,0 +1,100 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "shared_test_classes/single_op/search_sorted.hpp"
+
+#include "common_test_utils/ov_tensor_utils.hpp"
+#include "shared_test_classes/base/ov_subgraph.hpp"
+
+namespace ov {
+namespace test {
+
+static const int SEED = 7877;
+
+std::string SearchSortedLayerTest::getTestCaseName(testing::TestParamInfo<SearchSortedLayerTestParams> obj) {
+    SearchSortedLayerTestParams basicParamsSet;
+    basicParamsSet = obj.param;
+
+    SearchSortedSpecificParams searchSortedParams;
+
+    ElementType inputPrecision;
+    std::tie(searchSortedParams, inputPrecision) = basicParamsSet;
+
+    InputShape sortedInputShape;
+    InputShape valuesInputShape;
+    bool right_mode;
+    std::tie(sortedInputShape, valuesInputShape, right_mode) = searchSortedParams;
+
+    std::ostringstream result;
+    result << inputPrecision << "_IS=";
+    result << ov::test::utils::partialShape2str({sortedInputShape.first}) << ",";
+    result << ov::test::utils::partialShape2str({valuesInputShape.first}) << "_";
+    result << "TS=";
+    result << "(";
+    for (const auto& targetShape : sortedInputShape.second) {
+        result << ov::test::utils::vec2str(targetShape) << "_";
+    }
+    result << ", ";
+    for (const auto& targetShape : valuesInputShape.second) {
+        result << ov::test::utils::vec2str(targetShape) << "_";
+    }
+    result << ")_";
+    result << "right_mode=" << right_mode;
+
+    return result.str();
+}
+
+void SearchSortedLayerTest::generate_inputs(const std::vector<ov::Shape>& targetInputStaticShapes) {
+    inputs.clear();
+    const auto& funcInputs = function->inputs();
+
+    const auto dataPrecision = funcInputs[0].get_element_type();
+
+    auto sortedTensor =
+        ov::test::utils::create_and_fill_tensor_unique_sequence(dataPrecision, targetInputStaticShapes[0], 0, 8, SEED);
+    inputs.insert({funcInputs[0].get_node_shared_ptr(), sortedTensor});
+
+    auto valuesTensor =
+        ov::test::utils::create_and_fill_tensor_unique_sequence(dataPrecision, targetInputStaticShapes[1], 0, 8, SEED);
+
+    inputs.insert({funcInputs[1].get_node_shared_ptr(), valuesTensor});
+}
+
+void SearchSortedLayerTest::SetUp() {
+    SearchSortedLayerTestParams basicParamsSet;
+    basicParamsSet = this->GetParam();
+
+    SearchSortedSpecificParams searchSortedParams;
+
+    ElementType inputPrecision;
+    std::tie(searchSortedParams, inputPrecision) = basicParamsSet;
+
+    targetDevice = ov::test::utils::DEVICE_CPU;
+
+    InputShape sortedInputShape;
+    InputShape valuesInputShape;
+    bool right_mode;
+    std::tie(sortedInputShape, valuesInputShape, right_mode) = searchSortedParams;
+
+    init_input_shapes({sortedInputShape, valuesInputShape});
+    auto sortedParam = std::make_shared<ov::op::v0::Parameter>(inputPrecision, inputDynamicShapes[0]);
+    auto valuesParam = std::make_shared<ov::op::v0::Parameter>(inputPrecision, inputDynamicShapes[1]);
+
+    auto op = std::make_shared<ov::op::v15::SearchSorted>(sortedParam, valuesParam, right_mode);
+
+    ov::ParameterVector params{sortedParam, valuesParam};
+    function = std::make_shared<ov::Model>(op->outputs(), params, "SearchSorted");
+}
+
+const std::vector<SearchSortedSpecificParams> SearchSortedLayerTest::GenerateParams() {
+    const std::vector<SearchSortedSpecificParams> params = {
+        SearchSortedSpecificParams{InputShape{{}, {{1, 18, 104}}}, InputShape{{}, {{1, 18, 104}}}, true},
+        SearchSortedSpecificParams{InputShape{{}, {{50}}}, InputShape{{1, -1, 10}, {{1, 18, 10}}}, false},
+    };
+
+    return params;
+}
+
+}  // namespace test
+}  // namespace ov