From 2cb8222dd7bf443096f80e09bcba8766c223f680 Mon Sep 17 00:00:00 2001
From: Egor Duplenskii <egor.duplensky@gmail.com>
Date: Mon, 21 Oct 2024 10:38:52 +0200
Subject: [PATCH] [CPU] Use actual input shape to init desc for MemoryInputSDPA
 (#27143)

An output shape was previously used to create an input descriptor for
some reason
---
 src/plugins/intel_cpu/src/nodes/memory.cpp | 52 +++-------------------
 src/plugins/intel_cpu/src/nodes/memory.hpp |  2 -
 2 files changed, 7 insertions(+), 47 deletions(-)

diff --git a/src/plugins/intel_cpu/src/nodes/memory.cpp b/src/plugins/intel_cpu/src/nodes/memory.cpp
index 88693ebfa49fdf..756fbc5b578f61 100644
--- a/src/plugins/intel_cpu/src/nodes/memory.cpp
+++ b/src/plugins/intel_cpu/src/nodes/memory.cpp
@@ -427,29 +427,20 @@ void MemoryInputBase::initSupportedPrimitiveDescriptors() {
     if (!supportedPrimitiveDescriptors.empty())
         return;
 
-    auto&& shape = getOutputShapeAtPort(0);
     auto precision = getOriginalOutputPrecisionAtPort(0);
     auto&& descCreators = ov::intel_cpu::BlockedDescCreator::getCommonCreators();
-
     NodeConfig config;
 
     if (!getParentEdges().empty()) {
-        PortConfig inPortConfig;
-
-        inPortConfig.inPlace(-1);
-        inPortConfig.constant(false);
-        inPortConfig.setMemDesc(descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, shape));
-
-        config.inConfs.push_back(std::move(inPortConfig));
+        const auto& inputShape = getInputShapeAtPort(0);
+        config.inConfs.emplace_back(descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, inputShape));
     }
 
-    PortConfig outPortConfig;
-
-    outPortConfig.inPlace(0);
-    outPortConfig.constant(false);
-    outPortConfig.setMemDesc(descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, shape));
-
-    config.outConfs.push_back(std::move(outPortConfig));
+    const auto& outputShape = getOutputShapeAtPort(0);
+    config.outConfs.emplace_back(
+        descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, outputShape),
+        BlockedMemoryDesc::FULL_MASK,
+        0);
 
     supportedPrimitiveDescriptors.emplace_back(config, impl_desc_type::unknown);
 }
@@ -759,35 +750,6 @@ void MemoryInputSDPA::createPrimitive() {
     OPENVINO_ASSERT(m_child_port_idx != -1, getName(), " should be connected to SDPA node.");
 }
 
-void MemoryInputSDPA::initSupportedPrimitiveDescriptors() {
-    if (!supportedPrimitiveDescriptors.empty())
-        return;
-
-    auto&& shape = getOutputShapeAtPort(0);
-    auto precision = getOriginalOutputPrecisionAtPort(0);
-    auto&& descCreators = ov::intel_cpu::BlockedDescCreator::getCommonCreators();
-    NodeConfig config;
-    if (!getParentEdges().empty()) {
-        PortConfig inPortConfig;
-        inPortConfig.inPlace(-1);
-        inPortConfig.constant(false);
-        inPortConfig.setMemDesc(descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, shape));
-        config.inConfs.push_back(std::move(inPortConfig));
-    }
-
-    PortConfig outPortConfig;
-    outPortConfig.inPlace(0);
-    outPortConfig.constant(false);
-    // layout for fake memory obj, the child sdpa also does not use it
-    outPortConfig.setMemDesc(descCreators.at(LayoutType::ncsp)->createSharedDesc(precision, shape));
-    config.outConfs.push_back(std::move(outPortConfig));
-    supportedPrimitiveDescriptors.emplace_back(config, impl_desc_type::unknown);
-}
-
-void MemoryInputSDPA::initOptimalPrimitiveDescriptor() {
-    Node::initOptimalPrimitiveDescriptor();
-}
-
 void MemoryInputSDPA::assignStateHook() {
     auto currentState = getAssignedState();
     auto sdpaNode = m_sdpaNode.lock();
diff --git a/src/plugins/intel_cpu/src/nodes/memory.hpp b/src/plugins/intel_cpu/src/nodes/memory.hpp
index c5a83cfa5cad1a..c158d738a36148 100644
--- a/src/plugins/intel_cpu/src/nodes/memory.hpp
+++ b/src/plugins/intel_cpu/src/nodes/memory.hpp
@@ -204,8 +204,6 @@ class MemoryInputSDPA : public MemoryInputBase {
     static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
 
     void createPrimitive() override;
-    void initSupportedPrimitiveDescriptors() override;
-    void initOptimalPrimitiveDescriptor() override;
     void resolveInPlaceEdges(Edge::LOOK look) override;
 
     MemStatePtr makeState() const override;