Apply additional comments

src/plugins/intel_gpu/include/intel_gpu/graph/network.hpp

@@ -181,6 +181,7 @@ struct network {
     bool is_primary_stream() const { return _is_primary_stream; }
     bool is_dynamic() const { return _is_dynamic; }
     size_t get_weights_cache_capacity() const { return _weights_cache_capacity; }
+    bool contains_state(const std::string& variable_id);
 
     memory_pool& get_memory_pool() const {
         return *_memory_pool;
@@ -225,6 +226,8 @@ struct network {
 
     ov::intel_gpu::VariablesMap _variables_states;
     ov::intel_gpu::VariablesInfoMap _variables_state_info;
+    std::vector<std::shared_ptr<primitive_inst>> _read_values;
+    std::unordered_map<primitive_id, std::vector<std::shared_ptr<primitive_inst>>> _state_initializers;
 
     program::primitives_info _prims_info;
     size_t _weights_cache_capacity = 1;

src/plugins/intel_gpu/include/intel_gpu/graph/program.hpp

@@ -173,6 +173,11 @@ struct program {
     program_node const& get_node(primitive_id const& id) const;
     std::shared_ptr<program_node> get_node_ptr(const primitive_id& prim) { return nodes_map.at(prim); }
     std::shared_ptr<program_node> get_node_ptr(const primitive_id& prim) const { return nodes_map.at(prim); }
+    void set_state_initializers(const std::string& variable_id, const primitive_id& id);
+    bool has_state_initializers(const std::string& variable_id, const primitive_id& id);
+    bool contains_state(const std::string& variable_id);
+    const std::vector<primitive_id>& get_initializers(const std::string& variable_id) { return state_initializers.at(variable_id); }
+    const std::map<std::string, std::vector<primitive_id>>& get_state_initializers() const { return state_initializers; }
 
     // returns already existing program_node for given primitive 'prim' (lookup in 'nodes_map')
     // if it was previously created, otherwise creates and then returns program_node
@@ -322,6 +327,8 @@ struct program {
     primitives_info prim_info;
     graph_optimizer_info optimizer_passes_info;
 
+    std::map<std::string, std::vector<primitive_id>> state_initializers;
+
     primitives_info get_current_stage_info() const;
     /*
     ** High-level functions, in order of usage

src/plugins/intel_gpu/src/graph/graph_optimizer/mark_state_init_subgraphs.cpp

@@ -10,27 +10,27 @@
 
 using namespace cldnn;
 
-void mark_state_init_subgraphs::mark_node(program_node* node) {
-    if (node->is_in_state_init_subgraph())
+void mark_state_init_subgraphs::mark_init_subgraph(program& p, read_value_node& node) {
+    const auto& variable_id = node.get_primitive()->variable_id;
+    if (p.contains_state(variable_id))
         return;
-    if (!node->is_type<read_value>())
-        return;
-
-    const auto& variable_id = node->as<read_value>().get_primitive()->variable_id;
-    node->set_state_variable_id_of_init_subgraph(variable_id);
 
     std::queue<program_node*> q;
-    q.push(node);
+    q.push(&node);
+
+    auto can_be_marked = [&](const program_node* dep_node) {
+        if (p.has_state_initializers(variable_id, dep_node->id()))
+            return false;
 
-    auto can_be_marked = [&](const program_node* dep_node, const program_node* cur_node) {
         for (auto& u : dep_node->get_users()) {
-            if (u == cur_node)
+            if (u == &node)
                 continue;
-            if (u->get_state_variable_id_of_init_subgraph().compare(variable_id) != 0) {
+            if (p.has_state_initializers(variable_id, u->id()))
+                continue;
+            else
                 return false;
-            }
         }
-        GPU_DEBUG_TRACE_DETAIL << "marked " << dep_node->id() << " as node in a init_subgraph for " << node->id() << std::endl;
+        GPU_DEBUG_TRACE_DETAIL << "marked " << dep_node->id() << " as node in a init_subgraph for " << node.id() << std::endl;
         return true;
     };
 
@@ -40,10 +40,8 @@ void mark_state_init_subgraphs::mark_node(program_node* node) {
             auto& cur_node = q.front();
             q.pop();
             for (auto& dep : cur_node->get_dependencies()) {
-                if (can_be_marked(dep.first, cur_node)) {
-                    dep.first->set_state_variable_id_of_init_subgraph(variable_id);
-                    if (!dep.first->is_constant())
-                        node->add_state_initializer(dep.first->id());
+                if (can_be_marked(dep.first)) {
+                    p.set_state_initializers(variable_id, dep.first->id());
                     q.push(dep.first);
                 }
             }
@@ -53,10 +51,10 @@ void mark_state_init_subgraphs::mark_node(program_node* node) {
 
 void mark_state_init_subgraphs::run(program& p) {
     auto rit = p.get_processing_order().rbegin();
-    if (p.is_new_shape_infer()) {
-        for (; rit != p.get_processing_order().rend(); rit++) {
-            auto node = *rit;
-            mark_node(node);
+    for (; rit != p.get_processing_order().rend(); rit++) {
+        auto& node = *rit;
+        if (node->is_type<read_value>()) {
+            mark_init_subgraph(p, node->as<read_value>());
         }
     }
 }

src/plugins/intel_gpu/src/graph/impls/cpu/read_value.cpp

@@ -66,7 +66,7 @@ struct read_value_impl : public typed_primitive_impl<read_value> {
             if (!instance.get_user_insts().empty()) {
                 auto user_inst = instance.get_user_insts().front();
                 if (!(user_inst->get_node().is_type<assign>() || user_inst->get_node().is_type<kv_cache>()) &&
-                    !instance.get_state_initializers().empty()) {
+                    instance.get_network().contains_state(variable_id)) {
                     variable.set();
                 }
             }

src/plugins/intel_gpu/src/graph/include/pass_manager.h

@@ -9,6 +9,7 @@
 #include "quantize_inst.h"
 #include "eltwise_inst.h"
 #include "convolution_inst.h"
+#include "read_value_inst.h"
 #include <string>
 #include <vector>
 #include <memory>
@@ -96,7 +97,7 @@ class mark_state_init_subgraphs : public base_pass {
 
 private:
     void run(program& p) override;
-    void mark_node(program_node* node);
+    void mark_init_subgraph(program& p, read_value_node& node);
 };
 
 class mark_shape_of_subgraphs : public base_pass {

src/plugins/intel_gpu/src/graph/include/primitive_inst.h

@@ -201,7 +201,6 @@ class primitive_inst {
     uint32_t get_network_id() const;
     virtual event::ptr set_output_memory(memory::ptr mem, bool check = true, size_t idx = 0);
     void check_memory_to_set(const memory& mem, const layout& layout) const;
-    const std::vector<primitive_inst*>& get_state_initializers() const { return state_initializers; }
     const std::list<const cldnn::program_node *>& get_users() const { return _node->get_users(); }
     const std::vector<primitive_inst*>& get_user_insts() const { return _users; }
     void init_users() {
@@ -270,7 +269,6 @@ class primitive_inst {
     void do_runtime_in_place_crop();
     void do_runtime_skip_scatter_update();
     void configure_shape_of_dependencies();
-    void configure_state_initializers();
 
     memory::ptr fused_memory(size_t dep_id) const {
         return dep_memory_ptr(get_fused_mem_offset() + dep_id);
@@ -357,9 +355,6 @@ class primitive_inst {
     // List of depandant shape_of primitives for shape_of subgraphs
     std::vector<primitive_inst*> dependant_shape_of_insts;
 
-    // List of dependant primitives for state initializer subgraphs
-    std::vector<primitive_inst*> state_initializers;
-
     std::vector<primitive_inst*> _users;
     // this is a set of dependencies in terms of execution
     // execution of all primitives from this set should be enough to guarantee that all memory deps (see _deps)

src/plugins/intel_gpu/src/graph/include/program_node.h

@@ -301,10 +301,6 @@ struct program_node {
     void set_in_shape_of_subgraph(bool val = true) { in_shape_of_subgraph = val; }
     bool is_in_shape_of_subgraph() const { return in_shape_of_subgraph; }
 
-    void set_state_variable_id_of_init_subgraph(std::string id) { state_id_of_init_subgraph = id; }
-    bool is_in_state_init_subgraph() const { return !state_id_of_init_subgraph.empty(); }
-    std::string get_state_variable_id_of_init_subgraph() const { return state_id_of_init_subgraph; }
-
     // check/set if the node can be optimized out (removed from the network)
     bool can_be_optimized() const { return optimized; }
     void can_be_optimized(bool opt) { optimized = opt; }
@@ -372,12 +368,6 @@ struct program_node {
         return dependant_shape_of_nodes;
     }
 
-    void add_state_initializer(const primitive_id& id);
-
-    const std::vector<primitive_id>& get_state_initializers() const {
-        return state_initializers;
-    }
-
     void set_reused_memory_color(uint32_t color) const {
         has_reused_memory = true;
         reused_memory_color = color;
@@ -517,9 +507,6 @@ struct program_node {
 
     std::set<const program_node*> dependant_shape_of_nodes;
 
-    std::string state_id_of_init_subgraph;
-    std::vector<primitive_id> state_initializers;
-
     bool output = false;
     uint8_t user_mark = 0;
     bool optimized = false;

src/plugins/intel_gpu/src/graph/network.cpp

@@ -624,7 +624,6 @@ void network::build_insts_deps() {
         inst.second->build_deps();
         inst.second->init_users();
         inst.second->configure_shape_of_dependencies();
-        inst.second->configure_state_initializers();
     }
 }
 
@@ -662,6 +661,15 @@ void network::build_exec_order() {
         }
     }
 }
+
+bool network::contains_state(const std::string& variable_id) {
+    auto it = _state_initializers.find(variable_id);
+    if (it != _state_initializers.end())
+        return true;
+    else
+        return false;
+}
+
 void network::add_to_exec_order(const primitive_id& id) {
     auto inst = get_primitive(id);
     _exec_order.push_back(inst);
@@ -699,12 +707,13 @@ std::map<primitive_id, network_output> network::execute(const std::vector<event:
         }
     }
 
-    for (auto& inst : _exec_order) {
-        if (inst->get_node().is_type<read_value>() && !inst->get_state_initializers().empty()) {
-            auto prim = inst->get_node().as<read_value>().get_primitive();
+    for (auto& inst : _read_values) {
+        const auto& prim = inst->get_node().as<read_value>().get_primitive();
+        auto it = _state_initializers.find(prim->variable_id);
+        if (it != _state_initializers.end()) {
             const auto& variable = get_variable(prim->variable_id);
             if (variable.is_set()) {
-                for (auto& init_inst : inst->get_state_initializers()) {
+                for (auto& init_inst : it->second) {
                     init_inst->set_flag(ExecutionFlags::SKIP);
                 }
             }
@@ -926,6 +935,15 @@ void network::allocate_primitive_instance(program_node const& node) {
         if (node.is_type<data>())
             _data_outputs.push_back(inst);
     }
+    if (node.is_type<read_value>()) {
+        _read_values.push_back(inst);
+        const auto& variable_id = node.as<read_value>().get_primitive()->variable_id;
+        if (_program->contains_state(variable_id)) {
+            for (const auto& id : _program->get_initializers(variable_id)) {
+                _state_initializers[variable_id].push_back(get_primitive(id));
+            }
+        }
+    }
     if (auto state_prim = std::dynamic_pointer_cast<memory_state::variable>(inst)) {
         auto prim = inst->get_node().get_primitive();
         set_variables_state_info(state_prim->variable_id(), node.get_output_layout(0), state_prim->get_user_specified_type(), prim.get());

src/plugins/intel_gpu/src/graph/primitive_inst.cpp

@@ -2015,17 +2015,6 @@ void primitive_inst::configure_shape_of_dependencies() {
     }
 }
 
-void primitive_inst::configure_state_initializers() {
-    if (!state_initializers.empty())
-        return;
-
-    OPENVINO_ASSERT(_node != nullptr, "_node should not be nullptr for configure_state_initializers.");
-
-    if (state_initializers.empty() && !_node->get_state_initializers().empty()) {
-        state_initializers = _network.get_primitives(_node->get_state_initializers());
-    }
- }
-
 primitive_inst::primitive_inst(network& network)
     : _network(network)
     , _node(nullptr)

src/plugins/intel_gpu/src/graph/program.cpp

@@ -832,6 +832,27 @@ void program::reverse_connection(program_node& dep_node, program_node& user_node
     }
 }
 
+void program::set_state_initializers(const std::string& variable_id, const primitive_id& id) {
+    state_initializers[variable_id].push_back(id);
+}
+
+bool program::has_state_initializers(const std::string& variable_id, const primitive_id& id) {
+    auto it = state_initializers.find(variable_id);
+    if (it != state_initializers.end()) {
+        const auto& initializers = it->second;
+        return std::find(initializers.begin(), initializers.end(), id) != initializers.end();
+    }
+    return false;
+}
+
+bool program::contains_state(const std::string& variable_id) {
+    auto it = state_initializers.find(variable_id);
+    if (it != state_initializers.end())
+        return true;
+    else
+        return false;
+}
+
 program_node& program::get_or_create(std::shared_ptr<primitive> prim) {
     auto itr = nodes_map.lower_bound(prim->id);
     if (itr != nodes_map.end() && itr->first == prim->id)

src/plugins/intel_gpu/src/graph/program_node.cpp

@@ -323,9 +323,6 @@ std::unique_ptr<json_composite> program_node::desc_to_json() const {
     }
     node_info->add("dependant_shape_of_nodes_ids", dependant_shape_of_nodes_ids);
     node_info->add("in_shape_of_subgraph", in_shape_of_subgraph);
-
-    node_info->add("state_id_of_init_subgraph", state_id_of_init_subgraph);
-    node_info->add("state_initializers", state_initializers);
     return node_info;
 }
 
@@ -665,10 +662,6 @@ void program_node::add_dependant_shape_of_node(const program_node* node) {
     dependant_shape_of_nodes.insert(node);
 }
 
-void program_node::add_state_initializer(const primitive_id& id) {
-    state_initializers.push_back(id);
-}
-
 void program_node::save(cldnn::BinaryOutputBuffer& ob) const {
     ob << unique_id;
     ob << valid_output_layouts;
@@ -705,9 +698,6 @@ void program_node::save(cldnn::BinaryOutputBuffer& ob) const {
     ob << in_shape_of_subgraph;
     ob << runtime_skippable;
 
-    ob << state_id_of_init_subgraph;
-    ob << state_initializers;
-
     ob << output;
     ob << user_mark;
     ob << optimized;
@@ -883,9 +873,6 @@ void program_node::load(cldnn::BinaryInputBuffer& ib) {
     ib >> in_shape_of_subgraph;
     ib >> runtime_skippable;
 
-    ib >> state_id_of_init_subgraph;
-    ib >> state_initializers;
-
     ib >> output;
     ib >> user_mark;
     ib >> optimized;