[GPU] Fix set_state() when precision conversion is needed (#21874)

Author: Vladimir Paramuzov

src/plugins/intel_gpu/include/intel_gpu/plugin/common_utils.hpp

@@ -103,6 +103,7 @@ inline void ForceExit() {
     std::_Exit(-1);
 }
 
+void convert_and_copy(const ov::ITensor* src, cldnn::memory::ptr dst, cldnn::stream& stream);
 void convert_and_copy(const cldnn::memory::ptr src, ov::ITensor const* dst, const cldnn::stream& stream);
 void convert_and_copy(const ov::ITensor* src, ov::ITensor const* dst, const cldnn::stream& stream);
 

src/plugins/intel_gpu/src/plugin/common_utils.cpp

@@ -94,6 +94,27 @@ void convert_and_copy(const void* src_ptr, ov::element::Type src_et, void* dst_p
 namespace ov {
 namespace intel_gpu {
 
+void convert_and_copy(const ov::ITensor* src, cldnn::memory::ptr dst, cldnn::stream& stream) {
+    const bool blocking = true;
+    auto src_et = src->get_element_type();
+    auto dst_et = dst->get_layout().data_type;
+
+    if (dst_et == src_et) {
+        if (auto remote = dynamic_cast<const ov::intel_gpu::RemoteTensorImpl*>(src)) {
+            auto mem = remote->get_original_memory();
+            dst->copy_from(stream, *mem, blocking);
+        } else {
+            dst->copy_from(stream, src->data(), blocking);
+            return;
+        }
+    }
+
+    size_t size = ov::shape_size(src->get_shape());
+    ov::Tensor tmp_tensor(dst_et, src->get_shape());
+    ::convert_and_copy(src->data(), src_et, tmp_tensor.data(), dst_et, size, cldnn::layout({}, ov::element::undefined, cldnn::format::bfyx, cldnn::padding()));
+    dst->copy_from(stream, tmp_tensor.data(), blocking);
+}
+
 void convert_and_copy(const cldnn::memory::ptr src, ov::ITensor const* dst, const cldnn::stream& stream) {
     auto src_et = src->get_layout().data_type;
     auto dst_et = dst->get_element_type();

src/plugins/intel_gpu/src/plugin/variable_state.cpp

@@ -54,17 +54,9 @@ void VariableState::set_layout(const cldnn::layout& new_layout) {
 }
 
 void VariableState::set_state(const ov::SoPtr<ov::ITensor>& state) {
-    const bool blocking = true;
-    auto remote_ptr = std::dynamic_pointer_cast<RemoteTensorImpl>(state._ptr);
     m_layout.set_partial_shape(state->get_shape());
     update_device_buffer();
-    if (remote_ptr != nullptr) {
-        auto user_memory = remote_ptr->get_memory();
-        m_memory->copy_from(m_context->get_engine().get_service_stream(), *user_memory, blocking);
-    } else {
-        auto data = state->data();
-        m_memory->copy_from(m_context->get_engine().get_service_stream(), data, blocking);
-    }
+    convert_and_copy(state._ptr.get(), m_memory, m_context->get_engine().get_service_stream());
     set();
 }
 

src/plugins/intel_gpu/tests/functional/behavior/infer_request.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "common_test_utils/ov_tensor_utils.hpp"
 #include "common_test_utils/test_common.hpp"
 #include "common_test_utils/common_utils.hpp"
 #include "common_test_utils/node_builders/activation.hpp"
@@ -238,4 +239,39 @@ TEST(VariablesTest, smoke_canSetStateTensor) {
 
     ASSERT_NO_THROW(request.infer());
 }
-} // namespace
+
+TEST(VariablesTest, smoke_set_get_state_with_convert) {
+    auto build_model = [](ov::element::Type type, const ov::PartialShape& shape) {
+        auto param = std::make_shared<ov::op::v0::Parameter>(type, shape);
+        const ov::op::util::VariableInfo variable_info { shape, type, "v0" };
+        auto variable = std::make_shared<ov::op::util::Variable>(variable_info);
+        auto read_value = std::make_shared<ov::op::v6::ReadValue>(param, variable);
+        auto add = std::make_shared<ov::op::v1::Add>(read_value, param);
+        auto assign = std::make_shared<ov::op::v6::Assign>(add, variable);
+        auto res = std::make_shared<ov::op::v0::Result>(add);
+        return std::make_shared<ov::Model>(ov::ResultVector { res }, ov::SinkVector { assign }, ov::ParameterVector{param}, "StateTestModel");
+    };
+
+    auto ov = ov::Core();
+    const ov::Shape virable_shape = {1, 3, 2, 4};
+    const ov::Shape input_shape = {1, 3, 2, 4};
+    const ov::element::Type et = ov::element::f32;
+    auto model = build_model(et, input_shape);
+    auto compiled_model = ov.compile_model(model, ov::test::utils::DEVICE_GPU, ov::hint::inference_precision(ov::element::f16));
+    auto request = compiled_model.create_infer_request();
+
+    auto variables = request.query_state();
+    ASSERT_EQ(variables.size(), 1);
+    auto variable = variables.front();
+    ASSERT_EQ(variable.get_name(), "v0");
+    auto state_tensor = variable.get_state();
+    ASSERT_EQ(state_tensor.get_shape(), virable_shape);
+    ASSERT_EQ(state_tensor.get_element_type(), et);
+
+    auto tensor_to_set = ov::test::utils::create_and_fill_tensor(et, state_tensor.get_shape());
+    variable.set_state(tensor_to_set);
+    state_tensor = variable.get_state();
+
+    ov::test::utils::compare(tensor_to_set, state_tensor, 1e-5f, 1e-5f);
+}
+} // namespace