[CPU]whisper readvalue optimize

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -139,6 +139,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"
@@ -76,6 +77,7 @@ class TypeRelaxedExtension : public ov::OpExtension<ov::op::TypeRelaxed<Op>> {
     OP_EXTENSION(ov::intel_cpu::CausalMaskPreprocessNode)                   \
     OP_EXTENSION(ov::intel_cpu::SwishNode)                                  \
     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.cpp

@@ -940,6 +940,13 @@ void Graph::Allocate(const std::vector<size_t>& syncNodesInds) {
                 one_of(edge->getChild()->getType(), Type::Output, Type::MemoryOutput) &&
                 edge->inPlace(Edge::LOOK_DOWN)) {
                 edge->getChild()->resolveInPlaceEdges(Edge::LOOK_DOWN);
+            } else if (one_of(edge->getParent()->getType(), Type::MemoryInput)) {
+                auto memInp = std::dynamic_pointer_cast<node::MemoryInput>(edge->getParent());
+                if (memInp && memInp->haveSubgraph()) {
+                    // Since the ReadValueWithSubgraph is middle node, just add this branch in order to use
+                    // ProxyMemoryBlock to share memory
+                    edge->getParent()->resolveInPlaceEdges(Edge::LOOK_UP);
+                }
             }
         }
     }

src/plugins/intel_cpu/src/graph_optimizer.cpp

@@ -183,6 +183,10 @@ void GraphOptimizer::ApplyCommonGraphOptimizations(Graph &graph) {
     MatchSdpaKvCache(graph);
     graph.RemoveDroppedNodes();
 
+    OV_ITT_SCOPE_NEXT(FIRST_INFERENCE, taskChain, "ReplaceMemoryOutputWithMemoryOutputStub");
+    ReplaceMemoryOutputWithMemoryOutputStub(graph);
+    graph.RemoveDroppedNodes();
+
     OV_ITT_SCOPE_NEXT(FIRST_INFERENCE, taskChain, "RemoveDroppedEdges");
     graph.RemoveDroppedEdges();
 }
@@ -3064,6 +3068,90 @@ void GraphOptimizer::RemoveConvertMemoryOutput(Graph &graph) {
     }
 }
 
