[CPU] Enable CPU plugin cache for ShuffleChannels and SpaceToDepth (#9556)

* Space to depth.

* Shuffle Channels RT cache.

* Update

* Apllied reivew comments.

* Applied review comments.

src/plugins/intel_cpu/src/nodes/mkldnn_shuffle_channels_node.cpp

@@ -14,6 +14,7 @@
 
 #include <string>
 #include <cmath>
+#include <common/primitive_hashing_utils.hpp>
 
 #define THROW_SHCH_ERROR IE_THROW() << "ShuffleChannels layer with name '" << getName() << "' "
 
@@ -23,6 +24,31 @@ using namespace InferenceEngine;
 using namespace mkldnn::impl;
 using namespace mkldnn::impl::cpu::x64;
 
+size_t MKLDNNShuffleChannelsNode::ShuffleChannelsAttributes::hash() const {
+    using namespace dnnl::impl;
+    using namespace dnnl::impl::primitive_hashing;
+
+    size_t seed = 0;
+    seed = hash_combine(seed, layoutType);
+    seed = hash_combine(seed, dataRank);
+    seed = hash_combine(seed, axis);
+    seed = hash_combine(seed, spatialRank);
+    seed = hash_combine(seed, group);
+    seed = hash_combine(seed, dataSize);
+    seed = get_vector_hash(seed, srcDims);
+    seed = get_vector_hash(seed, srcBlockedDims);
+
+    return seed;
+}
+
+bool MKLDNNShuffleChannelsNode::ShuffleChannelsAttributes::operator==(const ShuffleChannelsAttributes& rhs) const {
+    bool result = layoutType == rhs.layoutType && dataRank == rhs.dataRank &&
+                  axis == rhs.axis && spatialRank == rhs.spatialRank &&
+                  group == rhs.group && dataSize == rhs.dataSize && srcDims == rhs.srcDims &&
+                  srcBlockedDims == rhs.srcBlockedDims;
+    return result;
+}
+
 bool MKLDNNShuffleChannelsNode::isSupportedOperation(const std::shared_ptr<const ngraph::Node>& op, std::string& errorMessage) noexcept {
     try {
         auto shuffleChannels = ov::as_type_ptr<const ngraph::op::v0::ShuffleChannels>(op);
@@ -123,17 +149,29 @@ void MKLDNNShuffleChannelsNode::createPrimitive() {
 
 void MKLDNNShuffleChannelsNode::prepareParams() {
     auto& srcMemPtr = getParentEdgeAt(0)->getMemoryPtr();
-    execPtr = std::make_shared<ShuffleChannelsExecutor>(attrs, srcMemPtr->getStaticDims(), srcMemPtr->GetDescWithType<BlockedMemoryDesc>()->getBlockDims());
+    auto builder = [](const ShuffleChannelsAttributes& key) -> std::shared_ptr<ShuffleChannelsExecutor> {
+        return std::make_shared<ShuffleChannelsExecutor>(key);
+    };
+    attrs.srcDims = srcMemPtr->getStaticDims();
+    attrs.srcBlockedDims = srcMemPtr->GetDescWithType<BlockedMemoryDesc>()->getBlockDims();
+
+    auto cache = getRuntimeCache();
+    auto result = cache->getOrCreate(attrs, builder);
+    if (!result.first) {
+        IE_THROW() << "ShuffleChannelsExecutor was not found for node " << getName() << ".";
+    }
+
+    execPtr = result.first;
 }
 
-MKLDNNShuffleChannelsNode::ShuffleChannelsExecutor::ShuffleChannelsExecutor(const ShuffleChannelsAttributes& attrs,
-                                                                            const VectorDims& srcDims,
-                                                                            const VectorDims& srcBlockedDims) {
+MKLDNNShuffleChannelsNode::ShuffleChannelsExecutor::ShuffleChannelsExecutor(const ShuffleChannelsAttributes& attrs) {
     if (!one_of(attrs.layoutType, LayoutType::nCsp16c, LayoutType::nCsp8c, LayoutType::nspc, LayoutType::ncsp))
         IE_THROW() << "ShuffleChannels executor supports only 'nCsp16c', 'nCsp8c', 'nspc' or 'ncsp' layouts.";
 
     const bool isBlocked = MKLDNNPlugin::one_of(attrs.layoutType, LayoutType::nCsp16c, LayoutType::nCsp8c);
     const bool isChannelsLast = attrs.layoutType == LayoutType::nspc;
+    const auto& srcDims = attrs.srcDims;
+    const auto& srcBlockedDims = attrs.srcBlockedDims;
 
     // 2 for decomposed axis dim, 1 for composed spatial dim
     const int batchRank = attrs.axis;

src/plugins/intel_cpu/src/nodes/mkldnn_shuffle_channels_node.h

@@ -26,11 +26,6 @@ public:
     bool created() const override;
 
     void prepareParams() override;
-
-protected:
-    void executeDynamicImpl(mkldnn::stream strm) override;
-
-private:
     struct ShuffleChannelsAttributes {
         LayoutType layoutType;
         int dataRank = 0;
@@ -38,10 +33,20 @@ private:
         int spatialRank = 0;
         size_t group = 0lu;
         size_t dataSize = 1lu;
-    } attrs;
+        VectorDims srcDims;
+        VectorDims srcBlockedDims;
+        size_t hash() const;
+        bool operator==(const ShuffleChannelsAttributes& rhs) const;
+    };
+
+protected:
+    void executeDynamicImpl(mkldnn::stream strm) override;
+
+private:
+    ShuffleChannelsAttributes attrs;
 
     struct ShuffleChannelsExecutor final {
-        ShuffleChannelsExecutor(const ShuffleChannelsAttributes& attrs, const VectorDims& srcDims, const VectorDims& srcBlockedDims);
+        ShuffleChannelsExecutor(const ShuffleChannelsAttributes& attrs);
         void exec(const uint8_t* srcData, uint8_t* dstData, const int MB);
         ~ShuffleChannelsExecutor() = default;
 

src/plugins/intel_cpu/src/nodes/mkldnn_space_to_depth_node.cpp

@@ -8,6 +8,7 @@
 #include <utils/general_utils.h>
 
 #include <cmath>
+#include <common/primitive_hashing_utils.hpp>
 #include <cpu/x64/jit_generator.hpp>
 #include <ngraph/opsets/opset1.hpp>
 #include <string>
@@ -21,6 +22,32 @@ using namespace InferenceEngine;
 using namespace mkldnn;
 using namespace mkldnn::impl;
 
+size_t MKLDNNSpaceToDepthNode::SpaceToDepthAttrs::hash() const {
+    using namespace dnnl::impl;
+    using namespace dnnl::impl::primitive_hashing;
+
+    size_t seed = 0;
+    seed = hash_combine(seed, layoutType);
+    seed = hash_combine(seed, mode);
+    seed = hash_combine(seed, blockSize);
+    seed = hash_combine(seed, blockStep);
+    seed = hash_combine(seed, dataSize);
+    seed = hash_combine(seed, nSpatialDims);
+    seed = get_vector_hash(seed, srcBlockedDims);
+    seed = get_vector_hash(seed, destBlockedDims);
+
+    return seed;
+}
+
+bool MKLDNNSpaceToDepthNode::SpaceToDepthAttrs::operator==(const SpaceToDepthAttrs& rhs) const {
+    bool result = layoutType == rhs.layoutType && mode == rhs.mode &&
+                  blockSize == rhs.blockSize && blockStep == rhs.blockStep &&
+                  dataSize == rhs.dataSize && nSpatialDims == rhs.nSpatialDims &&
+                  srcBlockedDims == rhs.srcBlockedDims && destBlockedDims == rhs.destBlockedDims;
+
+    return result;
+}
+
 bool MKLDNNSpaceToDepthNode::isSupportedOperation(const std::shared_ptr<const ngraph::Node>& op,
                                                   std::string& errorMessage) noexcept {
     try {
@@ -162,24 +189,36 @@ void MKLDNNSpaceToDepthNode::createPrimitive() {
 }
 
 void MKLDNNSpaceToDepthNode::prepareParams() {
-    const VectorDims& srcBlockedDims = getParentEdgeAt(0)->getMemoryPtr()->GetDescWithType<BlockedMemoryDesc>()->getBlockDims();
-    const VectorDims& dstBlockedDims = getChildEdgeAt(0)->getMemoryPtr()->GetDescWithType<BlockedMemoryDesc>()->getBlockDims();
-    execPtr = std::make_shared<SpaceToDepthExecutor>(attrs, srcBlockedDims, dstBlockedDims);
+    attrs.srcBlockedDims =
+        getParentEdgeAt(0)->getMemoryPtr()->GetDescWithType<BlockedMemoryDesc>()->getBlockDims();
+    attrs.destBlockedDims =
+        getChildEdgeAt(0)->getMemoryPtr()->GetDescWithType<BlockedMemoryDesc>()->getBlockDims();
+    auto builder = [](const SpaceToDepthAttrs& key) -> std::shared_ptr<SpaceToDepthExecutor> {
+        return std::make_shared<SpaceToDepthExecutor>(key);
+    };
+
+    auto cache = getRuntimeCache();
+    auto result = cache->getOrCreate(attrs, builder);
+    if (!result.first) {
+        IE_THROW() << "SpaceToDepthExecutor was not found for node " << getName() << ".";
+    }
+
+    execPtr = result.first;
 }
 
-MKLDNNSpaceToDepthNode::SpaceToDepthExecutor::SpaceToDepthExecutor(const SpaceToDepthAttrs& attrs,
-                                                                   const VectorDims& srcBlockedDims,
-                                                                   const VectorDims& dstBlockedDims) {
+MKLDNNSpaceToDepthNode::SpaceToDepthExecutor::SpaceToDepthExecutor(const SpaceToDepthAttrs& attrs) {
     if (!MKLDNNPlugin::one_of(attrs.layoutType,
                               LayoutType::nCsp16c,
                               LayoutType::nCsp8c,
                               LayoutType::nspc,
                               LayoutType::ncsp))
-        IE_THROW() << "DepthToSpace executor supports only 'nCsp16c', 'nCsp8c', "
+        IE_THROW() << "SpaceToDepth executor supports only 'nCsp16c', 'nCsp8c', "
                       "'nspc' or 'ncsp' layouts.";
 
     const bool isBlocked = MKLDNNPlugin::one_of(attrs.layoutType, LayoutType::nCsp16c, LayoutType::nCsp8c);
     const bool isChannelsFirst = attrs.layoutType == LayoutType::nspc;
+    const auto& srcBlockedDims = attrs.srcBlockedDims;
+    const auto& dstBlockedDims = attrs.destBlockedDims;
 
     size_t nDims = srcBlockedDims.size();
 

src/plugins/intel_cpu/src/nodes/mkldnn_space_to_depth_node.h

@@ -24,14 +24,7 @@ public:
 
     void prepareParams() override;
 
-protected:
-    void executeDynamicImpl(mkldnn::stream strm) override;
-
-private:
-    enum Mode {
-        BLOCKS_FIRST = 0,
-        DEPTH_FIRST = 1
-    };
+    enum Mode { BLOCKS_FIRST = 0, DEPTH_FIRST = 1 };
 
     struct SpaceToDepthAttrs {
         LayoutType layoutType;
@@ -40,12 +33,20 @@ private:
         size_t blockStep = 1lu;
         size_t dataSize = 1lu;
         size_t nSpatialDims = 0lu;
-    } attrs;
+        VectorDims srcBlockedDims;
+        VectorDims destBlockedDims;
+        size_t hash() const;
+        bool operator==(const SpaceToDepthAttrs& rhs) const;
+    };
+
+protected:
+    void executeDynamicImpl(mkldnn::stream strm) override;
+
+private:
+    SpaceToDepthAttrs attrs;
 
     struct SpaceToDepthExecutor final {
-        SpaceToDepthExecutor(const SpaceToDepthAttrs& attrs,
-                             const VectorDims& srcBlockedDims,
-                             const VectorDims& dstBlockedDims);
+        SpaceToDepthExecutor(const SpaceToDepthAttrs& attrs);
         void exec(const uint8_t* srcData, uint8_t* dstData, const int MB);
         ~SpaceToDepthExecutor() = default;
 

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

@@ -169,24 +169,24 @@ const auto shuffleChannelsParams5DBlock = ::testing::Combine(
 
 const std::vector<InputShape> inputShapesDynamic4D = {
         {{-1, -1, -1, -1},
-         {{8, 4, 4, 4}, {8, 16, 8, 4}}},
+         {{8, 4, 4, 4}, {8, 16, 8, 4}, {8, 4, 4, 4}}},
 
         {{-1, 8, -1, -1},
-         {{8, 8, 8, 8}, {8, 8, 4, 16}}},
+         {{8, 8, 8, 8}, {8, 8, 4, 16}, {8, 8, 8, 8}}},
 
         {{{4, 32}, {4, 32}, {4, 32}, {4, 32}},
-         {{4, 12, 8, 8}, {8, 32, 12, 4}}},
+         {{4, 12, 8, 8}, {8, 32, 12, 4}, {4, 12, 8, 8}}},
 };
 
 const std::vector<InputShape> inputShapesDynamic5D = {
         {{-1, -1, -1, -1, -1},
-         {{6, 6, 6, 6, 6}, {12, 6, 12, 12, 12}}},
+         {{6, 6, 6, 6, 6}, {12, 6, 12, 12, 12}, {6, 6, 6, 6, 6}}},
 
         {{-1, 18, -1, -1, -1},
-         {{6, 18, 12, 6, 12}, {6, 18, 6, 6, 6}}},
+         {{6, 18, 12, 6, 12}, {6, 18, 6, 6, 6}, {6, 18, 12, 6, 12}}},
 
         {{{6, 24}, {6, 24}, {6, 24}, {6, 24}, {6, 24}},
-         {{24, 12, 6, 6, 6}, {12, 24, 6, 12, 12}}},
+         {{24, 12, 6, 6, 6}, {12, 24, 6, 12, 12}, {24, 12, 6, 6, 6}}},
 };
 
 INSTANTIATE_TEST_SUITE_P(smoke_ShuffleChannelsStatic4D, ShuffleChannelsLayerCPUTest,

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

@@ -195,30 +195,30 @@ INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthStaticBS3_5D, SpaceToDepthLayerCPU
 namespace dynamic_shapes {
 
 const std::vector<InputShape> inputShapes4D = {
-        {{-1, -1, -1 , -1},                                  // dynamic
-         {{2, 12, 6, 12}, {1, 2, 18, 6}, {1, 10, 12, 24}}},     // target
+    {{-1, -1, -1, -1},                                                   // dynamic
+     {{2, 12, 6, 12}, {1, 2, 18, 6}, {1, 10, 12, 24}, {2, 12, 6, 12}}},  // target
 
-        {{-1, 32, -1 , -1},                                 // dynamic
-         {{1, 32, 6, 12}, {2, 32, 24, 24}, {3, 32, 30, 6}}},  // target
+    {{-1, 32, -1, -1},                                                     // dynamic
+     {{1, 32, 6, 12}, {2, 32, 24, 24}, {3, 32, 30, 6}, {2, 32, 24, 24}}},  // target
 
-        {{{2, 5}, {1, 50}, {1, 100}, {1, 100}},               // dynamic
-         {{3, 5, 12, 36}, {2, 3, 6, 12}, {5, 2, 12, 18}}},   // target
+    {{{2, 5}, {1, 50}, {1, 100}, {1, 100}},                            // dynamic
+     {{3, 5, 12, 36}, {2, 3, 6, 12}, {5, 2, 12, 18}, {2, 3, 6, 12}}},  // target
 };
 
 const std::vector<InputShape> inputShapes5D = {
-        {{-1, -1, -1, -1, -1},                                        // dynamic
-         {{2, 2, 6, 12, 24}, {1, 4, 24, 24, 36}, {1, 7, 6, 30, 18}}},  // target
+    {{-1, -1, -1, -1, -1},                                                            // dynamic
+     {{2, 2, 6, 12, 24}, {1, 4, 24, 24, 36}, {1, 7, 6, 30, 18}, {2, 2, 6, 12, 24}}},  // target
 
-        {{{1, 3}, {5, 16}, {1, 60}, {1, 60}, {1, 60}},                // dynamic
-         {{3, 5, 12, 6, 24}, {1, 6, 24, 18, 6}}},               // target
+    {{{1, 3}, {5, 16}, {1, 60}, {1, 60}, {1, 60}},                // dynamic
+     {{3, 5, 12, 6, 24}, {1, 6, 24, 18, 6}, {3, 5, 12, 6, 24}}},  // target
 };
 
 const std::vector<InputShape> inputShapesBlocked5D = {
-        {{-1, 16, -1, -1, -1},                                          // dynamic
-         {{1, 16, 4, 6, 10}, {1, 16, 12, 8, 2}, {3, 16, 2, 14, 24}}},     // target
+    {{-1, 16, -1, -1, -1},                                                            // dynamic
+     {{1, 16, 4, 6, 10}, {1, 16, 12, 8, 2}, {3, 16, 2, 14, 24}, {1, 16, 12, 8, 2}}},  // target
 
-        {{{1, 3}, 32, {1, 32}, {1, 32}, {1, 32}},                          // dynamic
-         {{1, 32, 4, 16, 10}, {1, 32, 18, 6, 14}, {3, 32, 2, 14, 12}}},     // target
+    {{{1, 3}, 32, {1, 32}, {1, 32}, {1, 32}},                                            // dynamic
+     {{1, 32, 4, 16, 10}, {1, 32, 18, 6, 14}, {3, 32, 2, 14, 12}, {1, 32, 18, 6, 14}}},  // target
 };
 
 INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamic4D, SpaceToDepthLayerCPUTest,