[CPU]Implement TopK-3 to CPU plugin (#6272)

src/plugins/intel_cpu/src/nodes/mkldnn_topk_node.h

@@ -4,105 +4,144 @@
 
 #pragma once
 
-#include "ie_common.h"
 #include <ie_common.h>
 #include <mkldnn_node.h>
+#include <string>
+#include <memory>
+#include <vector>
 
 namespace MKLDNNPlugin {
 
+enum TopKLayoutType {
+    topk_ncsp,
+    topk_nspc,
+    topk_blocked
+};
+
+enum TopKAlgorithm {
+    topk_bubble_sort,
+    topk_bitonic_sort,
+    topk_heap_sort
+};
+
+struct jit_topk_config_params {
+    bool mode_max;           // which of the two elements to select. ture: max; false: min
+    bool sort_index;         // sort by value or index. true: index; false: value
+    bool topk_innermost;     // if topk sorting is applied on innermost dimension or other dimension
+    bool bubble_inplace;     // all the elements in sorting is right in the register, no need to load and store for each comparison
+    TopKLayoutType layout;   // memory layout
+    TopKAlgorithm algorithm; // topk sorting algorithm
+    InferenceEngine::Precision precision; // precision
+    int data_size;           // data size
+    int blk_size;            // block size
+    int top_k;               // number of the output elements in the sorting dimension
+    int work_amount;         // how many elements are processed when call jit kernel once
+    int axis_dim;            // size of topk axis
+    int sort_stride;         // memory stride of adjacent elements in sorting
+    int bitonic_idx_cnt;     // the repeatedly counted total number of elements in sorting, which equal the total number of comparison x 2
+    int bitonic_k_idx_cnt;   // the counterpart of bitonic_idx_cnt, when sort_index == true
+};
+
+struct jit_topk_call_args {
+    const void *src;
+    void *process;
+    void *process_index;
+    void *dst;
+    void *index;
+    const int *bitonic_idx_buf;
+    const int *bitonic_k_idx_buf;
+    const int *idx_block_buf;// original idx sequence, repeated by block (eg. 00000000,11111111,...,77777777), only used in bubble sort
+    const int *idx_seq_buf;  // original idx sequence (eg. 01234567), only used in bubble sort and heap sort
+    size_t axis_dim;         // point to axis_dim, only used in heap sort with dynamic shapes to achieve axis_dim agnosic
+    size_t top_k;
+    size_t work_amount;
+    size_t sort_stride;
+};
+
+struct jit_uni_topk_kernel {
+    void (*ker_)(const jit_topk_call_args *);
+
+    void operator()(const jit_topk_call_args *args) {
+        assert(ker_);
+        ker_(args);
+    }
+
+    explicit jit_uni_topk_kernel(jit_topk_config_params jcp) : ker_(nullptr), jcp_(jcp) {}
+    virtual ~jit_uni_topk_kernel() {}
+
+    virtual void create_ker() = 0;
+
+    jit_topk_config_params jcp_;
+};
+
 class MKLDNNTopKNode : public MKLDNNNode {
 public:
-    MKLDNNTopKNode(const std::shared_ptr<ngraph::Node> &op, const mkldnn::engine &eng,
-                   MKLDNNWeightsSharing::Ptr &cache);
+    MKLDNNTopKNode(const std::shared_ptr<ngraph::Node>& op, const mkldnn::engine& eng, MKLDNNWeightsSharing::Ptr &cache);
+    ~MKLDNNTopKNode() override = default;
 
-    void getSupportedDescriptors() override {};
+    void getSupportedDescriptors() override;
     void initSupportedPrimitiveDescriptors() override;
-    void execute(mkldnn::stream strm) override;
-    bool created() const override;
-
-    bool needPrepareParams() const override;
-    void executeDynamicImpl(mkldnn::stream strm) override;
     bool needShapeInfer() const override;
     std::vector<VectorDims> shapeInfer() const override;
+    bool needPrepareParams() const override;
+    void prepareParams() override;
+    bool created() const override;
+    void execute(mkldnn::stream strm) override;
+    void executeDynamicImpl(mkldnn::stream strm) override;
+    bool canBeInPlace() const override {
+        return false;
+    }
 
     static bool isSupportedOperation(const std::shared_ptr<const ngraph::Node> &op, std::string &errorMessage) noexcept;
 
-#if defined(HAVE_AVX512F)
-    const int block_size = 16;
-    typedef __m512 vec_type_f;
-    typedef __m512i vec_type_i;
-    typedef __mmask16 vmask_type;
-#elif defined(HAVE_AVX2)
-    const int block_size = 8;
-    typedef __m256 vec_type_f;
-    typedef __m256i vec_type_i;
-    typedef __m256 vmask_type;
-#elif defined(HAVE_SSE)
-    const int block_size = 4;
-    typedef __m128 vec_type_f;
-    typedef __m128i vec_type_i;
-    typedef __m128 vmask_type;
-#else
-    typedef float vec_type_f;
-    typedef int vmask_type;
-#endif
-
-    struct cmpgt_ps {
-        static inline vmask_type cmp_ps(const vec_type_f _Left, const vec_type_f _Right) {
-#if defined(HAVE_SSE) || defined(HAVE_AVX2) || defined(HAVE_AVX512F)
-            return _mm_uni_cmpgt_ps(_Left, _Right);
-#else
-            return _Left > _Right ? _Left : _Right;
-#endif
-        }
-    };
-
-    struct cmplt_ps {
-        static inline vmask_type cmp_ps(const vec_type_f _Left, const vec_type_f _Right) {
-#if defined(HAVE_SSE) || defined(HAVE_AVX2) || defined(HAVE_AVX512F)
-            return _mm_uni_cmpgt_ps(_Right, _Left);
-#else
-            return _Right > _Left ? _Right : _Left;
-#endif
-        }
-    };
-
-    template<class Compare1, template<typename> class Compare2>
-    void top1_axis(const float *src_data, float *dst_data, int *dst_idx, InferenceEngine::SizeVector in_dims);
-
-    template<template<typename> class Compare>
-    void top1(const float *src_data, float *dst_data, int *dst_idx, InferenceEngine::SizeVector in_dims);
-
-    template<class Compare1, template<typename> class Compare2>
-    void topk_axis(const float *src_data, float *dst_data, int *dst_idx, InferenceEngine::SizeVector in_dims);
-
-    template<template<typename> class Compare>
-    void topk(const float *src_data, float *dst_data, int *dst_idx, InferenceEngine::SizeVector in_dims);
-
 private:
-    const size_t TOPK_DATA = 0;
-    const size_t TOPK_K = 1;
-    const size_t TOPK_VALUE = 0;
-    const size_t TOPK_INDEX = 1;
+    void topk_process(const uint8_t *in_ptr, uint8_t *out_ptr, uint8_t *dst_idx);
+    void topk_ref(const float *in_ptr, float *out_ptr, int32_t *dst_idx);
+    inline void topk_kernel_process(const uint8_t *in_p, uint8_t *out_p, uint8_t *src_idx,
+                                    uint8_t *process_p, uint8_t *process_idx_p, size_t work_amount);
+    inline static int count(InferenceEngine::SizeVector dims, size_t start_ind, size_t end_ind);
+    inline static int count(InferenceEngine::SizeVector dims, size_t start_ind = 0);
+    inline void bitonic_push_idx(int p, int n, std::vector<int> &vec, int &cnt, bool cmp_val = true);
+    void calc_bitonic_idx(size_t n, int &cnt, bool cmp_val);
+    void calc_dims_size(const InferenceEngine::SizeVector &layout_dims);
+    void topk_ref_process(const float* src_data, float* dst_data, int32_t* dst_idx,
+                   const InferenceEngine::SizeVector &in_dims, std::function<float(float, float)> compare) const;
+    void prepare_original_idx();
 
-    size_t axis;
-    int src_k = 1;
+    bool topk_innermost;
+    bool jit_mode;
+    bool sort_index;
+    bool mode_max;
+    int axis;
+    static const size_t TOPK_DATA = 0;
+    static const size_t TOPK_K = 1;
+    static const size_t TOPK_INDEX = 1;
+    size_t O, A, I;
+    size_t blk_size;
+    size_t count_xmm;
+    size_t data_size;
+    size_t axis_dim;
+    int top_k;
+    int dim, before_num;
+    bool bubble_inplace;
+    bool compile_kernel;
 
-    bool sort_value = false;
-    bool mode_max = true;
+    InferenceEngine::SizeVector src_dims, dst_dims;
+    TopKLayoutType layout;
+    TopKAlgorithm algorithm;
 
-    int dim = 0;
-    int before_num = 0;
+    std::vector<int> vec_bitonic_idx;
+    std::vector<int> vec_bitonic_k_idx;
 
-#if defined(HAVE_AVX512F)
-    const int count_vec = 32;
-#elif defined(HAVE_SSE) || defined(HAVE_AVX2)
-    const int count_vec = 16;
-#endif
+    std::vector<int> vec_idx_seq;
+    std::vector<int> vec_idx_block;
 
-    inline int count(InferenceEngine::SizeVector dims, size_t start_ind, size_t end_ind);
+    std::vector<uint8_t> vec_process_ptr;
+    std::vector<uint8_t> vec_process_idx_ptr;
 
-    inline int count(InferenceEngine::SizeVector dims, size_t start_ind = 0);
+    std::shared_ptr<jit_uni_topk_kernel> topk_kernel;
+
+    std::string errorPrefix;
 };
 
 }  // namespace MKLDNNPlugin

src/tests/functional/plugin/cpu/bfloat16/topk_inputs_i32.cpp

@@ -135,7 +135,7 @@ protected:
         expectedPrecisions["Add_4"] = "ndef";
         expectedPrecisions["Convolution_1"] = "BF16";
         expectedPrecisions["Convolution_2"] = "BF16";
-        expectedPrecisions["TopK_1"] = "FP32";
+        expectedPrecisions["TopK_1"] = netPrecision.name(); // tail kept in FP32 precision
     }
 };
 