+void GraphOptimizer::ReplaceMemoryOutputWithMemoryOutputStub(Graph& graph) {
+    auto& graphNodes = graph.GetNodes();
+
+    auto isSuitableMemInput = [](const NodePtr& node) -> bool {
+        if (Type::MemoryInput != node->getType()) {
+            return false;
+        }
+        auto memInput = std::dynamic_pointer_cast<node::MemoryInput>(node);
+        if (memInput) {
+            return memInput->haveSubgraph();
+        }
+
+        return false;
+    };
+
+    for (size_t i = 0; i < graphNodes.size(); i++) {
+        auto node = graphNodes[i];
+        if (!isSuitableMemInput(node)) {
+            continue;
+        }
+
+        CPU_GRAPH_OPTIMIZER_SCOPE(ReplaceMemoryOutputWithMemoryOutputStub);
+
+        auto memoryNode = std::dynamic_pointer_cast<node::MemoryNode>(node);
+        if (nullptr == memoryNode) {
+            continue;
+        }
+
+        // Find sibling MemoryOutput
+        std::shared_ptr<MemoryOutput> memOutput = nullptr;
+        bool isReplaced = false;
+        for (auto&& edge : node->getChildEdgesAtPort(0)) {
+            auto child = edge->getChild();
+            if (Type::MemoryOutput == child->getType()) {
+                memOutput = std::dynamic_pointer_cast<MemoryOutput>(child);
+                if (memOutput && memOutput->getId() == memoryNode->getId()) {
+                    break;
+                }
+
+                auto memOutputStub = std::dynamic_pointer_cast<MemoryOutputStub>(child);
+                if (memOutputStub && memOutputStub->getId() == memoryNode->getId()) {
+                    isReplaced = true;
+                    break;
+                }
+            }
+        }
+
+        if (isReplaced) {
+            continue;
+        }
+        if (memOutput == nullptr) {
+            OPENVINO_THROW("Can't find ", node->getName(), " corresponding sibling node.");
+        }
+
+        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;
+
+        if (!node->getParentEdges().empty()) {
+            input_shape = ov::optional<Shape>(node->getInputShapeAtPort(0));
+            input_prc = ov::optional<ov::element::Type>(node->getOriginalInputPrecisionAtPort(0));
+        }
+
+        // Capture reference to the original mem output before graph transformations
+        auto& memOutputBase = memInputNode->getOutputNode();
+
+        // Create a stub memory output
+        auto memOutputStub = std::make_shared<MemoryOutputStub>(memOutputBase.getId(),
+                                                                memOutputBase.getName() + "_MemoryOutputStub",
+                                                                memOutputBase.getTypeStr(),
+                                                                memOutputBase.getInputShapeAtPort(0),
+                                                                memOutputBase.getOriginalInputPrecisionAtPort(0),
+                                                                graph.getGraphContext());
+
+        auto memOutputEdge = memOutputBase.getParentEdgeAt(0);
+        const auto inputNum = memOutputEdge->getInputNum();
+        graph.RemoveEdge(memOutputEdge);
+        graph.CreateEdge(node, memOutputStub, inputNum, 0);
+        graph.AddNode(memOutputStub);
+    }
+}
+
 void GraphOptimizer::MatchSdpaKvCache(Graph &graph) {
     auto& graphNodes = graph.GetNodes();
 

src/plugins/intel_cpu/src/graph_optimizer.h

@@ -52,6 +52,7 @@ class GraphOptimizer {
     void RemoveMemoryInputConvert(Graph &graph);
     void RemoveConvertMemoryOutput(Graph &graph);
     void MatchSdpaKvCache(Graph &graph);
+    void ReplaceMemoryOutputWithMemoryOutputStub(Graph &graph);
 
     bool canBeInplaced(const NodePtr& parentNode, const NodePtr& childNode);
     // Method checks that after the sequential execution of Transpose and Reorder nodes,

src/plugins/intel_cpu/src/memory_state.cpp

@@ -156,6 +156,45 @@ MemoryPtr VariableStateDoubleBuffer::internal_state_mem() const {
     return prime_mem();
 }
 
+VariableStateSingleBuffer::VariableStateSingleBuffer(const std::string& name,
+                                                     const MemoryPtr& external_buffer,
+                                                     const MemoryDescPtr& external_desc)
+    : VariableStateBase(name, external_desc) {
+    OPENVINO_ASSERT(external_buffer);
+    m_internal_mem = external_buffer;
+    m_internal_desc = m_internal_mem->getDescPtr();
+    auto&& shape = m_internal_desc->getShape();
+
+    if (shape.isStatic()) {
+        m_internal_mem->nullify();
+    } else {
+        // in the case of the original desc has dynamic shape we create an empty tensor
+        auto new_desc = to_static(m_internal_desc);
+        m_internal_mem->redefineDesc(new_desc);
+    }
+}
+MemoryPtr VariableStateSingleBuffer::input_mem() {
+    return m_internal_mem;
+}
+MemoryPtr VariableStateSingleBuffer::output_mem() {
+    return m_internal_mem;
+}
+MemoryDescPtr VariableStateSingleBuffer::internal_desc() const {
+    return m_internal_desc;
+}
+
+void VariableStateSingleBuffer::reset_impl() {
+    auto new_desc = to_static(m_internal_desc);
+    if (m_internal_mem) {
+        m_internal_mem->redefineDesc(new_desc);
+        m_internal_mem->nullify();
+    }
+}
+
+MemoryPtr VariableStateSingleBuffer::internal_state_mem() const {
+    return m_internal_mem;
+}
+
 VariableStateKVcache::VariableStateKVcache(
     const std::string& name,
     const MemoryDescPtr& external_desc,

src/plugins/intel_cpu/src/memory_state.h

@@ -94,6 +94,27 @@ class VariableStateDoubleBuffer : public VariableStateBase {
     size_t buffer_num = 0;
 };
 
+class VariableStateSingleBuffer : public VariableStateBase {
+public:
+    VariableStateSingleBuffer(const std::string& name,
+                              const MemoryPtr& external_buffer,
+                              const MemoryDescPtr& external_desc);
+
+    MemoryPtr input_mem() override;
+    MemoryPtr output_mem() override;
+    MemoryDescPtr internal_desc() const override;
+
+private:
+    void reset_impl() override;
+    void commit_impl() override {};
+
+    MemoryPtr internal_state_mem() const override;
+
+private:
+    MemoryDescPtr m_internal_desc; //mem desc required by the graph internal tensor
+    MemoryPtr m_internal_mem;
+};
+
 class VariableStateKVcache : public VariableStateBase {
 public:
     VariableStateKVcache(const std::string& name,

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;
@@ -223,7 +224,8 @@ Input::Input(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr conte
                 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 ",