inference-engine/src/mkldnn_plugin/nodes/mkldnn_memory_node.cpp

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #include <string>
   6 #include <mkldnn_types.h>
   7 #include <mkldnn_extension_utils.h>
   8 #include "mkldnn_memory_node.hpp"
   9
  10 using namespace mkldnn;
  11 using namespace MKLDNNPlugin;
  12 using namespace InferenceEngine;
  13
  14 MKLDNNMemoryOutputNode::MKLDNNMemoryOutputNode(const InferenceEngine::CNNLayerPtr& layer, const mkldnn::engine& eng)
  15         : MKLDNNNode(layer, eng) , MKLDNNMemoryNode(layer) {
  16     if (created()) {
  17         MKLDNNMemoryNodeVirtualEdge::registerOutput(this);
  18     }
  19 }
  20
  21 MKLDNNMemoryOutputNode::~MKLDNNMemoryOutputNode() {
  22     MKLDNNMemoryNodeVirtualEdge::remove(this);
  23 }
  24
  25 void MKLDNNMemoryOutputNode::getSupportedDescriptors() {}
  26
  27 void MKLDNNMemoryOutputNode::initSupportedPrimitiveDescriptors() {
  28     if (!supportedPrimitiveDescriptors.empty())
  29         return;
  30
  31     InferenceEngine::Precision precision = getCnnLayer()->insData[0].lock()->getPrecision();
  32     if (precision != InferenceEngine::Precision::FP32)
  33         precision = InferenceEngine::Precision::FP32;
  34     auto inputDataType = MKLDNNExtensionUtils::IEPrecisionToDataType(precision);
  35     InferenceEngine::LayerConfig config;
  36     config.dynBatchSupport = true;
  37     config.inConfs.resize(1);
  38     config.inConfs[0].inPlace = -1;
  39     config.inConfs[0].constant = false;
  40     config.inConfs[0].desc = MKLDNNMemoryDesc(getParentEdgeAt(0)->getDims(), inputDataType, memory::format::any);
  41     supportedPrimitiveDescriptors.push_back({config, impl_desc_type::unknown});
  42 }
  43
  44 const MKLDNNEdgePtr MKLDNNMemoryOutputNode::getChildEdgeAt(size_t idx) const {
  45     if (inputNode != nullptr) {
  46         return inputNode->getChildEdgeAt(idx);
  47     }
  48     return MKLDNNNode::getChildEdgeAt(idx);
  49 }
  50
  51 void MKLDNNMemoryOutputNode::execute(mkldnn::stream strm)  {
  52     auto& srcMemory = getParentEdgeAt(0)->getMemory();
  53
  54     const float *src_ptr = reinterpret_cast<const float*>(srcMemory.GetData()) +
  55             srcMemory.GetDescriptor().data.layout_desc.blocking.offset_padding;
  56     float *dst_ptr = reinterpret_cast<float*>(getChildEdgeAt(0)->getMemory().GetData()) +
  57             getChildEdgeAt(0)->getMemory().GetDescriptor().data.layout_desc.blocking.offset_padding;
  58
  59     // TODO: this can be eliminated by completely removing MKLDNN memory output NODE, to fuse it with output of prev layer
  60     memcpy(dst_ptr, src_ptr, srcMemory.GetSize());
  61 }
  62
  63 std::string MKLDNNMemoryInputNode::nameFromCombinedName(std::string name) {
  64     auto idSplitter = name.find("/id=");
  65     return name.substr(0, idSplitter);
  66 }
  67
  68 std::string MKLDNNMemoryInputNode::idFromCombinedName(std::string name) {
  69     auto idSplitter = name.find("/id=");
  70     return name.substr(idSplitter == std::string::npos ? 0 : idSplitter + 4);
  71 }
  72
  73 MKLDNNMemoryInputNode::MKLDNNMemoryInputNode(const InferenceEngine::CNNLayerPtr& layer, const mkldnn::engine& eng)
  74         : MKLDNNInputNode(layer, eng), MKLDNNMemoryNode(layer) {
  75     if (created()) {
  76         MKLDNNMemoryNodeVirtualEdge::registerInput(this);
  77     }
  78 }
  79
  80 MKLDNNMemoryInputNode::~MKLDNNMemoryInputNode() {
  81     MKLDNNMemoryNodeVirtualEdge::remove(this);
  82 }
  83
  84 void MKLDNNMemoryNodeVirtualEdge::registerInput(MKLDNNMemoryInputNode * node) {
  85     // in case of output already registered
  86     auto sibling = MKLDNNMemoryNodeVirtualEdge::getByName(node->getId());
  87     if (sibling != nullptr) {
  88         auto outputNode = dynamic_cast<MKLDNNMemoryOutputNode*>(sibling);
  89         IE_ASSERT(outputNode != nullptr);
  90         outputNode->setInputNode(node);
  91     } else {
  92         getExisted()[node->getId()] = node;
  93     }
  94     // std::cout <<"[register] " << node << ", size="<< getExisted().size() <<"\n" << std::flush;
  95 }
  96
  97 void MKLDNNMemoryNodeVirtualEdge::registerOutput(MKLDNNMemoryOutputNode * node) {
  98     // in case of output layer
  99     auto sibling = MKLDNNMemoryNodeVirtualEdge::getByName(node->getId());
 100     if (sibling != nullptr) {
 101         auto inputNode = dynamic_cast<MKLDNNMemoryInputNode*>(sibling);
 102         IE_ASSERT(inputNode != nullptr);
 103         node->setInputNode(inputNode);
 104     } else {
 105         getExisted()[node->getId()] = node;
 106     }
 107     // std::cout <<"[register] " << node << ", size="<< getExisted().size() <<"\n" << std::flush;
 108 }