src/tests/functional/plugin/cpu/shared_tests_instances/single_layer_tests/topk.cpp

@@ -19,12 +19,15 @@ const std::vector<int64_t> axes = {
         0,
         1,
         2,
+        3
 };
 
 const std::vector<int64_t> k = {
         1,
         5,
-        10,
+        7,
+        18,
+        21
 };
 
 const std::vector<ngraph::opset4::TopK::Mode> modes = {
@@ -48,7 +51,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_TopK, TopKLayerTest,
                 ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                 ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                 ::testing::Values(InferenceEngine::Layout::ANY),
-                ::testing::Values(std::vector<size_t>({10, 10, 10})),
+                ::testing::Values(std::vector<size_t>({21, 21, 21, 21})),
                 ::testing::Values(CommonTestUtils::DEVICE_CPU)),
         TopKLayerTest::getTestCaseName);
 }  // namespace

src/tests/functional/plugin/cpu/single_layer_tests/topk.cpp

@@ -3,104 +3,209 @@
 //
 
 #include "functional_test_utils/ov_tensor_utils.hpp"
-#include "common_test_utils/data_utils.hpp"
 #include "test_utils/cpu_test_utils.hpp"
-
 #include "shared_test_classes/base/ov_subgraph.hpp"
 #include "ngraph_functions/builders.hpp"
 
