[GPU] Extend shape_of subgraphs markup logic to include PagedAttention input

Author: sshlyapn

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

@@ -10,14 +10,34 @@
 #include "select_inst.h"
 #include "strided_slice_inst.h"
 #include "gather_inst.h"
+#include "input_layout_inst.h"
+#include "paged_attention_inst.h"
 #include "pass_manager.h"
 
 #include "intel_gpu/graph/program.hpp"
 
 using namespace cldnn;
 
-void mark_shape_of_subgraphs::look_for_shape_of_subgraph(program_node& node) {
+static bool is_shape_of_subgraph_root(program_node& node) {
     if (node.is_type<shape_of>()) {
+        return true;
+    }
+
+    if (node.is_type<input_layout>()) {
+        const auto& users = node.get_users();
+        for (const auto& user : users) {
+            const auto max_context_len_input_id = 12;
+            if (user->is_type<paged_attention>() && user->get_dependency_index(node) == max_context_len_input_id) {
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
+void mark_shape_of_subgraphs::look_for_shape_of_subgraph(program_node& node) {
+    if (is_shape_of_subgraph_root(node)) {
         mark_node(node);
         return;
     }
@@ -102,7 +122,7 @@ void mark_shape_of_subgraphs::mark_node(program_node& node) {
 
     // If current node has shape_of type add it to dependant shape_of nodes for
     // correct dependency propagation for users
-    if (node.is_type<shape_of>())
+    if (is_shape_of_subgraph_root(node))
         node.add_dependant_shape_of_node(&node);
 
     // Add parent shape_of nodes from other dependencies if there are any

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

@@ -658,7 +658,7 @@ void program_node::select_preferred_formats(impl_types impl_type) {
 }
 
 void program_node::add_dependant_shape_of_node(const program_node* node) {
-    OPENVINO_ASSERT(node->is_type<shape_of>(), "[GPU] Expected node type is shape_of");
+    OPENVINO_ASSERT(node->is_type<shape_of>() || node->is_type<input_layout>(), "[GPU] Expected node type is shape_of");
     dependant_shape_of_nodes.insert(node);
 }