[CPU] Optimize concat node for dynamic shapes (#13019)

This commit is contained in:
Zhang Yi authored and GitHub committed 2022-11-11 13:46:32 +04:00
1 parent b44eb83551
commit b8ada02cba
3 files changed
+232 -34

No files matched your search

+60
View File
@@ -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
+166 -32
View File
@@ -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
+6 -2
View File
@@ -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