Update evaluate methods of ReadValue/Assign operations (#21372)

* Update evaluate methods of ReadValue/Assign operations

* Update evaluate methods of ReadValue/Assign operations

* Prepare IVariable interface

* Support initializing subgraph

* Add exception types to the edge conflict searching algo

* Add test for init graph edge conflict

* Remove ShapeOf from consumer nodes in memory conflicts detection algo

* update template plugin

* Fix template tests; synchronize reset flag between VariableState and VariableValue objects

* fix legacy (Assign/ReadValue v3 + CNN Network) tests

* fix python tests

* fix python tests

* resolve review comments, apply cpu patch

---------

Co-authored-by: Maksim Kutakov <maksim.kutakov@intel.com>

src/bindings/python/tests/test_graph/test_basic.py

@@ -564,8 +564,8 @@ def test_sink_model_ctor():
     sinks = model.get_sinks()
     assert ["Assign"] == [sink.get_type_name() for sink in sinks]
     assert model.sinks[0].get_output_shape(0) == Shape([2, 2])
-    assert op_types == ["Parameter", "ReadValue", "Add", "Assign", "Result"]
-    assert len(model.get_ops()) == 5
+    assert op_types == ["Parameter", "Constant", "ReadValue", "Add", "Assign", "Result"]
+    assert len(model.get_ops()) == 6
     assert model.get_output_size() == 1
     assert model.get_output_op(0).get_type_name() == "Result"
     assert model.get_output_element_type(0) == model.get_parameters()[0].get_element_type()

src/bindings/python/tests/utils/helpers.py

@@ -211,7 +211,8 @@ def generate_add_model() -> openvino._pyopenvino.Model:
 
 def generate_model_with_memory(input_shape, data_type) -> openvino._pyopenvino.Model:
     input_data = ops.parameter(input_shape, name="input_data", dtype=data_type)
-    rv = ops.read_value(input_data, "var_id_667", data_type, input_shape)
+    init_val = ops.constant(np.zeros(input_shape), data_type)
+    rv = ops.read_value(init_val, "var_id_667", data_type, input_shape)
     add = ops.add(rv, input_data, name="MemoryAdd")
     node = ops.assign(add, "var_id_667")
     res = ops.result(add, "res")

src/bindings/python/tests_compatibility/test_inference_engine/test_InferRequest.py

@@ -19,7 +19,8 @@ test_net_xml, test_net_bin = model_path()
 
 def create_function_with_memory(input_shape, data_type):
     input_data = ng.parameter(input_shape, name="input_data", dtype=data_type)
-    rv = ng.read_value(input_data, "var_id_667")
+    init_val = ng.constant(np.zeros(input_shape), data_type)
+    rv = ng.read_value(init_val, "var_id_667")
     add = ng.add(rv, input_data, name="MemoryAdd")
     node = ng.assign(add, "var_id_667")
     res = ng.result(add, "res")

src/core/src/op/assign.cpp

@@ -119,12 +119,12 @@ bool Assign::evaluate(TensorVector& outputs,
 
     const auto var_value = variable_values.find(m_variable)->second;
     var_value->set_reset(false);
-    auto buffer = var_value->get_state();
-    buffer.set_shape(inputs[0].get_shape());
+    auto memory_buffer = var_value->get_state();
+    memory_buffer.set_shape(inputs[0].get_shape());
     outputs[0].set_shape(inputs[0].get_shape());
 
     std::memcpy(outputs[0].data(), inputs[0].data(), inputs[0].get_byte_size());
-    std::memcpy(buffer.data(), inputs[0].data(), inputs[0].get_byte_size());
+    std::memcpy(memory_buffer.data(), inputs[0].data(), inputs[0].get_byte_size());
 
     return true;
 }

src/core/src/op/read_value.cpp

@@ -153,14 +153,9 @@ bool ReadValue::evaluate(TensorVector& outputs,
     const auto use_context = var_value != variable_values.end() && !var_value->second->get_reset();
 
     auto& output = outputs[0];
-    if (use_context) {
-        const auto& ctx_tensor = var_value->second->get_state();
-        output.set_shape(ctx_tensor.get_shape());
-        std::memcpy(output.data(), ctx_tensor.data(), output.get_byte_size());
-    } else {
-        output.set_shape(inputs[0].get_shape());
-        std::memset(output.data(), 0, output.get_byte_size());
-    }
+    const auto& input = use_context ? var_value->second->get_state() : inputs[0];
+    output.set_shape(input.get_shape());
+    std::memcpy(output.data(), input.data(), output.get_byte_size());
     return true;
 }
 

src/plugins/intel_cpu/src/edge.cpp

@@ -76,7 +76,12 @@ void Edge::collectConsumers(std::vector<NodePtr>& result) const {
             }
         }
     } else {
-        result.push_back(this->getChild());
+        auto childNode = this->getChild();
+        if (Type::ShapeOf == childNode->getType()) {
+            // ShapeOf doesn't actually read the data, it only reads shape
+            return;
+        }
+        result.push_back(childNode);
     }
 }
 
