[CPU] Optimize concat node for dynamic shapes (#13019)
This commit is contained in:
1 parent
b44eb83551
commit
b8ada02cba
3 files changed
+232
-34
No files matched your search
@@ -640,4 +640,64 @@ void parallel_for5d(const T0& D0, const T1& D1, const T2& D2, const T3& D3, cons
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename F>
|
||||
void for_6d(const int& ithr,
|
||||
const int& nthr,
|
||||
const T0& D0,
|
||||
const T1& D1,
|
||||
const T2& D2,
|
||||
const T3& D3,
|
||||
const T4& D4,
|
||||
const T5& D5,
|
||||
const F& func) {
|
||||
const size_t work_amount = (size_t)D0 * D1 * D2 * D3 * D4 * D5;
|
||||
if (work_amount == 0)
|
||||
return;
|
||||
size_t start{0}, end{0};
|
||||
splitter(work_amount, nthr, ithr, start, end);
|
||||
|
||||
T0 d0{0};
|
||||
T1 d1{0};
|
||||
T2 d2{0};
|
||||
T3 d3{0};
|
||||
T4 d4{0};
|
||||
T5 d5{0};
|
||||
parallel_it_init(start, d0, D0, d1, D1, d2, D2, d3, D3, d4, D4, d5, D5);
|
||||
for (size_t iwork = start; iwork < end; ++iwork) {
|
||||
details::call_with_args(func, ithr, iwork, d0, d1, d2, d3, d4, d5);
|
||||
parallel_it_step(d0, D0, d1, D1, d2, D2, d3, D3, d4, D4, d5, D5);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename F>
|
||||
void parallel_for6d(const T0& D0, const T1& D1, const T2& D2, const T3& D3, const T4& D4, const T5& D5, const F& func) {
|
||||
#if IE_THREAD == IE_THREAD_TBB
|
||||
auto work_amount = static_cast<size_t>(D0 * D1 * D2 * D3 * D4 * D5);
|
||||
int nthr = parallel_get_max_threads();
|
||||
if (static_cast<size_t>(nthr) > work_amount)
|
||||
nthr = static_cast<int>(work_amount);
|
||||
if (nthr == 1) {
|
||||
for_6d(0, 1, D0, D1, D2, D3, D4, D5, func);
|
||||
} else {
|
||||
tbb::parallel_for(
|
||||
0,
|
||||
nthr,
|
||||
[&](int ithr) {
|
||||
for_6d(ithr, nthr, D0, D1, D2, D3, D4, D5, func);
|
||||
},
|
||||
tbb::static_partitioner());
|
||||
}
|
||||
#elif IE_THREAD == IE_THREAD_TBB_AUTO
|
||||
const int nthr = parallel_get_max_threads();
|
||||
tbb::parallel_for(0, nthr, [&](int ithr) {
|
||||
for_6d(ithr, nthr, D0, D1, D2, D3, D4, D5, func);
|
||||
});
|
||||
#elif IE_THREAD == IE_THREAD_OMP
|
||||
# pragma omp parallel
|
||||
for_6d(parallel_get_thread_num(), parallel_get_num_threads(), D0, D1, D2, D3, D4, D5, func);
|
||||
#elif IE_THREAD == IE_THREAD_SEQ
|
||||
for_6d(0, 1, D0, D1, D2, D3, D4, D5, func);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace InferenceEngine
|
||||
@@ -22,7 +22,6 @@
|
||||
#include "common/cpu_memcpy.h"
|
||||
#include "common/blocked_desc_creator.h"
|
||||
#include <memory_desc/cpu_memory_desc_utils.h>
|
||||
|
||||
using namespace dnnl;
|
||||
using namespace InferenceEngine;
|
||||
|
||||
@@ -58,6 +57,7 @@ Concat::Concat(const std::shared_ptr<ngraph::Node>& op, const dnnl::engine& eng,
|
||||
}
|
||||
|
||||
const auto inRank = getInputShapeAtPort(0).getRank();
|
||||
canExecRef = inRank <= 6;
|
||||
auto concatOp = ngraph::as_type_ptr<ngraph::op::v0::Concat>(op);
|
||||
auto axis = concatOp->get_axis();
|
||||
if (axis < 0) {
|
||||
@@ -94,6 +94,9 @@ void Concat::getSupportedDescriptors() {
|
||||
if (std::all_of(childDims.begin(), childDims.begin() + axis, [](size_t dim) { return dim == 1; }))
|
||||
canBeInPlace = true;
|
||||
}
|
||||
nelemToCopy.resize(getParentEdges().size(), 0);
|
||||
dstOffset.resize(getParentEdges().size());
|
||||
inputStrides.resize(getParentEdges().size());
|
||||
}
|
||||
|
||||
void Concat::initSupportedPrimitiveDescriptors() {
|
||||
@@ -359,42 +362,66 @@ void Concat::prepareParams() {
|
||||
return;
|
||||
|
||||
const auto& dstMemPtr = getChildEdgesAtPort(0)[0]->getMemoryPtr();
|
||||
auto dstMemDesc = dstMemPtr->GetDescWithType<BlockedMemoryDesc>();
|
||||
if (!dstMemPtr || !dstMemPtr->isAllocated())
|
||||
IE_THROW() << "Destination memory didn't allocate.";
|
||||
if (getSelectedPrimitiveDescriptor() == nullptr)
|
||||
IE_THROW() << "Preferable primitive descriptor is not set.";
|
||||
|
||||
const auto& outputStrides = dstMemDesc->getStrides();
|
||||
size_t curConcatOffset = 0;
|
||||
const size_t elemSize = DnnlExtensionUtils::sizeOfDataType(dstMemPtr->GetDataType());
|
||||
const auto& src0BlkMemDesc = getParentEdgesAtPort(0)[0]->getMemoryPtr()->getDescPtr()->as<BlockedMemoryDesc>();
|
||||
const auto& outputOrder = src0BlkMemDesc->getOrder();
|
||||
for (size_t i = 0; i < outputOrder.size(); i++) {
|
||||
if (outputOrder[i] == axis) {
|
||||
reorderedAxis = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
std::vector<memory::desc> srcs_d;
|
||||
for (size_t i = 0; i < getParentEdges().size(); i++) {
|
||||
const auto& srcMemPtr = getParentEdgesAtPort(i)[0]->getMemoryPtr();
|
||||
if (!srcMemPtr || !srcMemPtr->isAllocated()) {
|
||||
auto parent = getParentEdgeAt(i)->getParent();
|
||||
IE_THROW() << "Source memory from " << parent->getName() << " didn't allocate for node "
|
||||
<< getName() << ".";
|
||||
<< getName() << ".";
|
||||
}
|
||||
|
||||
if (srcMemPtr->GetShape().hasZeroDims()) {
|
||||
continue;
|
||||
if (canExecRef) {
|
||||
const auto srcMemDesc = srcMemPtr->getDescPtr()->as<BlockedMemoryDesc>();
|
||||
const auto& inputShape = srcMemDesc->getBlockDims();
|
||||
inputStrides[i] = srcMemDesc->getStrides();
|
||||
size_t nElem = 1;
|
||||
for (size_t j = reorderedAxis; j < inputShape.size(); j++) {
|
||||
nElem *= inputShape[j];
|
||||
}
|
||||
nelemToCopy[i] = nElem * elemSize;
|
||||
dstOffset[i] = outputStrides[reorderedAxis] * curConcatOffset * elemSize;
|
||||
curConcatOffset += inputShape[reorderedAxis];
|
||||
} else {
|
||||
if (srcMemPtr->GetShape().hasZeroDims()) {
|
||||
continue;
|
||||
}
|
||||
auto desc = srcMemPtr->GetDescWithType<DnnlMemoryDesc>()->getDnnlDesc();
|
||||
const auto& dims = srcMemPtr->getStaticDims();
|
||||
for (size_t j = 0; j < dims.size(); j++) {
|
||||
desc.data.dims[j] = dims[j];
|
||||
}
|
||||
srcs_d.emplace_back(desc);
|
||||
}
|
||||
|
||||
auto desc = srcMemPtr->GetDescWithType<DnnlMemoryDesc>()->getDnnlDesc();
|
||||
const auto& dims = srcMemPtr->getStaticDims();
|
||||
for (size_t j = 0; j < dims.size(); j++) {
|
||||
desc.data.dims[j] = dims[j];
|
||||
}
|
||||
|
||||
srcs_d.emplace_back(desc);
|
||||
}
|
||||
if (!canExecRef) {
|
||||
auto desc = dstMemPtr->GetDescWithType<DnnlMemoryDesc>()->getDnnlDesc();
|
||||
const auto& dims = dstMemPtr->getStaticDims();
|
||||
for (size_t i = 0; i < dims.size(); i++) {
|
||||
desc.data.dims[i] = dims[i];
|
||||
desc.data.padded_dims[i] = dims[i];
|
||||
}
|
||||
|
||||
auto desc = dstMemPtr->GetDescWithType<DnnlMemoryDesc>()->getDnnlDesc();
|
||||
const auto& dims = dstMemPtr->getStaticDims();
|
||||
for (size_t i = 0; i < dims.size(); i++) {
|
||||
desc.data.dims[i] = dims[i];
|
||||
desc.data.padded_dims[i] = dims[i];
|
||||
auto primitive_desc = concat::primitive_desc(desc, static_cast<int>(axis), srcs_d, getEngine());
|
||||
prim.reset(new concat(primitive_desc));
|
||||
}
|
||||
|
||||
auto primitive_desc = concat::primitive_desc(desc, static_cast<int>(axis), srcs_d, getEngine());
|
||||
prim.reset(new concat(primitive_desc));
|
||||
}
|
||||
|
||||
size_t Concat::inverseOrder(const SizeVector& order, size_t axis) {
|
||||
@@ -493,25 +520,28 @@ void Concat::execute(dnnl::stream strm) {
|
||||
return;
|
||||
}
|
||||
|
||||
const Memory& dst_memory = getChildEdgeAt(0)->getMemory();
|
||||
if (canOptimizeNspc) {
|
||||
execNspcSpecCase();
|
||||
return;
|
||||
}
|
||||
|
||||
const size_t num_src = getParentEdges().size();
|
||||
std::unordered_map<int, memory> mem_ags {{DNNL_ARG_DST, dst_memory.GetPrimitive()}};
|
||||
size_t nonZeroInShapes = 0;
|
||||
for (int i = 0; i < num_src; i++) {
|
||||
const auto& srcMem = getParentEdgesAtPort(i)[0]->getMemory();
|
||||
if (srcMem.GetShape().hasZeroDims()) {
|
||||
continue;
|
||||
if (canExecRef) {
|
||||
execRef();
|
||||
} else {
|
||||
const Memory& dst_memory = getChildEdgeAt(0)->getMemory();
|
||||
const size_t num_src = getParentEdges().size();
|
||||
std::unordered_map<int, memory> mem_ags {{DNNL_ARG_DST, dst_memory.GetPrimitive()}};
|
||||
size_t nonZeroInShapes = 0;
|
||||
for (int i = 0; i < num_src; i++) {
|
||||
const auto& srcMem = getParentEdgesAtPort(i)[0]->getMemory();
|
||||
if (srcMem.GetShape().hasZeroDims()) {
|
||||
continue;
|
||||
}
|
||||
mem_ags[DNNL_ARG_MULTIPLE_SRC + nonZeroInShapes] = srcMem.GetPrimitive();
|
||||
nonZeroInShapes++;
|
||||
}
|
||||
mem_ags[DNNL_ARG_MULTIPLE_SRC + nonZeroInShapes] = srcMem.GetPrimitive();
|
||||
nonZeroInShapes++;
|
||||
(*prim).execute(strm, mem_ags);
|
||||
}
|
||||
|
||||
(*prim).execute(strm, mem_ags);
|
||||
}
|
||||
|
||||
InferenceEngine::Precision Concat::getRuntimePrecision() const {
|
||||
@@ -560,6 +590,110 @@ void Concat::execNspcSpecCase() {
|
||||
});
|
||||
}
|
||||
|
||||
void Concat::execRef() {
|
||||
const size_t numSrc = getParentEdges().size();
|
||||
std::vector<const uint8_t*> srcPtrs;
|
||||
const Memory& dstMemory = getChildEdgeAt(0)->getMemory();
|
||||
const size_t elemSize = DnnlExtensionUtils::sizeOfDataType(dstMemory.GetDataType());
|
||||
const auto dstMemBlkDesc = dstMemory.getDescPtr()->as<BlockedMemoryDesc>();
|
||||
const auto& outputShape = dstMemBlkDesc->getBlockDims();
|
||||
uint8_t* dstPtr = reinterpret_cast<uint8_t*>(dstMemory.GetData());
|
||||
for (size_t i = 0; i < numSrc; i++) {
|
||||
const Memory& srcMem = getParentEdgesAtPort(i)[0]->getMemory();
|
||||
srcPtrs.push_back(reinterpret_cast<const uint8_t*>(srcMem.GetPtr()));
|
||||
}
|
||||
const auto& outputStrides = dstMemBlkDesc->getStrides();
|
||||
bool hasOuterLoop = false;
|
||||
for (size_t i = 0; i < reorderedAxis; i++) {
|
||||
if (outputShape[i] != 1) {
|
||||
hasOuterLoop = true;
|
||||
}
|
||||
}
|
||||
if (!hasOuterLoop) {
|
||||
int nthr = parallel_get_max_threads();
|
||||
if (nthr == 1) {
|
||||
for (size_t a = 0; a < srcPtrs.size(); ++a) {
|
||||
const auto inData = srcPtrs[a];
|
||||
auto outputData = &dstPtr[dstOffset[a]];
|
||||
std::memcpy(outputData, inData, nelemToCopy[a]);
|
||||
}
|
||||
} else {
|
||||
parallel_nt(nthr, [&](int ithr, int nthr) {
|
||||
for (size_t a = 0; a < srcPtrs.size(); ++a) {
|
||||
size_t start = 0, end = 0;
|
||||
splitter(nelemToCopy[a], nthr, ithr, start, end);
|
||||
const uint8_t* i = srcPtrs[a] + start;
|
||||
uint8_t* o = dstPtr + dstOffset[a] + start;
|
||||
std::memcpy(o, i, end - start);
|
||||
}
|
||||
});
|
||||
}
|
||||
} else {
|
||||
size_t physDims[5] = {1, 1, 1, 1, 1};
|
||||
for (size_t i = 0; i < reorderedAxis; i++) {
|
||||
physDims[i] = outputShape[i];
|
||||
}
|
||||
const auto L1Size = dnnl::utils::get_cache_size(1, true);
|
||||
UNUSED(L1Size); // for Windows
|
||||
parallel_for6d(physDims[0], physDims[1], physDims[2], physDims[3], physDims[4], numSrc,
|
||||
[&](size_t n0, size_t n1, size_t n2, size_t n3, size_t n4, size_t a) {
|
||||
// check if zero memory
|
||||
if (srcPtrs[a] == nullptr) return;
|
||||
|
||||
size_t inOff = inputStrides[a][0] * n0 + inputStrides[a][1] * n1 + inputStrides[a][2] * n2
|
||||
+ inputStrides[a][3] * n3 + inputStrides[a][4] * n4;
|
||||
size_t outOff = outputStrides[0] * n0 + outputStrides[1] * n1 + outputStrides[2] * n2
|
||||
+ outputStrides[3] * n3 + outputStrides[4] * n4;
|
||||
const uint8_t *i = srcPtrs[a] + inOff * elemSize;
|
||||
uint8_t *o = dstPtr + dstOffset[a] + outOff * elemSize;
|
||||
|
||||
#if defined(__GNUC__)
|
||||
// Heuristic:
|
||||
// memcpy works generally faster for data sizes not
|
||||
// exceeding L1 cache.
|
||||
if (nelemToCopy[a] > L1Size) {
|
||||
// The code below performs data copying: o[e] = i[e]
|
||||
// and uses a workaround to make GNU compilers optimize it
|
||||
uint8_t *ptro = o;
|
||||
const uint8_t *ptri = i;
|
||||
// head part: bytes before 4 byte-align's address
|
||||
const size_t headPart = sizeof(uint32_t)
|
||||
- reinterpret_cast<uint64_t>(ptro)
|
||||
% sizeof(uint32_t);
|
||||
|
||||
// main part: bytes in 4 byte-align
|
||||
const size_t mainPart
|
||||
= (nelemToCopy[a] - headPart) / sizeof(uint32_t);
|
||||
// tail part: bytes after 4 byte-align
|
||||
const size_t tailPart
|
||||
= (nelemToCopy[a]) - headPart
|
||||
- (mainPart * sizeof(uint32_t));
|
||||
// copy head part
|
||||
for (size_t e = 0; e < headPart; ++e) {
|
||||
*ptro = *ptri;
|
||||
++ptro;
|
||||
++ptri;
|
||||
}
|
||||
// copy main part
|
||||
std::memcpy(ptro, ptri, mainPart * sizeof(uint32_t));
|
||||
ptro += mainPart * sizeof(uint32_t);
|
||||
ptri += mainPart * sizeof(uint32_t);
|
||||
// copy tail part
|
||||
for (size_t e = 0; e < tailPart; ++e) {
|
||||
*ptro = *ptri;
|
||||
++ptro;
|
||||
++ptri;
|
||||
}
|
||||
} else {
|
||||
std::memcpy(o, i, nelemToCopy[a]);
|
||||
}
|
||||
#else
|
||||
std::memcpy(o, i, nelemToCopy[a]);
|
||||
#endif
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace node
|
||||
} // namespace intel_cpu
|
||||
} // namespace ov
|
||||
@@ -36,14 +36,18 @@ public:
|
||||
|
||||
private:
|
||||
size_t axis = 0;
|
||||
size_t reorderedAxis = 0;
|
||||
bool canBeInPlace = false;
|
||||
bool canOptimizeNspc = false;
|
||||
|
||||
void execRef();
|
||||
size_t inverseOrder(const InferenceEngine::SizeVector& order, size_t axis);
|
||||
void execNspcSpecCase();
|
||||
|
||||
std::vector<VectorDims> inputStrides;
|
||||
std::vector<size_t> nelemToCopy; // byte moved in each iter
|
||||
std::vector<size_t> dstOffset; // dst offset for each input
|
||||
InferenceEngine::Precision inputPrecision = InferenceEngine::Precision::FP32;
|
||||
InferenceEngine::Precision outputPrecision = InferenceEngine::Precision::FP32;
|
||||
bool canExecRef = false;
|
||||
};
|
||||
|
||||
} // namespace node
|
||||
|
||||
Reference in new issue
Block a user