CR fix

Author: ceciliapeng2011

src/plugins/intel_cpu/src/graph.cpp

@@ -27,7 +27,6 @@
 #include "nodes/convert.h"
 #include "nodes/subgraph.h"
 #include "nodes/fullyconnected.h"
-#include "output_mem_mgr.h"
 
 #include <ie_algorithm.hpp>
 #include <blob_factory.hpp>
@@ -903,7 +902,7 @@ void Graph::AllocateWithReuse() {
             if (isOutGrp) {
                 IE_ASSERT(isOutGrp==1);  // reuse_io_tensors false
                 grpMemMngr =
-                    std::make_shared<OutputMemoryMngr>(grpMemMngr);
+                    std::make_shared<ProxyMemoryMngr>(grpMemMngr);
                 DEBUG_LOG(grpMemMngr, " ", this);
 
                 // Store the output memory managers.
@@ -912,7 +911,7 @@ void Graph::AllocateWithReuse() {
                     const auto child = edge->getChild();
                     if (child->getType() == Type::Output) {
                         for (auto &output : outputNodesMap) {
-                            if (output.second == child) outputNodesMemMngrMap[output.first] = grpMemMngr;
+                            if (output.second == child) outputNodesMemMngrMap[output.first] = std::static_pointer_cast<ProxyMemoryMngr>(grpMemMngr);
                         }
                     }
                 }

src/plugins/intel_cpu/src/graph.h

@@ -19,6 +19,8 @@
 #include <memory>
 #include <atomic>
 
+#include "proxy_mem_mgr.h"
+
 namespace ov {
 namespace intel_cpu {
 
@@ -248,7 +250,7 @@ class Graph {
     std::map<std::string, NodePtr> inputNodesMap;
     std::map<std::string, NodePtr> outputNodesMap;
 
-    std::map<std::string, MemoryMngrPtr> outputNodesMemMngrMap;
+    std::map<std::string, ProxyMemoryMngrPtr> outputNodesMemMngrMap;
 
     // these node pointers (from graphNodes) are to avoid regular checking for
     // constantness of nodes in Infer methods and calls of

src/plugins/intel_cpu/src/infer_request.cpp

@@ -25,7 +25,7 @@
 #include "memory_desc/dnnl_blocked_memory_desc.h"
 #include <transformations/utils/utils.hpp>
 #include <ie_ngraph_utils.hpp>
-#include "output_mem_mgr.h"
+#include "proxy_mem_mgr.h"
 
 namespace ov {
 namespace intel_cpu {
@@ -269,26 +269,24 @@ void InferRequestBase::changeDefaultPtr() {
                 bool canBeInPlace = true;
                 // TODO: filter
 
-                auto inferrequest = std::dynamic_pointer_cast<InferRequest>(this->shared_from_this());
-                OPENVINO_ASSERT(inferrequest, "should be a InferRequest instance");
-
-                // share intel_cpu::Tensor to Graph by injecting to corresponding OutputMemoryMngr instance.
-                OutputMemoryMngrPtr outputMemMngr;
+                // share intel_cpu::Tensor to Graph by injecting to corresponding ProxyMemoryMngr instance.
+                ProxyMemoryMngrPtr outputMemMngr;
                 const auto &outMemMngrMap = graph->outputNodesMemMngrMap;
                 auto itr = outMemMngrMap.find(it.first);
                 if (itr != outMemMngrMap.end()) {
-                    outputMemMngr = std::dynamic_pointer_cast<OutputMemoryMngr>(itr->second);
+                    outputMemMngr = itr->second;
                     OPENVINO_ASSERT(outputMemMngr, "output memmanager should not be empty.");
                 } else {
                     OPENVINO_THROW("Cannot find output memmanager for output " + it.first + " !");
                 }
 
                 if (canBeInPlace) {
                     auto tt = std::get<0>(outputsTensor2BlobMap[it.first]);
-                    outputMemMngr->setTensor(tt);
+                    auto memptr = tt->get_memory();
+                    outputMemMngr->setManager(memptr->getMemoryMngr());
                     DEBUG_LOG("setTensor ", tt, " graph ", graph, " inferrequest ", this);
                 } else {
-                    outputMemMngr->setTensor(nullptr);
+                    outputMemMngr->setManager(nullptr);
                     DEBUG_LOG("setTensor nullptr", " graph ", graph, " inferrequest ", this);
                 }
 

src/plugins/intel_cpu/src/output_mem_mgr.cpp

@@ -0,0 +1,50 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "proxy_mem_mgr.h"
+#include "utils/debug_capabilities.h"
+
+using namespace ov::intel_cpu;
+
+void ProxyMemoryMngr::setManager(MemoryMngrPtr _pMngr) {
+    m_validated = (_pMngr != m_pMngr);
+    if (_pMngr) {
+        m_pMngr = _pMngr;
+    } else {
+        m_pMngr = m_pOrigMngr;
+    }
+}
+
+void* ProxyMemoryMngr::getRawPtr() const noexcept {
+    {
+        std::lock_guard<std::mutex> guard(m_lock);
+        if (m_validated) {
+            m_pMngr->resize(m_Size);
+            m_validated = false;
+        }
+    }
+    return m_pMngr->getRawPtr();
+}
+
+void ProxyMemoryMngr::setExtBuff(void* ptr, size_t size) {
+    return m_pMngr->setExtBuff(ptr, size);
+}
+
+bool ProxyMemoryMngr::resize(size_t size) {
+    m_Size = size;
+    m_validated = false;
+    return m_pMngr->resize(size);
+}
+
+bool ProxyMemoryMngr::hasExtBuffer() const noexcept {
+    return m_pMngr->hasExtBuffer();
+}
+
+void ProxyMemoryMngr::registerMemory(Memory* memPtr) {
+    m_pMngr->registerMemory(memPtr);
+}
+
+void ProxyMemoryMngr::unregisterMemory(Memory* memPtr) {
+    m_pMngr->unregisterMemory(memPtr);
+}

src/plugins/intel_cpu/src/proxy_mem_mgr.cpp

@@ -12,9 +12,9 @@
 namespace ov {
 namespace intel_cpu {
 
-class OutputMemoryMngr : public IMemoryMngrObserver {
+class ProxyMemoryMngr : public IMemoryMngrObserver {
 public:
-    explicit OutputMemoryMngr(MemoryMngrPtr pMngr) : m_pMngr(pMngr), m_pOrigMngr(pMngr) {
+    explicit ProxyMemoryMngr(MemoryMngrPtr pMngr) : m_pMngr(pMngr), m_pOrigMngr(pMngr) {
         OPENVINO_ASSERT(m_pOrigMngr, "Memory manager is uninitialized");
     }
 
@@ -26,15 +26,21 @@ class OutputMemoryMngr : public IMemoryMngrObserver {
     void registerMemory(Memory* memPtr) override;
     void unregisterMemory(Memory* memPtr) override;
 
-    void setTensor(std::shared_ptr<Tensor> tensor);
+    void setManager(MemoryMngrPtr _pMngr);
 
 private:
     // We keep the original MemMngr as may fallback to copy output.
     const MemoryMngrPtr m_pOrigMngr;
     MemoryMngrPtr m_pMngr;
+
+    // WA: resize stage might work because there is no shape change,
+    // but the underlying actual memory manager changes.
+    mutable std::mutex m_lock;
+    mutable size_t m_Size = 0ul;
+    mutable bool m_validated = false;
 };
-using OutputMemoryMngrPtr = std::shared_ptr<OutputMemoryMngr>;
-using OutputMemoryMngrCPtr = std::shared_ptr<const OutputMemoryMngr>;
+using ProxyMemoryMngrPtr = std::shared_ptr<ProxyMemoryMngr>;
+using ProxyMemoryMngrCPtr = std::shared_ptr<const ProxyMemoryMngr>;
 
 }   // namespace intel_cpu
 }   // namespace ov