@@ -554,7 +559,8 @@ NodePtr Edge::modifiedInPlace() const {
         auto& outConfs = childSPD->getConfig().outConfs;
         for (size_t i = 0; i < outConfs.size(); ++i) {
             const auto& conf = outConfs[i];
-            if (childPort < 0 || conf.inPlace() != childPort) {
+            if (childPort < 0 || conf.inPlace() != childPort ||
+                Type::MemoryInput == childNode->getType()) { //exception type, it doesn't modify memory
                 continue;
             }
             if (childNode->isExecutable()) {

src/plugins/intel_cpu/src/memory_state.cpp

@@ -49,6 +49,7 @@ void VariableStateBase::set_state(const ov::SoPtr<ov::ITensor>& state) {
 
     Memory mem(get_engine(), state_desc, src);
     input_mem()->load(mem);
+    reset_state_flag = false;
 }
 
 ov::SoPtr<ov::ITensor> VariableStateBase::get_state() const {
@@ -80,11 +81,24 @@ ov::SoPtr<ov::ITensor> VariableStateBase::get_state() const {
     return std::make_shared<Tensor>(mem);
 }
 
+void VariableStateBase::reset() {
+    reset_impl();
+    reset_state_flag = true;
+}
+
+bool VariableStateBase::is_reset_state() const {
+    return reset_state_flag;
+}
+
+void VariableStateBase::commit() {
+    commit_impl();
+    reset_state_flag = false;
+}
+
 VariableStateDoubleBuffer::VariableStateDoubleBuffer(const std::string& name,
                                                      const MemoryPtr& first_buffer,
                                                      const MemoryPtr& second_buffer,
-                                                     const MemoryDescPtr& external_desc,
-                                                     const MemoryCPtr& init_val) :
+                                                     const MemoryDescPtr& external_desc) :
     VariableStateBase(name, external_desc) {
     OPENVINO_ASSERT(first_buffer && second_buffer);
     reset_prime_mem(first_buffer);
@@ -94,11 +108,7 @@ VariableStateDoubleBuffer::VariableStateDoubleBuffer(const std::string& name,
     //TODO what if by some reason we already have internal static state while the node is dynamic, is it even possible?
 
     if (shape.isStatic()) {
-        if (init_val) {
-            prime_mem()->load(*init_val);
-        } else {
-            prime_mem()->nullify();
-        }
+        prime_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);
@@ -106,7 +116,7 @@ VariableStateDoubleBuffer::VariableStateDoubleBuffer(const std::string& name,
     }
 }
 
-void VariableStateDoubleBuffer::reset() {
+void VariableStateDoubleBuffer::reset_impl() {
     auto new_desc = to_static(m_internal_desc);
     for (auto&& mem : m_internal_mem) {
         if (mem) {
@@ -116,7 +126,7 @@ void VariableStateDoubleBuffer::reset() {
     }
 }
 
-void VariableStateDoubleBuffer::commit() {
+void VariableStateDoubleBuffer::commit_impl() {
     buffer_num ^= 0x01;
 }
 
@@ -138,8 +148,7 @@ MemoryPtr VariableStateDoubleBuffer::internal_state_mem() const {
 
 VariableStateSingleBuffer::VariableStateSingleBuffer(const std::string& name,
                                                      const MemoryPtr& buffer,
-                                                     const MemoryDescPtr& external_desc,
-                                                     const MemoryCPtr& init_val) :
+                                                     const MemoryDescPtr& external_desc) :
     VariableStateBase(name, external_desc) {
     OPENVINO_ASSERT(buffer);
     m_internal_mem = buffer;
@@ -148,11 +157,7 @@ VariableStateSingleBuffer::VariableStateSingleBuffer(const std::string& name,
     //TODO what if by some reason we already have internal static state while the node is dynamic, is it even possible?
 
     if (shape.isStatic()) {
-        if (init_val) {
-            m_internal_mem->load(*init_val);
-        } else {
-            m_internal_mem->nullify();
-        }
+        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);
@@ -160,7 +165,7 @@ VariableStateSingleBuffer::VariableStateSingleBuffer(const std::string& name,
     }
 }
 
-void VariableStateSingleBuffer::reset() {
+void VariableStateSingleBuffer::reset_impl() {
     auto new_desc = to_static(m_internal_desc);
     m_internal_mem->redefineDesc(new_desc);
     m_internal_mem->nullify();
@@ -182,7 +187,7 @@ MemoryPtr VariableStateSingleBuffer::internal_state_mem() const {
     return m_internal_mem;
 }
 
-void VariableStateSingleBuffer::commit() {
+void VariableStateSingleBuffer::commit_impl() {
     //nothing to do
 }
 

src/plugins/intel_cpu/src/memory_state.h

@@ -26,6 +26,7 @@ public:
     virtual MemoryPtr input_mem() = 0;
     virtual MemoryPtr output_mem() = 0;
     virtual MemoryDescPtr internal_desc() const = 0;
+    virtual bool is_reset_state() const = 0;
 };
 
 class VariableStateBase : public IVariableState {
@@ -35,15 +36,21 @@ public:
     //ov::IVariableState
     void set_state(const ov::SoPtr<ov::ITensor>& state) override;
     ov::SoPtr<ov::ITensor> get_state() const override;
+    void reset() override final; // NOLINT
+    bool is_reset_state() const override final; // NOLINT
+    void commit() override final; // NOLINT
 
 protected:
     virtual MemoryPtr internal_state_mem() const = 0;
+    virtual void reset_impl() = 0;
+    virtual void commit_impl() = 0;
 
     static MemoryDescPtr to_static(const MemoryDescPtr& desc);
     static const dnnl::engine& get_engine();
 
-protected:
+private:
     MemoryDescPtr m_external_desc;
+    bool reset_state_flag = true;
 };
 
 class VariableStateDoubleBuffer : public VariableStateBase {
@@ -51,20 +58,17 @@ public:
     VariableStateDoubleBuffer(const std::string& name,
                               const MemoryPtr& first_buffer,
                               const MemoryPtr& second_buffer,
-                              const MemoryDescPtr& external_desc,
-                              const MemoryCPtr& init_val);
-
-    //ov::IVariableState
-    void reset() override;
-
-    //ov::intel_cpu::IVariableState
-    void commit() override;
+                              const MemoryDescPtr& external_desc);
 
     MemoryPtr input_mem() override;
     MemoryPtr output_mem() override;
     MemoryDescPtr internal_desc() const override;
 
 private:
+    //ov::intel_cpu::VariableStateBase
+    void reset_impl() override;
+    void commit_impl() override;
+
     void reset_prime_mem(const MemoryPtr& mem) {
         m_internal_mem[buffer_num] = mem;
     }
@@ -93,19 +97,17 @@ class VariableStateSingleBuffer : public VariableStateBase {
 public:
     VariableStateSingleBuffer(const std::string& name,
                               const MemoryPtr& buffer,
-                              const MemoryDescPtr& external_desc,
-                              const MemoryCPtr& init_val);
-    //ov::IVariableState
-    void reset() override;
-
-    //ov::intel_cpu::IVariableState
-    void commit() override;
+                              const MemoryDescPtr& external_desc);
 
     MemoryPtr input_mem() override;
     MemoryPtr output_mem() override;
     MemoryDescPtr internal_desc() const override;
 
 private:
+    //ov::intel_cpu::VariableStateBase
+    void reset_impl() override;
+    void commit_impl() override;
+
     MemoryPtr internal_state_mem() const override;
 
 private:

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

@@ -501,10 +501,6 @@ void Input::initSupportedPdFromMemDesc() {
     supportedPrimitiveDescriptors.emplace_back(std::move(config), impl_desc_type::unknown);
 }
 
-void Input::resetMemoryPtr(const MemoryCPtr& mem) {
-    memoryPtr = mem;
-}
-
 }   // namespace node
 }   // namespace intel_cpu
 }   // namespace ov

src/plugins/intel_cpu/src/nodes/input.h

@@ -38,9 +38,6 @@ public:
     bool needShapeInfer() const override { return false; }
     bool needPrepareParams() const override { return false; }
 
-protected:
-    void resetMemoryPtr(const MemoryCPtr& mem);
-
 private:
     void cloneBlobIfRequired();
     void initSupportedPdDefault();

src/plugins/intel_cpu/src/nodes/memory.cpp

@@ -260,17 +260,6 @@ MemoryInputBase::MemoryInputBase(const std::string id,
     // this is their responsibility to link the input/output nodes properly
 }
 
-void MemoryInputBase::createPrimitive() {
-    Input::createPrimitive();
-    if (!inputShapes.empty()) {
-        auto parentEdge = getParentEdgeAt(0);
-
-        if (parentEdge->getParent()->isConstant()) {
-            Input::resetMemoryPtr(parentEdge->getMemoryPtr());
-        }
-    }
-}
-
 void MemoryInputBase::resolveInPlaceEdges(Edge::LOOK look) {
     if (!(look & Edge::LOOK_UP)) {
         Node::resolveInPlaceEdges(look);
@@ -308,6 +297,12 @@ MemoryOutput& MemoryInputBase::getOutputNode() {
 void MemoryInputBase::assignState(MemStatePtr newState) {
     assignedMem = newState->input_mem();
 
+    if (!getParentEdges().empty() && newState->is_reset_state()) {
+        isExecutableFlag = true;
+    } else {
+        isExecutableFlag = false;
+    }
+
     OPENVINO_ASSERT(assignedMem,
         "MemoryInput ",
         getName(),
@@ -342,22 +337,44 @@ void MemoryInputBase::assignState(MemStatePtr newState) {
         memMngr->reset();
     }
 
-    const auto& edges = getChildEdgesAtPort(0);
-    if (isDynamicNode()) {
-        for (auto&& edge : edges) {
-            edge->getMemoryPtr()->redefineDesc(internDesc);
+    if (!isExecutableFlag) {
+        const auto& edges = getChildEdgesAtPort(0);
+        if (isDynamicNode()) {
+            for (auto&& edge : edges) {
+                edge->getMemoryPtr()->redefineDesc(internDesc);
+            }
+        }
+
+        auto outMem = edges.front()->getMemoryPtr();
+
+        if (outMem->getData() != assignedMem->getData()) {
+            outMem->load(*assignedMem);
         }
     }
 
-    auto outMem = edges.front()->getMemoryPtr();
-
-    if (outMem->getData() != assignedMem->getData()) {
-        outMem->load(*assignedMem);
-    }
-
     getOutputNode().assignExtMemory(newState->output_mem(), newState->internal_desc());
 }
 
+bool MemoryInputBase::needShapeInfer() const {
+    return isExecutableFlag;
+}
+
+bool MemoryInputBase::isExecutable() const {
+    return isExecutableFlag && Node::isExecutable();
+}
+
+void MemoryInputBase::executeDynamicImpl(dnnl::stream strm) {
+    execute(strm);
+}
+
+void MemoryInputBase::execute(dnnl::stream strm) {
+    if (!isExecutableFlag) return;
+
+    auto&& src = getParentEdgeAt(0)->getMemory();
+    auto&& dst = getChildEdgesAtPort(0).front()->getMemoryPtr();
+    dst->load(src);
+}
+
 void MemoryInputBase::registerOutputNode(MemoryOutput* node) {
     if (outputNode == node) { return; }
     if (outputNode) { outputNode->deregisterSibling(this); }
@@ -514,8 +531,7 @@ MemStatePtr MemoryInput::makeState() const {
     return std::make_shared<VariableStateDoubleBuffer>(state_name,
         std::make_shared<Memory>(eng, mem_desc),
         std::make_shared<Memory>(eng, mem_desc),
-        original_desc,
-        getMemoryPtr());
+        original_desc);
 }
 
 bool MemoryInput::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept {
@@ -606,10 +622,10 @@ MemStatePtr MemoryInputSDPA::makeState() const {
         state_name = state_name.substr(0, suffix_idx);
     }
 
-    return std::make_shared<VariableStateSingleBuffer>(state_name,
-        std::make_shared<Memory>(eng, mem_desc, std::make_shared<DnnlMemoryMngr>(make_unique<MemoryMngrRealloc>())),
-        original_desc,
-        getMemoryPtr());
+    auto internal_memory =
+        std::make_shared<Memory>(eng, mem_desc, std::make_shared<DnnlMemoryMngr>(make_unique<MemoryMngrRealloc>()));
+
+    return std::make_shared<VariableStateSingleBuffer>(state_name, internal_memory, original_desc);
 }
 
 }   // namespace node

src/plugins/intel_cpu/src/nodes/memory.hpp

@@ -127,10 +127,10 @@ public:
         return getType() == Type::MemoryInput;
     }
 
-    void execute(dnnl::stream strm) override {/*pass*/}
-    void executeDynamicImpl(dnnl::stream strm) override {/*pass*/}
-
-    void createPrimitive() override;
+    bool needShapeInfer() const override;
+    bool isExecutable() const override;
+    void execute(dnnl::stream strm) override;
+    void executeDynamicImpl(dnnl::stream strm) override;
 
     void resolveInPlaceEdges(Edge::LOOK look) override;
 
@@ -149,6 +149,7 @@ private:
     MemoryPtr assignedMem = nullptr;
     MemoryNodeVirtualEdge::Holder* holder = nullptr;
     ProxyMemoryMngrPtr memMngr = nullptr;
+    bool isExecutableFlag = true;
 };
 
 class MemoryInput : public MemoryInputBase {

src/plugins/intel_cpu/tests/functional/subgraph_tests/src/init_state_inplace_conflicts.cpp

@@ -0,0 +1,65 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <shared_test_classes/base/ov_subgraph.hpp>
+#include "test_utils/cpu_test_utils.hpp"
+
+using namespace CPUTestUtils;
+using namespace ov::test;
+
+// ┌─────────┐    ┌───────────┐
+// │ Param1  │────| ReadValue │..
+// └───┬──┬──┘    └─────┬─────┘ .
+//     │  │             │       .
+//     │  │             │       .
+//     │  └──┬─────┬────┘       .
+//     │     │ Add │            .
+//     │     └──┬──┘            .
+//     │        │               .
+//     │    ┌───┴────┐          .
+//     └────┤ Concat │          .
+//          └────────┘          .
+//             / \              .
+//            /   \             .
+//   ┌────────┐   ┌────────┐    .
+//   │ Result │   │ Assign │.....
+//   └────────┘   └────────┘
+
+namespace CPUSubgraphTestsDefinitions {
+
+class DynamicShapeStatefulModel : public SubgraphBaseTest, public CPUTestsBase {
+public:
+    void SetUp() override {
+        targetDevice = utils::DEVICE_CPU;
+        ov::element::Type netPrc = ElementType::f32;
+
+        const ov::Shape inpShape = {1, 1};
+        const InputShape input_shape = {{-1, 1}, {}};
+        init_input_shapes({input_shape});
+
+        auto arg = std::make_shared<ov::op::v0::Parameter>(netPrc, inputDynamicShapes.front());
+
+        // The ReadValue/Assign operations must be used in pairs in the model.
+        // For each such a pair, its own variable object must be created.
+        const std::string variable_name("variable0");
+        auto variable = std::make_shared<ov::op::util::Variable>(
+            ov::op::util::VariableInfo{inputDynamicShapes.front(), netPrc, variable_name});
+
+        // Creating ov::Model
+        auto read = std::make_shared<ov::op::v6::ReadValue>(arg, variable);
+        std::vector<std::shared_ptr<ov::Node>> args = {arg, read};
+        auto add = std::make_shared<ov::op::v1::Add>(arg, read);
+        constexpr int concat_axis = 0;
+        auto concat = std::make_shared<ov::op::v0::Concat>(ov::NodeVector{arg, add}, concat_axis);
+        auto assign = std::make_shared<ov::op::v6::Assign>(concat, variable);
+        auto res = std::make_shared<ov::op::v0::Result>(concat);
+        function = std::make_shared<ov::Model>(ov::ResultVector({res}), ov::SinkVector({assign}), ov::ParameterVector({arg}));
+    }
+};
+
+TEST_F(DynamicShapeStatefulModel, smoke_StateInitGraphIplaceConflicts) {
+    compile_model();
+    CheckNumberOfNodesWithType(compiledModel, "Reorder", 0);
+}
+} // namespace CPUSubgraphTestsDefinitions

src/plugins/intel_cpu/tests/unit/graph/memory_state.cpp

@@ -8,6 +8,8 @@
 #include "ov_models/builders.hpp"
 #include "nodes/memory.hpp"
 #include "nodes/softmax.h"
+#include "nodes/shapeof.h"
+#include "nodes/convert.h"
 
 using namespace ov::intel_cpu;
 
@@ -199,4 +201,110 @@ TEST(MemStateGraphTest, smoke_Check_Memory_Modification_Guard) {
         auto state_out_mem = state->output_mem()->getData();
         ASSERT_NE(state_out_mem, memory_output_inp_mem);
     }
+}
+
+/*
+
+           ┌────────┐
+           │ Param  │
+           └───┬────┘
+               │
+               │
+           ┌───┴────┐
+           │MemInput│
+           └────────┘
+            /     \
+           /       \
+   ┌────────┐     ┌────────┐
+   │DummyUp │     │ ShapeOf│
+   └───┬────┘     └───┬────┘
+       │              │
+       │              │
+   ┌───┴────┐     ┌───┴────┐
+   │ MemOut │     │Subgraph│
+   └────────┘     └───┬────┘
+                      │
+                      │
+                  ┌───┴────┐
+                  │ Output │
+                  └────────┘
+
+*/
+
+TEST(MemStateGraphTest, smoke_ShapeOf_no_Inplace_Conflicts) {
+    const ov::element::Type_t test_prec = ov::element::Type_t::f32;
+    const ov::Shape test_shape{1, 3, 8, 16};
+
+    auto param = std::make_shared<ov::op::v0::Parameter>(test_prec, test_shape);
+
+    // The ReadValue/Assign operations must be used in pairs in the model.
+    // For each such a pair, its own variable object must be created.
+    const std::string variable_name("variable0");
+    auto variable = std::make_shared<ov::op::util::Variable>(
+        ov::op::util::VariableInfo{test_shape, test_prec, variable_name});
+
+    // creat ngraph ops to build CPU nodes
+    auto read = std::make_shared<ov::op::v6::ReadValue>(param, variable);
+    auto shapeof = std::make_shared<ov::op::v0::ShapeOf>(read);
+    auto convert = std::make_shared<ov::op::v0::Convert>(shapeof, test_prec);
+    auto softmax = std::make_shared<ov::op::v1::Softmax>(convert, 0);
+    auto assign = std::make_shared<ov::op::v6::Assign>(read, variable);
+    auto res = std::make_shared<ov::op::v0::Result>(softmax);
+
+    std::unordered_set<NodePtr> nodes_set;
+    std::vector<EdgePtr> graph_edges;
+
+    auto add_edge = [&](const NodePtr& parent, const NodePtr& child, size_t parentPort, size_t childPort) -> void {
+        auto edge = std::make_shared<Edge>(parent, child, parentPort, childPort);
+        child->addEdge(edge);
+        graph_edges.push_back(edge);
+        nodes_set.insert(parent);
+        nodes_set.insert(child);
+    };
+
+    //create graph context
+
+    Config conf;
+    conf.rtCacheCapacity = 0;
+    auto context = std::make_shared<GraphContext>(conf, nullptr, nullptr, false);
+
+    auto input_node = std::make_shared<node::Input>(param, context);
+    auto memory_input = std::make_shared<node::MemoryInput>(read, context);
+    auto dummy = std::make_shared<cpu_unit_test::DummyNode>(
+        test_shape, test_prec,
+        "first_dummy",
+        "DummyNode",
+        context,
+        LayoutType::ncsp,
+        Edge::LOOK::LOOK_UP,
+        true);
+
+    auto memory_output = std::make_shared<node::MemoryOutput>(assign, context);
+    auto shapeof_node = std::make_shared<node::ShapeOf>(shapeof, context);
+    auto softmax_node = std::make_shared<node::SoftMax>(softmax, context);
+    auto convert_node = std::make_shared<node::Convert>(convert, context);
+    auto output_node = std::make_shared<node::Input>(res, context);
+
+    convert_node->setOriginalInputPrecisionAtPort(0, ov::element::i32);
+
+    add_edge(input_node, memory_input, 0, 0);
+    add_edge(memory_input, shapeof_node, 0, 0);
+    add_edge(memory_input, dummy, 0, 0);
+
+    add_edge(dummy, memory_output, 0, 0);
+
+    add_edge(shapeof_node, convert_node, 0, 0);
+    add_edge(convert_node, softmax_node, 0, 0);
+    add_edge(softmax_node, output_node, 0, 0);
+
+    std::vector<NodePtr> graph_nodes(nodes_set.begin(), nodes_set.end());
+
+    Graph graph;
+    graph.CreateGraph(graph_nodes, graph_edges, context, "test_graph");
+
+    auto&& nodes = graph.GetNodes();
+    auto itr = std::find_if(nodes.begin(), nodes.end(),
+        [](const NodePtr& node){ return Type::Reorder == node->getType(); });
+
+    ASSERT_EQ(itr, nodes.end());
 }

src/plugins/intel_gna/tests/functional/shared_tests_instances/skip_tests_config.cpp

@@ -109,5 +109,8 @@ std::vector<std::string> disabledTestPatterns() {
         R"(.*OVCompiledModelBaseTest.*CanSetInputPrecisionForNetwork.*)",
         R"(.*OVCompiledModelBaseTest.*CanCreateTwoCompiledModelsAndCheckRuntimeModel.*)",
         R"(.*OVInferRequestWaitTests.*FailedAsyncInferWithNegativeTimeForWait.*)",
+        // TODO: Issue 127801
+        R"(.*smoke_MemoryTest.*transformation=NONE.*)",
+        R"(.*smoke_concat_memory.*)",
     };
 }

src/plugins/template/backend/int_executable.cpp

@@ -65,7 +65,7 @@ bool ov::runtime::interpreter::INTExecutable::call(std::vector<ov::Tensor>& outp
         if (auto var_extension = std::dynamic_pointer_cast<ov::op::util::VariableExtension>(op)) {
             auto variable = var_extension->get_variable();
             if (!variable_context.get_variable_value(variable)) {
-                auto h_tensor = ov::Tensor(op->get_input_element_type(0), op->get_input_shape(0));
+                auto h_tensor = ov::Tensor(op->get_output_element_type(0), op->get_output_shape(0));
                 variable_context.set_variable_value(variable, std::make_shared<ov::op::util::VariableValue>(h_tensor));
             }
         }

src/plugins/template/src/sync_infer_request.cpp

@@ -91,9 +91,8 @@ ov::template_plugin::InferRequest::InferRequest(const std::shared_ptr<const ov::
             ov::Tensor tensor = ov::Tensor(variable->get_info().data_type, shape);
             variable_context.set_variable_value(variable, std::make_shared<ov::op::util::VariableValue>(tensor));
         }
-        auto state = std::make_shared<VariableState>(
-            variable->get_info().variable_id,
-            get_tensor_impl(variable_context.get_variable_value(variable)->get_state()));
+        auto state = std::make_shared<VariableState>(variable->get_info().variable_id,
+                                                     variable_context.get_variable_value(variable));
         m_variable_states.emplace_back(state);
     }
     m_eval_context.emplace("VariableContext", variable_context);

src/plugins/template/src/variable_state.hpp

@@ -13,8 +13,10 @@ namespace template_plugin {
 
 class VariableState : public ov::IVariableState {
 public:
-    VariableState(const std::string& name, const ov::SoPtr<ov::ITensor>& tensor) : ov::IVariableState(name) {
-        m_state = tensor;
+    VariableState(const std::string& name, const std::shared_ptr<op::util::VariableValue>& variable_value)
+        : ov::IVariableState(name),
+          m_variable_value(variable_value) {
+        m_state = get_tensor_impl(variable_value->get_state());
     }
     void set_state(const ov::SoPtr<ov::ITensor>& state) override {
         OPENVINO_ASSERT(state->get_shape() == m_state->get_shape(), "Wrong tensor shape.");
@@ -25,10 +27,13 @@ public:
 
     void reset() override {
         std::memset(m_state->data(), 0, m_state->get_byte_size());
+        m_variable_value->set_reset(true);
     }
 
     ~VariableState() override = default;
-};
 
+private:
+    std::shared_ptr<op::util::VariableValue> m_variable_value;
+};
 }  // namespace template_plugin
 }  // namespace ov

src/plugins/template/tests/functional/op_reference/memory.cpp

@@ -464,65 +464,64 @@ TEST_P(ReferenceReadValueAssignAddMultiVariableLayerTest, MemoryWithHardcodedRef
 template <ov::element::Type_t IN_ET>
 std::vector<MemoryTestParams> ReadValueAssignAddMultiVariableLayer() {
     using T = typename ov::element_type_traits<IN_ET>::value_type;
+    const size_t num_tests = 3;
     size_t count_runs = 10;
 
-    std::vector<T> first_result_shape1 = {1};
-    std::vector<T> first_result_shape22 = {1, 2, 3, 4};
-    std::vector<T> first_result_shape123 = {1, 2, 3, 4, 5, 6};
+    std::vector<std::vector<T>> parameter_value(num_tests);
+    parameter_value[0] = {1};
+    parameter_value[1] = {1, 2, 3, 4};
+    parameter_value[2] = {1, 2, 3, 4, 5, 6};
 
-    std::vector<T> new_result_shape1(1, T(0));
-    std::vector<T> new_result_shape22(4, T(0));
-    std::vector<T> new_result_shape123(6, T(0));
+    std::vector<Shape> in_out_shapes = {{1}, {1, 2}, {1, 2, 3}};
 
-    std::vector<std::vector<T>> result_shape1;
-    std::vector<std::vector<T>> result_shape22;
-    std::vector<std::vector<T>> result_shape123;
+    // the initial value for the buffers is equal to the params values on the 1st iteration
+    auto state_buffer_value = parameter_value;
 
-    for (size_t i = 0; i < count_runs; i++) {
-        std::transform(new_result_shape1.begin(),
-                       new_result_shape1.end(),
-                       first_result_shape1.begin(),
-                       new_result_shape1.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape22.begin(),
-                       new_result_shape22.end(),
-                       first_result_shape22.begin(),
-                       new_result_shape22.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape123.begin(),
-                       new_result_shape123.end(),
-                       first_result_shape123.begin(),
-                       new_result_shape123.begin(),
-                       std::plus<T>());
-        result_shape1.push_back(new_result_shape1);
-        result_shape22.push_back(new_result_shape22);
-        result_shape123.push_back(new_result_shape123);
+    // the result contain values after each inference request
+    // number of inferences = count_runs
+    std::vector<std::vector<std::vector<T>>> expected_res(num_tests);
+
+    // the reference for ov::Model:
+    //   ___________
+    //  |           | -> [ReadValue 1] ->   ________
+    //                                     |  Add 1 |        _______
+    //  | Parameter | -> [ReadValue 2] ->  |________| ----> |       |
+    //                                                      |  Add 2| -> Assign_1 -> Result
+    //  | __________| ---------------------------------->   |_______|
+    //  Note: Assign_2 is not shown in the graph here, it exists and connected to ReadValue2 directly,
+    //  but we don't check its value.
+    std::vector<std::vector<T>> add_1(num_tests);
+    for (size_t i = 0; i < num_tests; ++i) {
+        add_1[i].resize(parameter_value[i].size(), 0);
     }
 
-    std::vector<MemoryTestParams> params{MemoryTestParams(ov::Shape{1},
-                                                          ov::Shape{1},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1},
-                                                          result_shape1,
-                                                          count_runs,
-                                                          {"v0", "v1"}),
-                                         MemoryTestParams(ov::Shape{2, 2},
-                                                          ov::Shape{2, 2},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1, 2, 3, 4},
-                                                          result_shape22,
-                                                          count_runs,
-                                                          {"v0", "v1"}),
-                                         MemoryTestParams(ov::Shape{1, 2, 3},
-                                                          ov::Shape{1, 2, 3},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1, 2, 3, 4, 5, 6},
-                                                          result_shape123,
-                                                          count_runs,
-                                                          {"v0", "v1"})};
+    std::vector<MemoryTestParams> params;
+    for (size_t test_i = 0; test_i < num_tests; ++test_i) {
+        for (size_t i = 0; i < count_runs; i++) {
+            // Add1 = ReadValue1 + ReadValue2
+            std::transform(state_buffer_value[test_i].begin(),
+                           state_buffer_value[test_i].end(),
+                           parameter_value[test_i].begin(),
+                           add_1[test_i].begin(),
+                           std::plus<T>());
+            // Res = Add1 + Parameter
+            std::transform(add_1[test_i].begin(),
+                           add_1[test_i].end(),
+                           parameter_value[test_i].begin(),
+                           state_buffer_value[test_i].begin(),
+                           std::plus<T>());
+
+            expected_res[test_i].push_back(state_buffer_value[test_i]);
+        }
+        params.push_back(MemoryTestParams(in_out_shapes[test_i],
+                                          in_out_shapes[test_i],
+                                          IN_ET,
+                                          IN_ET,
+                                          parameter_value[test_i],
+                                          expected_res[test_i],
+                                          count_runs,
+                                          {"v0", "v1"}));
+    }
     return params;
 }
 
@@ -568,10 +567,7 @@ TEST_P(ReferenceReadValueAssignAddResetLayerTest, MemoryResetWithHardcodedRefs)
 
     auto reset_var = [&](size_t iter, ov::InferRequest& inferRequest) {
         if (params.m_reset_on_run == iter) {
-            auto vars = inferRequest.query_state();
-            for (auto& var : vars) {
-                var.reset();
-            }
+            inferRequest.reset_state();
         }
     };
     CommonTestSteps(reset_var);
@@ -874,8 +870,7 @@ TEST_P(ReferenceReadValueAssignAddMultiVariableModifyLayerTest, MemoryResetWithH
 
     auto reset_var = [&](size_t iter, ov::InferRequest& inferRequest) {
         if (params.m_reset_on_run == iter) {
-            auto vars = inferRequest.query_state();
-            vars[1].set_state(params.m_input_data);
+            inferRequest.reset_state();
         }
     };
     CommonTestSteps(reset_var);
@@ -884,106 +879,69 @@ TEST_P(ReferenceReadValueAssignAddMultiVariableModifyLayerTest, MemoryResetWithH
 template <ov::element::Type_t IN_ET>
 std::vector<MemoryTestParams> generateParamsForReadValueAssignAddMultiVariableModify() {
     using T = typename ov::element_type_traits<IN_ET>::value_type;
+    const size_t num_tests = 3;
     size_t count_runs = 10;
     size_t reset_on_run = 5;
 
-    std::vector<T> first_result_shape1 = {1};
-    std::vector<T> first_result_shape22 = {1, 2, 3, 4};
-    std::vector<T> first_result_shape123 = {1, 2, 3, 4, 5, 6};
+    std::vector<std::vector<T>> parameter_value(num_tests);
+    parameter_value[0] = {1};
+    parameter_value[1] = {1, 2, 3, 4};
+    parameter_value[2] = {1, 2, 3, 4, 5, 6};
 
-    std::vector<T> new_result_shape1(1, T(0));
-    std::vector<T> new_result_shape22(4, T(0));
-    std::vector<T> new_result_shape123(6, T(0));
+    std::vector<Shape> in_out_shapes = {{1}, {1, 2}, {1, 2, 3}};
 
-    std::vector<std::vector<T>> result_shape1;
-    std::vector<std::vector<T>> result_shape22;
-    std::vector<std::vector<T>> result_shape123;
+    // the initial value for the buffers is equal to the params values on the 1st iteration
+    auto state_buffer_value = parameter_value;
 
-    for (size_t i = 0; i < count_runs - reset_on_run; i++) {
-        std::transform(new_result_shape1.begin(),
-                       new_result_shape1.end(),
-                       first_result_shape1.begin(),
-                       new_result_shape1.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape22.begin(),
-                       new_result_shape22.end(),
-                       first_result_shape22.begin(),
-                       new_result_shape22.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape123.begin(),
-                       new_result_shape123.end(),
-                       first_result_shape123.begin(),
-                       new_result_shape123.begin(),
-                       std::plus<T>());
-        result_shape1.push_back(new_result_shape1);
-        result_shape22.push_back(new_result_shape22);
-        result_shape123.push_back(new_result_shape123);
+    // the result contain values after each inference request
+    // number of inferences = count_runs
+    std::vector<std::vector<std::vector<T>>> expected_res(num_tests);
+
+    // the reference for ov::Model:
+    //   ___________
+    //  |           | -> [ReadValue 1] ->   ________
+    //                                     |  Add 1 |        _______
+    //  | Parameter | -> [ReadValue 2] ->  |________| ----> |       |
+    //                                                      |  Add 2| -> Assign_1 -> Result
+    //  | __________| ---------------------------------->   |_______|
+    //  Note: Assign_2 is not shown in the graph here, it exists and connected to ReadValue2 directly,
+    //  but we don't check its value.
+    std::vector<std::vector<T>> add_1(num_tests);
+    for (size_t i = 0; i < num_tests; ++i) {
+        add_1[i].resize(parameter_value[i].size(), 0);
     }
 
-    std::transform(first_result_shape1.begin(),
-                   first_result_shape1.end(),
-                   first_result_shape1.begin(),
-                   first_result_shape1.begin(),
-                   std::plus<T>());
-    std::transform(first_result_shape22.begin(),
-                   first_result_shape22.end(),
-                   first_result_shape22.begin(),
-                   first_result_shape22.begin(),
-                   std::plus<T>());
-    std::transform(first_result_shape123.begin(),
-                   first_result_shape123.end(),
-                   first_result_shape123.begin(),
-                   first_result_shape123.begin(),
-                   std::plus<T>());
+    std::vector<MemoryTestParams> params;
+    for (size_t test_i = 0; test_i < num_tests; ++test_i) {
+        for (size_t i = 0; i < count_runs; i++) {
+            if (i == reset_on_run) {
+                state_buffer_value[test_i] = parameter_value[test_i];
+            }
+            // Add1 = ReadValue1 + ReadValue2
+            std::transform(state_buffer_value[test_i].begin(),
+                           state_buffer_value[test_i].end(),
+                           parameter_value[test_i].begin(),
+                           add_1[test_i].begin(),
+                           std::plus<T>());
+            // Res = Add1 + Parameter
+            std::transform(add_1[test_i].begin(),
+                           add_1[test_i].end(),
+                           parameter_value[test_i].begin(),
+                           state_buffer_value[test_i].begin(),
+                           std::plus<T>());
 
-    for (size_t i = count_runs - reset_on_run; i < count_runs; i++) {
-        std::transform(new_result_shape1.begin(),
-                       new_result_shape1.end(),
-                       first_result_shape1.begin(),
-                       new_result_shape1.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape22.begin(),
-                       new_result_shape22.end(),
-                       first_result_shape22.begin(),
-                       new_result_shape22.begin(),
-                       std::plus<T>());
-        std::transform(new_result_shape123.begin(),
-                       new_result_shape123.end(),
-                       first_result_shape123.begin(),
-                       new_result_shape123.begin(),
-                       std::plus<T>());
-        result_shape1.push_back(new_result_shape1);
-        result_shape22.push_back(new_result_shape22);
-        result_shape123.push_back(new_result_shape123);
+            expected_res[test_i].push_back(state_buffer_value[test_i]);
+        }
+        params.push_back(MemoryTestParams(in_out_shapes[test_i],
+                                          in_out_shapes[test_i],
+                                          IN_ET,
+                                          IN_ET,
+                                          parameter_value[test_i],
+                                          expected_res[test_i],
+                                          count_runs,
+                                          {"v0", "v1"},
+                                          reset_on_run));
     }
-
-    std::vector<MemoryTestParams> params{MemoryTestParams(ov::Shape{1},
-                                                          ov::Shape{1},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1},
-                                                          result_shape1,
-                                                          count_runs,
-                                                          {"v0", "v1"},
-                                                          reset_on_run),
-                                         MemoryTestParams(ov::Shape{2, 2},
-                                                          ov::Shape{2, 2},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1, 2, 3, 4},
-                                                          result_shape22,
-                                                          count_runs,
-                                                          {"v0", "v1"},
-                                                          reset_on_run),
-                                         MemoryTestParams(ov::Shape{1, 2, 3},
-                                                          ov::Shape{1, 2, 3},
-                                                          IN_ET,
-                                                          IN_ET,
-                                                          std::vector<T>{1, 2, 3, 4, 5, 6},
-                                                          result_shape123,
-                                                          count_runs,
-                                                          {"v0", "v1"},
-                                                          reset_on_run)};
     return params;
 }