-using namespace ngraph;
 using namespace InferenceEngine;
 using namespace CPUTestUtils;
 using namespace ov::test;
 
 namespace CPULayerTestsDefinitions {
-using topKParams = std::tuple<
-    ngraph::element::Type, // data precision
-    InputShape, // data shape
-    std::int64_t, // axis
-    ngraph::opset4::TopK::Mode,
-    ngraph::opset4::TopK::SortType>;
 
-class TopKLayerCPUTest : public testing::WithParamInterface<topKParams>, public ov::test::SubgraphBaseTest, public CPUTestsBase {
+typedef std::tuple<
+        int64_t,                        // keepK
+        int64_t,                        // axis
+        ngraph::opset4::TopK::Mode,     // mode
+        ngraph::opset4::TopK::SortType, // sort
+        ElementType,                    // Net precision
+        ElementType,                    // Input precision
+        ElementType,                    // Output precision
+        InputShape                      // inputShape
+> basicTopKParams;
+
+typedef std::tuple<
+        basicTopKParams,
+        CPUSpecificParams,
+        std::map<std::string, std::string>> TopKLayerCPUTestParamsSet;
+
+class TopKLayerCPUTest : public testing::WithParamInterface<TopKLayerCPUTestParamsSet>,
+                         virtual public SubgraphBaseTest, public CPUTestsBase {
 public:
-    static std::string getTestCaseName(testing::TestParamInfo<topKParams> obj) {
-        ngraph::element::Type inputPrecision;
-        InputShape inputShape;
-        std::int64_t axis;
+    static std::string getTestCaseName(testing::TestParamInfo<TopKLayerCPUTestParamsSet> obj) {
+        basicTopKParams basicParamsSet;
+        CPUSpecificParams cpuParams;
+        std::map<std::string, std::string> additionalConfig;
+        std::tie(basicParamsSet, cpuParams, additionalConfig) = obj.param;
+
+        int64_t keepK, axis;
         ngraph::opset4::TopK::Mode mode;
-        ngraph::opset4::TopK::SortType sortType;
-        std::tie(inputPrecision, inputShape, axis, mode, sortType) = obj.param;
+        ngraph::opset4::TopK::SortType sort;
+        ElementType netPrecision, inPrc, outPrc;
+        InputShape inputShape;
+        std::tie(keepK, axis, mode, sort, netPrecision, inPrc, outPrc, inputShape) = basicParamsSet;
 
         std::ostringstream result;
-        result << inputPrecision << "_" << "IS=" << CommonTestUtils::partialShape2str({ inputShape.first }) << "_" << "TS=(";
+        bool staticShape = inputShape.first.rank() == 0;
+        if (staticShape)
+            result << "k=" << keepK << "_";
+        result << "axis=" << axis << "_";
+        result << "mode=" << mode << "_";
+        result << "sort=" << sort << "_";
+        result << "netPRC=" << netPrecision << "_";
+        result << "inPRC=" << inPrc << "_";
+        result << "outPRC=" << outPrc << "_";
+        result << "IS=" << CommonTestUtils::partialShape2str({inputShape.first}) << "_" << "TS=(";
         for (const auto& shape : inputShape.second) {
             result << CommonTestUtils::vec2str(shape) << "_";
         }
 
-        result << ")_axis=" << axis << "_mode=" << mode << "_sortType=" << sortType;
+        result << CPUTestsBase::getTestCaseName(cpuParams);
+
+        if (!additionalConfig.empty()) {
+            result << "_PluginConf";
+            for (auto& item : additionalConfig) {
+                if (item.second == PluginConfigParams::YES)
+                    result << "_" << item.first << "=" << item.second;
+            }
+        }
+
         return result.str();
     }
 
 protected:
+    void SetUp() override {
+        targetDevice = CommonTestUtils::DEVICE_CPU;
+
+        basicTopKParams basicParamsSet;
+        CPUSpecificParams cpuParams;
+        std::map<std::string, std::string> additionalConfig;
+        std::tie(basicParamsSet, cpuParams, additionalConfig) = this->GetParam();
+
+        std::tie(inFmts, outFmts, priority, selectedType) = cpuParams;
+
+        int64_t keepK;
+        ngraph::opset4::TopK::Mode mode;
+        ngraph::opset4::TopK::SortType sort;
+        ElementType inPrc, outPrc;
+        InputShape inputShape;
+        std::tie(keepK, axis, mode, sort, netPrecision, inPrc, outPrc, inputShape) = basicParamsSet;
+
+        if (additionalConfig[PluginConfigParams::KEY_ENFORCE_BF16] == PluginConfigParams::YES)
+            inPrc = outPrc = netPrecision = ElementType::bf16;
+        else
+            inPrc = outPrc = netPrecision;
+        configuration.insert(additionalConfig.begin(), additionalConfig.end());
+
+        selectedType = getPrimitiveType() + "_" + InferenceEngine::details::convertPrecision(netPrecision).name();
+
+        staticShape = inputShape.first.rank() == 0;
+        if (staticShape) {
+            init_input_shapes({inputShape});
+        } else {
+            inputDynamicShapes = {inputShape.first, {}};
+            for (size_t i = 0; i < inputShape.second.size(); ++i) {
+                targetStaticShapes.push_back({inputShape.second[i], {}});
+            }
+        }
+
+        auto params = ngraph::builder::makeDynamicParams(netPrecision, {inputDynamicShapes[0]});
+
+        // static shape need specific const k to test different sorting algorithms, dynamic shape tests random param k
+        std::shared_ptr<ngraph::opset4::TopK> topk;
+        if (staticShape) {
+            auto k = std::make_shared<ngraph::opset3::Constant>(ngraph::element::Type_t::i64, ngraph::Shape{}, &keepK);
+            topk = std::dynamic_pointer_cast<ngraph::opset4::TopK>(
+                    std::make_shared<ngraph::opset4::TopK>(params[0], k, axis, mode, sort));
+        } else {
+            auto k = std::make_shared<ngraph::opset3::Parameter>(ngraph::element::Type_t::i64, inputDynamicShapes[1]);
+            params.push_back(k);
+            topk = std::dynamic_pointer_cast<ngraph::opset4::TopK>(
+                    std::make_shared<ngraph::opset4::TopK>(params[0], k, axis, mode, sort));
+        }
+
+        topk->get_rt_info() = getCPUInfo();
+
+        ngraph::ResultVector results;
+        for (size_t i = 0; i < topk->get_output_size(); i++) {
+            results.push_back(std::make_shared<ngraph::opset4::Result>(topk->output(i)));
+        }
+
+        function = std::make_shared<ngraph::Function>(results, params, "TopK");
+    }
+
     void generate_inputs(const std::vector<ngraph::Shape>& targetInputStaticShapes) override {
         inputs.clear();
         const auto& funcInputs = function->inputs();
 
-        ov::Tensor data_tensor;
-        const auto& dataPrecision = funcInputs[0].get_element_type();
-        const auto& dataShape = targetInputStaticShapes.front();
-        if (funcInputs[0].get_element_type().is_real()) {
-            data_tensor = ov::test::utils::create_and_fill_tensor(dataPrecision, dataShape, 10, 0, 1000);
+        // Spec TopK_3.md allows to use unstable sorting, thus generate unreapeated input data to avoid a. and b.
+        // a. Skip comparing of index results, because an element in actual index tensor can be different with
+        //    its counterpart in expected index tensor
+        // b. If SortType is SORT_INDICES or NONE, the test program still needs to apply std::sort for all pairs
+        //    of 1xk value vectors in expected and actual output tensor before comparing them
+        auto shape = targetInputStaticShapes.front();
+        ov::Tensor tensor;
+        tensor = ov::test::utils::create_and_fill_tensor(funcInputs[0].get_element_type(), shape);
+        size_t size = tensor.get_size();
+
+        if (netPrecision == ElementType::f32) {
+            std::vector<int> data(size);
+            int start = - static_cast<int>(size / 2);
+            std::iota(data.begin(), data.end(), start);
+            std::mt19937 gen(0);
+            std::shuffle(data.begin(), data.end(), gen);
+
+            auto *rawBlobDataPtr = static_cast<float *>(tensor.data());
+            for (size_t i = 0; i < size; ++i) {
+                rawBlobDataPtr[i] = static_cast<float>(data[i]);
+            }
+        } else if (netPrecision == ElementType::bf16) {
+            size_t O = 1, A = 1, I = 1;
+            A = shape[axis];
+            for (size_t i = 0; i < axis; i++)
+                O *= shape[i];
+            for (size_t i = axis + 1; i < shape.size(); i++)
+                I *= shape[i];
+            if (O * A * I != size)
+                FAIL() << "Incorrect blob shape " << shape;
+
+            auto *rawBlobDataPtr = static_cast<ngraph::bfloat16 *>(tensor.data());
+            for (size_t o = 0; o < O; o++) {
+                for (size_t i = 0; i < I; i++) {
+                    std::vector<int> data(A);
+                    int start = - static_cast<int>(A / 2);
+                    std::iota(data.begin(), data.end(), start);
+                    const size_t seed = (o + 1) * (i + 1);
+                    std::mt19937 gen(seed);
+                    std::shuffle(data.begin(), data.end(), gen);
+                    for (size_t a = 0; a < A; a++) {
+                        rawBlobDataPtr[o * A * I + a * I + i] = static_cast<ngraph::bfloat16>(data[a]);
+                    }
+                }
+            }
         } else {
-            data_tensor = ov::test::utils::create_and_fill_tensor(dataPrecision, dataShape);
+            FAIL() << "generate_inputs for " << netPrecision << " precision isn't supported";
         }
+        inputs.insert({funcInputs[0].get_node_shared_ptr(), tensor});
 
-        const auto axis = std::get<2>(this->GetParam());
-        const auto& kPrecision = funcInputs[1].get_element_type();
-        const auto& kShape = targetInputStaticShapes[1];
+        if (!staticShape) {
+            const auto& kPrecision = funcInputs[1].get_element_type();
+            const auto& kShape = targetInputStaticShapes[1];
 
-        const size_t startFrom = 1;
-        const size_t range = targetInputStaticShapes[0][axis] - 1;
-        const size_t seed = inferRequestNum++;
-        const auto kTensor = ov::test::utils::create_and_fill_tensor(kPrecision, kShape, range, startFrom, 1, seed);
+            const size_t startFrom = 1;
+            const size_t range = targetInputStaticShapes[0][axis] - 1;
+            const size_t seed = inferRequestNum++;
+            const auto kTensor = ov::test::utils::create_and_fill_tensor(kPrecision, kShape, range, startFrom, 1, seed);
 
-        inputs.insert({ funcInputs[0].get_node_shared_ptr(), data_tensor });
-        inputs.insert({ funcInputs[1].get_node_shared_ptr(), kTensor });
-    }
-
-    void SetUp() override {
-        targetDevice = CommonTestUtils::DEVICE_CPU;
-        ngraph::element::Type inputPrecision;
-        InputShape inputShape;
-        std::int64_t axis;
-        ngraph::opset4::TopK::Mode mode;
-        ngraph::opset4::TopK::SortType sortType;
-        std::tie(inputPrecision, inputShape, axis, mode, sortType) = this->GetParam();
-
-        inputDynamicShapes = { inputShape.first, {} };
-        for (size_t i = 0; i < inputShape.second.size(); ++i) {
-            targetStaticShapes.push_back({ inputShape.second[i], {} });
+            inputs.insert({funcInputs[1].get_node_shared_ptr(), kTensor});
         }
-
-        selectedType = makeSelectedTypeStr("ref_any", inputPrecision);
-
-        auto params = ngraph::builder::makeDynamicParams(inputPrecision, { inputDynamicShapes[0] });
-        auto k_param = std::make_shared<ngraph::opset1::Parameter>(ngraph::element::i32, inputDynamicShapes[1]);
-        params.push_back(k_param);
-        auto topk = std::make_shared<ngraph::opset1::TopK>(params[0], k_param, axis, mode, sortType);
-
-        ngraph::ResultVector results;
-        for (size_t i = 0; i < topk->get_output_size(); ++i) {
-            results.push_back(std::make_shared<ngraph::opset1::Result>(topk->output(i)));
-        }
-
-        function = std::make_shared<ngraph::Function>(results, params, "TopKLayerCPUTest");
-        functionRefs = ngraph::clone_function(*function);
     }
 
 private:
+    int64_t axis;
     size_t inferRequestNum = 0;
+    ElementType netPrecision;
+    bool staticShape;
 };
 
 TEST_P(TopKLayerCPUTest, CompareWithRefs) {
@@ -110,11 +215,19 @@ TEST_P(TopKLayerCPUTest, CompareWithRefs) {
     CheckPluginRelatedResults(executableNetwork, "TopK");
 }
 
-const ngraph::element::TypeVector inputPrecisions = {
-    ngraph::element::f32,
+namespace {
+
+const std::vector<ElementType> netPrecisions = {
+    ElementType::f32,
 };
 
-const std::vector<int64_t> axes = { 0, 1, 2 };
+std::vector<std::map<std::string, std::string>> additionalConfig = {
+    {{PluginConfigParams::KEY_ENFORCE_BF16, PluginConfigParams::NO}},
+    {{PluginConfigParams::KEY_ENFORCE_BF16, PluginConfigParams::YES}}
+};
+
+const std::vector<int64_t> axes = {0, 1, 2, 3};
+const std::vector<int64_t> k = {1, 5, 7, 18, 21};
 
 const std::vector<ngraph::opset4::TopK::Mode> modes = {
     ngraph::opset4::TopK::Mode::MIN,
@@ -122,51 +235,54 @@ const std::vector<ngraph::opset4::TopK::Mode> modes = {
 };
 
 const std::vector<ngraph::opset4::TopK::SortType> sortTypes = {
-    ngraph::opset4::TopK::SortType::SORT_INDICES,
     ngraph::opset4::TopK::SortType::SORT_VALUES,
+    ngraph::opset4::TopK::SortType::SORT_INDICES,
 };
 
-const std::vector<InputShape> inShapes = {
-    InputShape{
-        // dynamic
-        {-1, -1, -1},
-        // target
-        {
-            {10, 13, 11},
-            {15, 10, 5},
-            {5, 5, 8}
-        }
-    },
-    InputShape{
-        // dynamic
-        {-1, { 7, 25 }, -1},
-        // target
-        {
-            {10, 7, 10},
-            {15, 20, 5},
-            {5, 25, 7}
-        }
-    },
-    InputShape{
-        // dynamic
-        {{8, 10}, {7, 17}, {12, 27}},
-        // target
-        {
-            {8, 7, 12},
-            {9, 12, 25},
-            {10, 17, 27}
-        }
-    },
+std::vector<ov::test::InputShape> inputShapes = {
+    {{}, {{21, 21, 21, 21}}},
 };
 
-const auto testCases = ::testing::Combine(
-    ::testing::ValuesIn(inputPrecisions),
-    ::testing::ValuesIn(inShapes),
-    ::testing::ValuesIn(axes),
-    ::testing::ValuesIn(modes),
-    ::testing::ValuesIn(sortTypes)
-);
+std::vector<ov::test::InputShape> inputShapesDynamic = {
+    {{21, {20, 25}, 21, {20, 25}}, {{21, 21, 21, 21}, {21, 22, 21, 23}}}
+};
 
-INSTANTIATE_TEST_SUITE_P(smoke_CompareWithRefs, TopKLayerCPUTest, testCases, TopKLayerCPUTest::getTestCaseName);
+std::vector<CPUSpecificParams> cpuParams = {
+    CPUSpecificParams({nChw16c, x}, {nChw16c, nChw16c}, {}, {}),
+    CPUSpecificParams({nchw, x}, {nchw, nchw}, {}, {}),
+    CPUSpecificParams({nhwc, x}, {nhwc, nhwc}, {}, {})
+};
+
+INSTANTIATE_TEST_CASE_P(smoke_TopK, TopKLayerCPUTest,
+    ::testing::Combine(
+        ::testing::Combine(
+            ::testing::ValuesIn(k),
+            ::testing::ValuesIn(axes),
+            ::testing::ValuesIn(modes),
+            ::testing::ValuesIn(sortTypes),
+            ::testing::ValuesIn(netPrecisions),
+            ::testing::Values(ElementType::undefined),
+            ::testing::Values(ElementType::undefined),
+            ::testing::ValuesIn(inputShapes)),
+        ::testing::ValuesIn(filterCPUSpecificParams(cpuParams)),
+        ::testing::ValuesIn(additionalConfig)),
+    TopKLayerCPUTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_TopK_dynamic, TopKLayerCPUTest,
+    ::testing::Combine(
+        ::testing::Combine(
+            ::testing::Values(1),
+            ::testing::ValuesIn(axes),
+            ::testing::ValuesIn(modes),
+            ::testing::ValuesIn(sortTypes),
+            ::testing::ValuesIn(netPrecisions),
+            ::testing::Values(ElementType::undefined),
+            ::testing::Values(ElementType::undefined),
+            ::testing::ValuesIn(inputShapesDynamic)),
+        ::testing::ValuesIn(filterCPUSpecificParams(cpuParams)),
+        ::testing::ValuesIn(additionalConfig)),
+    TopKLayerCPUTest::getTestCaseName);
+
+} // namespace
 
 } // namespace CPULayerTestsDefinitions

src/tests/functional/shared_test_classes/include/shared_test_classes/single_layer/topk.hpp

@@ -28,6 +28,7 @@ class TopKLayerTest : public testing::WithParamInterface<TopKParams>,
                       virtual public LayerTestsUtils::LayerTestsCommon {
 public:
     static std::string getTestCaseName(const testing::TestParamInfo<TopKParams>& obj);
+    InferenceEngine::Blob::Ptr GenerateInput(const InferenceEngine::InputInfo &info) const override;
 
 protected:
     void SetUp() override;

src/tests/functional/shared_test_classes/src/single_layer/topk.cpp

@@ -52,4 +52,40 @@ void TopKLayerTest::SetUp() {
     }
     function = std::make_shared<ngraph::Function>(results, params, "TopK");
 }
+
+InferenceEngine::Blob::Ptr TopKLayerTest::GenerateInput(const InferenceEngine::InputInfo &info) const {
+    IE_ASSERT(InferenceEngine::Precision::FP32 == info.getTensorDesc().getPrecision()
+           || InferenceEngine::Precision::BF16 == info.getTensorDesc().getPrecision()
+           || InferenceEngine::Precision::FP16 == info.getTensorDesc().getPrecision());
+
+    InferenceEngine::Blob::Ptr blob = make_blob_with_precision(info.getTensorDesc());
+    blob->allocate();
+
+    size_t size = blob->size();
+    int start = - static_cast<int>(size / 2);
+    std::vector<int> data(size);
+    std::iota(data.begin(), data.end(), start);
+    std::mt19937 gen(0);
+    std::shuffle(data.begin(), data.end(), gen);
+
+    float divisor = size / 10.0;
+    if (InferenceEngine::Precision::FP32 == info.getTensorDesc().getPrecision()) {
+        auto *rawBlobDataPtr = blob->buffer().as<float *>();
+        for (size_t i = 0; i < size; i++) {
+            rawBlobDataPtr[i] = static_cast<float>(data[i] / divisor);
+        }
+    } else if (InferenceEngine::Precision::BF16 == info.getTensorDesc().getPrecision()) {
+        auto *rawBlobDataPtr = blob->buffer().as<ngraph::bfloat16 *>();
+        for (size_t i = 0; i < size; i++) {
+            rawBlobDataPtr[i] = static_cast<ngraph::bfloat16>(data[i] / divisor);
+        }
+    } else if (InferenceEngine::Precision::FP16 == info.getTensorDesc().getPrecision()) {
+        auto *rawBlobDataPtr = blob->buffer().as<ngraph::float16 *>();
+        for (size_t i = 0; i < size; i++) {
+            rawBlobDataPtr[i] = static_cast<ngraph::float16>(data[i] / divisor);
+        }
+    }
+
+    return blob;
+}
 }  // namespace LayerTestsDefinitions