Remove already released ids

yeonbok · Feb 3, 2024 · 4ba7ddc · 4ba7ddc
1 parent 4b135ae
commit 4ba7ddc
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/src/plugins/intel_gpu/src/graph/network.cpp b/src/plugins/intel_gpu/src/graph/network.cpp
@@ -1033,38 +1033,39 @@ void network::execute_impl(const std::vector<event::ptr>& events) {
                 }
             }
         }
+        execute_primitive(inst, events);
+        executed_prims++;
+        if (needs_flushing && executed_prims % flush_frequency == 0)
+            get_stream().flush();
         if (inst->get_node().is_type<kv_cache>()) {
             // Garbage collection of kv cache meories :
             // Once the corresponding kv cache's execution is done, the input mems are no
             // longer needed and can be released.
             GPU_DEBUG_TRACE_DETAIL << ": Check releasable kv cache memories" << std::endl;
-            std::vector<primitive_id> mem_deps_eol;
+            std::vector<std::pair<primitive_id, primitive_id>> mem_deps_eol;
             for (auto kms : _kv_cache_mem_deps) {
                 const auto kv_cache_id = kms.first;
                 auto queue_type = get_stream().get_queue_type();
                 if (queue_type == QueueTypes::in_order ||
                     (has_event(kv_cache_id) && get_primitive_event(kv_cache_id)->is_set())) {
-                    for (auto mem_deps : kms.second) {
-                        mem_deps_eol.push_back(mem_deps);
+                    for (auto mem_dep : kms.second) {
+                        mem_deps_eol.push_back(std::make_pair(kv_cache_id, mem_dep));
                     }
                 }
             }
             for (auto mem_dep : mem_deps_eol) {
-                auto mem_dep_inst = get_primitive(mem_dep);
+                auto mem_dep_inst = get_primitive(mem_dep.second);
                 GPU_DEBUG_TRACE_DETAIL << "Release output memory of " << mem_dep_inst->id() << ": "
                                        << ((mem_dep_inst->output_memory_ptr())
                                                ? mem_dep_inst->output_memory_ptr()->buffer_ptr()
                                                : " 0xffffffff")
                                        << std::endl;
 
                 mem_dep_inst->release_output_memory();
+                auto& kv_map = _kv_cache_mem_deps[mem_dep.first];
+                kv_map.erase(std::find(kv_map.begin(), kv_map.end(), mem_dep.second));
             }
         }
-        execute_primitive(inst, events);
-        executed_prims++;
-        if (needs_flushing && executed_prims % flush_frequency == 0)
-            get_stream().flush();
-
         // Dump output buffers of 'inst'
         GPU_DEBUG_IF(debug_config->dump_layers_path.length() > 0) {
             get_stream().finish();

diff --git a/src/plugins/intel_gpu/src/graph/primitive_inst.cpp b/src/plugins/intel_gpu/src/graph/primitive_inst.cpp
@@ -224,7 +224,7 @@ void primitive_inst::check_memory_to_set(const memory& mem, const layout& layout
 }
 
 void primitive_inst::release_output_memory() {
-    for(auto o : _outputs) {
+    for(auto& o : _outputs) {
         o = nullptr;
     }
     _mem_allocated = false;