[CPU] Refactor memory control and allocation

openvinotoolkit · Oct 26, 2024 · 28be812 · 28be812
1 parent 0c07136
commit 28be812
Show file tree

Hide file tree

Showing 19 changed files with 364 additions and 265 deletions.
diff --git a/src/plugins/intel_cpu/src/compiled_model.cpp b/src/plugins/intel_cpu/src/compiled_model.cpp
@@ -19,6 +19,7 @@
 #include "openvino/runtime/threading/cpu_streams_info.hpp"
 #include "openvino/runtime/threading/cpu_message.hpp"
 #include "utils/serialize.hpp"
+#include "memory_control.hpp"
 
 #include "cpu/x64/cpu_isa_traits.hpp"
 #include <cstring>
@@ -52,7 +53,8 @@ CompiledModel::CompiledModel(const std::shared_ptr<ov::Model>& model,
       m_cfg{cfg},
       m_name{model->get_name()},
       m_loaded_from_cache(loaded_from_cache),
-      m_sub_memory_manager(sub_memory_manager) {
+      m_sub_memory_manager(sub_memory_manager),
+      m_networkMemoryControl(std::make_shared<NetworkMemoryControl>()) {
     m_mutex = std::make_shared<std::mutex>();
     const auto& core = m_plugin->get_core();
     if (!core)
@@ -157,13 +159,14 @@ CompiledModel::GraphGuard::Lock CompiledModel::get_graph() const {
                     std::lock_guard<std::mutex> lock{*m_mutex.get()};
                     auto isQuantizedFlag = (m_cfg.lpTransformsMode == Config::On) &&
                                            ov::pass::low_precision::LowPrecision::isFunctionQuantized(m_model);
-
                     ctx = std::make_shared<GraphContext>(m_cfg,
                                                          m_socketWeights[socketId],
                                                          isQuantizedFlag,
+                                                         m_networkMemoryControl->createMemoryControlUnit(),
                                                          streamsExecutor,
                                                          m_sub_memory_manager);
                 }
+
                 const std::shared_ptr<const ov::Model> model = m_model;
                 graphLock._graph.CreateGraph(model, ctx);
             } catch (...) {
@@ -346,7 +349,7 @@ void CompiledModel::release_memory() {
     for (auto&& graph : m_graphs) {
         GraphGuard::Lock graph_lock{graph};
         auto ctx = graph_lock._graph.getGraphContext();
-        ctx->getNetworkMemoryControl()->releaseMemory();
+        m_networkMemoryControl->releaseMemory();
     }
 }
 

diff --git a/src/plugins/intel_cpu/src/compiled_model.h b/src/plugins/intel_cpu/src/compiled_model.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -19,6 +20,8 @@
 namespace ov {
 namespace intel_cpu {
 
+class NetworkMemoryControl;
+
 class CompiledModel : public ov::ICompiledModel {
 public:
     typedef std::shared_ptr<CompiledModel> Ptr;
@@ -51,6 +54,10 @@ class CompiledModel : public ov::ICompiledModel {
 
     void release_memory() override;
 
+    std::shared_ptr<NetworkMemoryControl> get_network_memory_control() const {
+        return m_networkMemoryControl;
+    }
+
 private:
     std::shared_ptr<ov::ISyncInferRequest> create_sync_infer_request() const override;
     friend class SyncInferRequest;
@@ -91,6 +98,7 @@ class CompiledModel : public ov::ICompiledModel {
 
     std::vector<std::shared_ptr<CompiledModel>> m_sub_compiled_models;
     std::shared_ptr<SubMemoryManager> m_sub_memory_manager = nullptr;
+    std::shared_ptr<NetworkMemoryControl> m_networkMemoryControl = nullptr;
     bool m_has_sub_compiled_models = false;
 };
 

diff --git a/src/plugins/intel_cpu/src/edge.cpp b/src/plugins/intel_cpu/src/edge.cpp
@@ -441,7 +441,7 @@ void Edge::validate() {
     getChild();
 
     if (status != Status::Allocated || !memoryPtr) {
-        OPENVINO_THROW("Error memory is not allocated!");
+        OPENVINO_THROW("Error memory is not allocated for edge: ", name());
     }
     status = Status::Validated;
 }

diff --git a/src/plugins/intel_cpu/src/edge.h b/src/plugins/intel_cpu/src/edge.h
@@ -82,6 +82,7 @@ class Edge {
     }
 
     std::string name() const;
+    const MemoryDesc& getDesc() const;
 
 private:
     std::weak_ptr<Node> parent;
@@ -99,7 +100,6 @@ class Edge {
     PortDescBaseCPtr getInputPortDesc() const;
     PortDescBaseCPtr getOutputPortDesc() const;
 
-    const MemoryDesc& getDesc() const;
     bool enforceReorder();
 
     void collectConsumers(std::vector<std::shared_ptr<Node>>& result) const;