refactor stateful model impl by removing storage from MemoryInput nod…

…e and decrease memory copy from 4 times to 2 times.
slyalin · Oct 7, 2023 · 86ddb40 · 86ddb40
1 parent a13b503
commit 86ddb40
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 21 deletions.
diff --git a/src/plugins/intel_cpu/src/infer_request.cpp b/src/plugins/intel_cpu/src/infer_request.cpp
@@ -107,9 +107,13 @@ void InferRequestBase::PushStates() {
             auto cur_id = cur_node->getId();
             for (const auto& state : memoryStates) {
                 if (state->GetName() == cur_id) {
+                    auto storage = cur_node->getStore();
                     auto state_blob = state->GetState();
+                    if (storage->getData() == state_blob->cbuffer().as<const void *>())
+                        continue;  // there is no inferrequest switch
+
                     auto state_desc = MemoryDescUtils::convertToDnnlBlockedMemoryDesc(state_blob->getTensorDesc());
-                    Memory state_mem(eng, state_desc, state_blob->cbuffer(), false);
+                    auto state_mem = std::make_shared<Memory>(eng, state_desc, state_blob->cbuffer(), false);
                     cur_node->storeState(state_mem);
                 }
             }
@@ -129,11 +133,7 @@ void InferRequestBase::PullStates() {
             for (const auto& state : memoryStates) {
                 if (state->GetName() == cur_id) {
                     auto storage = cur_node->getStore();
-
-                    //redefine state
-                    auto blob = make_blob_with_precision(MemoryDescUtils::convertToTensorDesc(storage->getDesc()));
-                    blob->allocate();
-                    cpu_memcpy(blob->buffer(), storage->getData(), storage->getSize());
+                    auto blob = make_blob_with_precision(MemoryDescUtils::convertToTensorDesc(storage->getDesc()), storage->getData());
                     state->SetState(blob);
                 }
             }

diff --git a/src/plugins/intel_cpu/src/memory_state.h b/src/plugins/intel_cpu/src/memory_state.h
@@ -20,6 +20,7 @@ class VariableState : public InferenceEngine::IVariableStateInternal {
     VariableState(std::string name, MemoryPtr storage)
         : InferenceEngine::IVariableStateInternal{name} {
         tensor_desc = MemoryDescUtils::convertToTensorDesc(storage->getDesc());
+        Reset();
     }
 
     void Reset() override;

diff --git a/src/plugins/intel_cpu/src/nodes/memory.cpp b/src/plugins/intel_cpu/src/nodes/memory.cpp
@@ -19,6 +19,9 @@ namespace ov {
 namespace intel_cpu {
 namespace node {
 
+inline
+static void simple_copy(const IMemory& dst, const IMemory& src);
+
 std::mutex MemoryNodeVirtualEdge::holderMutex;
 
 MemoryNode::MemoryNode(const std::shared_ptr<ngraph::Node>& op) {
@@ -79,7 +82,15 @@ void MemoryOutput::execute(dnnl::stream strm)  {
     auto inputMemoryNode = dynamic_cast<MemoryInput*>(inputNode);
     IE_ASSERT(inputMemoryNode != nullptr);
 
-    inputMemoryNode->storeState(srcMemory);
+    //redefine storage memory before copy to it
+    auto storage = inputMemoryNode->getStore();
+
+    auto desc = storage->getDescPtr();
+    const auto new_shape = srcMemory.getStaticDims();
+    const auto newDesc = desc->cloneWithNewDims(new_shape, true);
+    storage->redefineDesc(newDesc);
+
+    simple_copy(*storage, srcMemory);
 }
 
 void MemoryOutput::executeDynamicImpl(dnnl::stream strm) {
@@ -162,18 +173,8 @@ MemoryPtr MemoryInput::getStore() {
     return dataStore;
 }
 
-void MemoryInput::storeState(const IMemory &new_state) {
-    //redefine storage memory before copy to it
-    const auto &_desc = new_state.getDesc();
-    auto _shape = _desc.getShape();
-    VectorDims dims = _shape.getStaticDims();
-    auto desc = _desc.cloneWithNewDims(dims);
-    dataStore = std::make_shared<Memory>(getEngine(), desc);
-
-    // TODO: Should be next one call:
-    //           dataStore.load(new_state, false);
-    //       But because of performance reason we use simple manual copy
-    simple_copy(*dataStore, new_state);
+void MemoryInput::storeState(const MemoryPtr new_state) {
+    dataStore = new_state;
 }
 
 void MemoryInput::execute(dnnl::stream strm) {

diff --git a/src/plugins/intel_cpu/src/nodes/memory.hpp b/src/plugins/intel_cpu/src/nodes/memory.hpp
@@ -105,10 +105,10 @@ class MemoryInput : public Input, public MemoryNode {
     void createPrimitive() override;
 
     void setInputNode(Node* node) override {}
-    void storeState(const IMemory& mem);
+    void storeState(const MemoryPtr mem);
     MemoryPtr getStore();
  private:
-    MemoryPtr dataStore;
+    MemoryPtr dataStore = nullptr;
     MemoryNodeVirtualEdge::Holder* holder = nullptr;
 };