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

* Space to depth.

* Shuffle Channels RT cache.

* Update

* Apllied reivew comments.

* Applied review comments.
This commit is contained in:
Luwei Zhou authored and GitHub committed 2022-01-17 12:28:30 +03:00
1 parent 0ed3da0e86
commit a76483ddb2
6 files changed
+133 -50

No files matched your search

@@ -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;
@@ -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;
@@ -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();
@@ -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;