[CPU]whisper readvalue optimize

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -140,6 +140,7 @@ static const TypeToNameMap& get_type_to_name_tbl() {
         {"Loop", Type::TensorIterator},
         {"ReadValue", Type::MemoryInput},  // for construction from name ctor, arbitrary name is used
         {"Assign", Type::MemoryOutput},    // for construction from layer ctor
+        {"ReadValueWithSubgraph", Type::MemoryInput},
         {"Convert", Type::Convert},
         {"NV12toRGB", Type::ColorConvert},
         {"NV12toBGR", Type::ColorConvert},

src/plugins/intel_cpu/src/extension.cpp

@@ -22,6 +22,7 @@
 #include "transformations/cpu_opset/common/op/power_static.hpp"
 #include "transformations/cpu_opset/common/op/sdpa.hpp"
 #include "transformations/cpu_opset/common/op/swish_cpu.hpp"
+#include "transformations/cpu_opset/common/op/read_value_with_subgraph.hpp"
 #include "transformations/cpu_opset/x64/op/interaction.hpp"
 #include "transformations/cpu_opset/x64/op/mha.hpp"
 #include "transformations/cpu_opset/x64/op/llm_mlp.hpp"
@@ -77,6 +78,7 @@ class TypeRelaxedExtension : public ov::OpExtension<ov::op::TypeRelaxed<Op>> {
     OP_EXTENSION(ov::intel_cpu::SwishNode)                                  \
     OP_EXTENSION(ov::intel_cpu::SDPAWithTransposeReshape)                   \
     OP_EXTENSION(ov::intel_cpu::NgramNode)                                  \
+    OP_EXTENSION(ov::intel_cpu::ReadValueWithSubgraph)                      \
     OP_EXTENSION(ov::op::internal::GatherCompressed)                        \
     OP_EXTENSION(ov::op::internal::NonMaxSuppressionIEInternal)             \
     OP_EXTENSION(ov::op::internal::MulticlassNmsIEInternal)                 \

src/plugins/intel_cpu/src/graph_optimizer.cpp

@@ -3120,12 +3120,20 @@ void GraphOptimizer::MatchSdpaKvCache(Graph &graph) {
         auto memInputNode = std::dynamic_pointer_cast<node::MemoryInputBase>(node);
         OPENVINO_ASSERT(memInputNode, "MemoryInput node ", node->getName(), " has unexpected dynamic type");
 
-        ov::optional<Shape> input_shape;
-        ov::optional<ov::element::Type> input_prc;
+        ov::optional<std::vector<Shape>> input_shape;
+        ov::optional<std::vector<ov::element::Type>> input_prc;
 
+        std::vector<Shape> input_shape_vec;
+        std::vector<ov::element::Type> input_prc_vec;
         if (!node->getParentEdges().empty()) {
-            input_shape = ov::optional<Shape>(node->getInputShapeAtPort(0));
-            input_prc = ov::optional<ov::element::Type>(node->getOriginalInputPrecisionAtPort(0));
+            for (size_t i = 0; i < node->getParentEdges().size(); i++) {
+                input_shape_vec.push_back(node->getInputShapeAtPort(i));
+                input_prc_vec.push_back(node->getOriginalInputPrecisionAtPort(i));
+            }
+        }
+        if (input_shape_vec.size() > 0) {
+            input_shape = ov::optional<std::vector<Shape>>(input_shape_vec);
+            input_prc = ov::optional<std::vector<ov::element::Type>>(input_prc_vec);
         }
 
         //search for SDPA
@@ -3251,12 +3259,18 @@ void GraphOptimizer::DropRedundantMemoryOutput(Graph &graph) {
         auto memInputNode = std::dynamic_pointer_cast<node::MemoryInputBase>(node);
         OPENVINO_ASSERT(memInputNode, "MemoryInput node ", node->getName(), " has unexpected dynamic type");
 
-        ov::optional<Shape> inputShape;
-        ov::optional<ov::element::Type> inputPrc;
-
+        ov::optional<std::vector<Shape>> inputShapes;
+        ov::optional<std::vector<ov::element::Type>> inputPrcs;
         if (!node->getParentEdges().empty()) {
-            inputShape = ov::optional<Shape>(node->getInputShapeAtPort(0));
-            inputPrc = ov::optional<ov::element::Type>(node->getOriginalInputPrecisionAtPort(0));
+            inputShapes = ov::optional<std::vector<Shape>>(std::vector<Shape>{});
+            inputPrcs = ov::optional<std::vector<ov::element::Type>>(std::vector<ov::element::Type>{});
+
+            auto& input_shape_vec = *inputShapes;
+            auto& input_prc_vec = *inputPrcs;
+            for (size_t i = 0; i < node->getParentEdges().size(); i++) {
+                input_shape_vec.push_back(node->getInputShapeAtPort(i));
+                input_prc_vec.push_back(node->getOriginalInputPrecisionAtPort(i));
+            }
         }
 
         //search for the MemoryOutputNode
@@ -3273,24 +3287,35 @@ void GraphOptimizer::DropRedundantMemoryOutput(Graph &graph) {
         graph.RemoveEdge(memoryOutputNode->getParentEdgeAt(0));
         // there are no output edges from MemoryOutput nodes
 
+        CPU_GRAPH_OPTIMIZER_SCOPE(DropRedundantMemoryOutput_SubGraph);
+        auto memInpNd = std::dynamic_pointer_cast<node::MemoryInput>(node);
+        auto subGraph = memInpNd ? memInpNd->getSubGraph() : nullptr;
+
         // now replace the existing MemoryInput with a special type that works without the corresponding MemoryOutput
         auto memInputSingle = std::make_shared<MemoryInputSingle>(memInputNode->getId(),
                                                                   memInputNode->getName(),
                                                                   memInputNode->getTypeStr(),
                                                                   memInputNode->getOutputShapeAtPort(0),
                                                                   memInputNode->getOriginalOutputPrecisionAtPort(0),
                                                                   graph.getGraphContext(),
-                                                                  inputShape,
-                                                                  inputPrc);
-
+                                                                  inputShapes,
+                                                                  inputPrcs,
+                                                                  subGraph);
         graph.AddNode(memInputSingle);
 
         if (!memInputNode->getParentEdges().empty()) {
-            auto parentEdge = memInputNode->getParentEdgeAt(0);
-            auto parent = parentEdge->getParent();
-            const auto inputNum = parentEdge->getInputNum();
-            graph.RemoveEdge(parentEdge);
-            graph.CreateEdge(parent, memInputSingle, inputNum, 0);
+            auto parentEdgeNum = memInputNode->getParentEdges().size();
+            std::vector<ov::intel_cpu::EdgePtr> parentEdges;
+            for (size_t i = 0; i < parentEdgeNum; i++) {
+                auto parentEdge = memInputNode->getParentEdgeAt(i);
+                auto parent = parentEdge->getParent();
+                const auto inputNum = parentEdge->getInputNum();
+                parentEdges.push_back(parentEdge);
+                graph.CreateEdge(parent, memInputSingle, inputNum, parentEdge->getOutputNum());
+            }
+            for (auto parentEdge : parentEdges) {
+                graph.RemoveEdge(parentEdge);
+            }
         }
 
         for (auto&& edge : memInputNode->getChildEdgesAtPort(0)) {

src/plugins/intel_cpu/src/nodes/input.cpp

@@ -8,6 +8,7 @@
 #include "nodes/node_config.h"
 #include "openvino/core/parallel.hpp"
 #include "shape_inference/shape_inference_pass_through.hpp"
+#include "transformations/cpu_opset/common/op/read_value_with_subgraph.hpp"
 
 using namespace dnnl;
 using namespace dnnl::impl::cpu::x64;
@@ -217,13 +218,14 @@ jit_has_subnormals_base::fn_t jit_has_subnormals_function() {
 #endif
 
 Input::Input(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
-        : Node(op, context, PassThroughShapeInferFactory()) {
+        : Node(op, context, PassThroughShapeInferFactory(op)) {
     if (!one_of(op->get_type_info(),
                 op::v0::Parameter::get_type_info_static(),
                 op::v0::Constant::get_type_info_static(),
                 op::v0::Result::get_type_info_static(),
                 op::v3::ReadValue::get_type_info_static(),
-                op::v6::ReadValue::get_type_info_static()))
+                op::v6::ReadValue::get_type_info_static(),
+                ov::intel_cpu::ReadValueWithSubgraph::get_type_info_static()))
         OPENVINO_THROW_NOT_IMPLEMENTED("CPU Input node doesn't support ngraph operation ",
                                        op->get_type_name(),
                                        " with name ",
@@ -479,7 +481,11 @@ void Input::selectOptimalPrimitiveDescriptor() {
     supportedPrimitiveDescriptors.clear();
 
     // and just use parent memory descriptor for Output node to avoid reorders insertion
-    NodeConfig config({PortConfig(getParentOutputMemDesc(getParentEdgeAt(0)), BlockedMemoryDesc::FULL_MASK, 0)}, {});
+    std::vector<PortConfig> inConfs;
+    for (size_t i = 0; i < getParentEdges().size(); i++) {
+        inConfs.push_back({PortConfig(getParentOutputMemDesc(getParentEdgeAt(i)), BlockedMemoryDesc::FULL_MASK, 0)});
+    }
+    NodeConfig config(inConfs, {});
 
     supportedPrimitiveDescriptors.emplace_back(config, impl_desc_type::unknown);
     selectPrimitiveDescriptorByIndex(0);