Loop/If/TensorIterator - fix dynamic input cases (#9752)

* Loop/If/TensorIterator - fix dynamic input cases
Reference evaluate for body uses Model::evaluate instead of custom evaluation
Loop/TensorIterator additional fix - set result shape according to body execution result

Only op_eval test verifies issues, template tests were added just in case (these passed even without fix)

* Fix clang-format

* rename ti.cpp

docs/template_plugin/tests/functional/op_reference/if.cpp

@@ -224,6 +224,40 @@ struct IfConditionIsDynamic : public IfFunctionalBase {
     }
 };
 
+struct IfDynamicInputs : public IfFunctionalBase {
+    std::shared_ptr<Model> create_function(const std::vector<reference_tests::Tensor>& if_inputs,
+                                           const std::vector<reference_tests::Tensor>& results) override {
+        NGRAPH_CHECK(if_inputs.size() == 3, "Incorrect test case! Number of inputs is not 3.");
+        NGRAPH_CHECK(results.size() == 1, "Incorrect test case! Number of outputs is not 1.");
+
+        auto X = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        auto Y = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape{Dimension::dynamic()});
+        // Set up the cell body, a function from (Xi, Yi) -> (Zo)
+        // Body parameters
+        auto Xt = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        auto Yt = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        auto Xe = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        auto Ye = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+        // Body
+        auto then_op = std::make_shared<op::v1::Multiply>(Xt, Yt);
+        auto else_op = std::make_shared<op::v1::Add>(Xe, Ye);
+        auto then_op_result = std::make_shared<op::v0::Result>(then_op);
+        auto else_op_result = std::make_shared<op::v0::Result>(else_op);
+        auto then_body = std::make_shared<ov::Model>(OutputVector{then_op_result}, ParameterVector{Xt, Yt});
+        auto else_body = std::make_shared<ov::Model>(OutputVector{else_op_result}, ParameterVector{Xe, Ye});
+        auto if_op = std::make_shared<op::v8::If>(cond);
+        if_op->set_then_body(then_body);
+        if_op->set_else_body(else_body);
+        if_op->set_input(X, Xt, Xe);
+        if_op->set_input(Y, Yt, Ye);
+        auto rs = if_op->set_output(then_op_result, else_op_result);
+        auto result = std::make_shared<op::v0::Result>(rs);
+        auto fun = std::make_shared<Model>(OutputVector{result}, ParameterVector{cond, X, Y});
+        return fun;
+    }
+};
+
 struct IfParams {
     IfParams(const std::shared_ptr<IfFunctionalBase>& functional,
              const std::vector<reference_tests::Tensor>& if_inputs,
@@ -365,4 +399,18 @@ INSTANTIATE_TEST_SUITE_P(
                                 reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{1.0, 2.0, 3.0, 4.0}),
                                 reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{2.0, 1.0, 2.0, 3.0})},
             std::vector<reference_tests::Tensor>{reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{3.0, 3.0, 5.0, 7.0})},
-            "if_condition_is_dynamic_cond_false")));
+            "if_condition_is_dynamic_cond_false"),
+        IfParams(
+            std::make_shared<IfDynamicInputs>(),
+            std::vector<reference_tests::Tensor>{reference_tests::Tensor(Shape{}, ngraph::element::boolean, std::vector<unsigned char>{1}),
+                                reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{1.0, 2.0, 3.0, 4.0}),
+                                reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{2.0, 1.0, 2.0, 3.0})},
+            std::vector<reference_tests::Tensor>{reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{2.0, 2.0, 6.0, 12.0})},
+            "if_dynamic_inputs_cond_true"),
+        IfParams(
+            std::make_shared<IfDynamicInputs>(),
+            std::vector<reference_tests::Tensor>{reference_tests::Tensor(Shape{}, ngraph::element::boolean, std::vector<unsigned char>{0}),
+                                reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{1.0, 2.0, 3.0, 4.0}),
+                                reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{2.0, 1.0, 2.0, 3.0})},
+            std::vector<reference_tests::Tensor>{reference_tests::Tensor(Shape{1, 2, 2}, ngraph::element::f32, std::vector<float>{3.0, 3.0, 5.0, 7.0})},
+            "if_dynamic_inputs_cond_false")));

docs/template_plugin/tests/functional/op_reference/loop.cpp

@@ -0,0 +1,115 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <openvino/core/model.hpp>
+#include <openvino/opsets/opset8.hpp>
+
+#include "base_reference_test.hpp"
+#include "functional_test_utils/skip_tests_config.hpp"
+
+struct LoopFunctionalBase {
+    virtual std::shared_ptr<ov::Model> create_function(const std::vector<reference_tests::Tensor>& loop_inputs,
+                                                       const std::vector<reference_tests::Tensor>& results) = 0;
+    LoopFunctionalBase() = default;
+    virtual ~LoopFunctionalBase() = default;
+};
+
+struct LoopDynamicInputs : public LoopFunctionalBase {
+    std::shared_ptr<ov::Model> create_function(const std::vector<reference_tests::Tensor>& loop_inputs,
+                                               const std::vector<reference_tests::Tensor>& results) override {
+        auto X = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto Y = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto M = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+
+        // Set up the cell body, a function from (Xi, Yi) -> (Zo)
+        // Body parameters
+        auto Xi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto Yi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto M_body = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto body_condition = std::make_shared<ov::opset8::Constant>(ov::element::boolean, ov::Shape{1}, true);
+
+        auto trip_count = std::make_shared<ov::opset8::Constant>(ngraph::element::i64, ov::Shape{1}, 3);
+        auto exec_condition = std::make_shared<ov::opset8::Constant>(ngraph::element::boolean, ov::Shape{1}, true);
+        // Body
+        auto sum = std::make_shared<ov::opset8::Add>(Xi, Yi);
+        auto Zo = std::make_shared<ov::opset8::Multiply>(sum, M_body);
+        auto body = std::make_shared<ov::Model>(ov::OutputVector{body_condition, Zo},
+                                                ov::ParameterVector{Xi, Yi, M_body});
+
+        auto loop = std::make_shared<ov::opset8::Loop>(trip_count, exec_condition);
+        loop->set_function(body);
+
+        loop->set_invariant_input(Xi, X);
+        loop->set_invariant_input(Yi, Y);
+        loop->set_merged_input(M_body, M, Zo);
+
+        loop->set_special_body_ports(ov::opset8::Loop::SpecialBodyPorts{-1, 0});
+
+        // Output is last Zo
+        auto result = std::make_shared<ov::opset8::Result>(loop->get_iter_value(Zo, -1));
+        return std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{X, Y, M});
+    }
+};
+
+struct LoopParams {
+    LoopParams(const std::shared_ptr<LoopFunctionalBase>& functional,
+               const std::vector<reference_tests::Tensor>& loop_inputs,
+               const std::vector<reference_tests::Tensor>& expected_results,
+               const std::string& test_case_name)
+            : function(functional),
+              inputs(loop_inputs),
+              expected_results(expected_results),
+              test_case_name(test_case_name) {}
+
+    std::shared_ptr<LoopFunctionalBase> function;
+    std::vector<reference_tests::Tensor> inputs;
+    std::vector<reference_tests::Tensor> expected_results;
+    std::string test_case_name;
+};
+
+class ReferenceLoopLayerTest : public testing::TestWithParam<LoopParams>, public reference_tests::CommonReferenceTest {
+public:
+    void SetUp() override {
+        SKIP_IF_CURRENT_TEST_IS_DISABLED()
+        auto params = GetParam();
+        function = params.function->create_function(params.inputs, params.expected_results);
+        inputData.reserve(params.inputs.size());
+        refOutData.reserve(params.expected_results.size());
+        for (auto& input_tensor : params.inputs) {
+            inputData.push_back(input_tensor.data);
+        }
+        for (auto& expected_tensor : params.expected_results) {
+            refOutData.push_back(expected_tensor.data);
+        }
+    }
+    static std::string getTestCaseName(const testing::TestParamInfo<LoopParams>& obj) {
+        auto param = obj.param;
+        return param.test_case_name;
+    }
+};
+
+TEST_P(ReferenceLoopLayerTest, TensorIteratorWithHardcodedRefs) {
+    Exec();
+}
+
+INSTANTIATE_TEST_SUITE_P(
+        smoke_TensorIterator_With_Hardcoded_Refs,
+        ReferenceLoopLayerTest,
+        ::testing::Values(
+                LoopParams(
+                        std::make_shared<LoopDynamicInputs>(),
+                        std::vector<reference_tests::Tensor>{
+                                reference_tests::Tensor(ov::element::f32, ov::Shape{2, 2}, std::vector<float>{0, 1, 2, 3}),
+                                reference_tests::Tensor(ov::element::f32, ov::Shape{2, 2}, std::vector<float>{1, 2, 3, 4}),
+                                reference_tests::Tensor(ov::element::f32, ov::Shape{2, 2}, std::vector<float>{5, 4, 3, 2})},
+                        // 5*(0+1)*(0+1)*(0+1) = 5
+                        // 4*(1+2)*(1+2)*(1+2) = 108
+                        // 3*(2+3)*(2+3)*(2+3) = 375
+                        // 2*(3+4)*(3+4)*(3+4) = 686
+                        std::vector<reference_tests::Tensor>{
+                                reference_tests::Tensor(ov::element::f32, ov::Shape{2, 2}, std::vector<float>{5, 108, 375, 686})},
+                        "loop_dynamic_inputs")),
+        ReferenceLoopLayerTest::getTestCaseName);

docs/template_plugin/tests/functional/op_reference/tensor_iterator.cpp

@@ -0,0 +1,110 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <openvino/core/model.hpp>
+#include <openvino/opsets/opset8.hpp>
+
+#include "base_reference_test.hpp"
+#include "functional_test_utils/skip_tests_config.hpp"
+
+struct TIFunctionalBase {
+    virtual std::shared_ptr<ov::Model> create_function(const std::vector<reference_tests::Tensor>& ti_inputs,
+                                                   const std::vector<reference_tests::Tensor>& results) = 0;
+    TIFunctionalBase() = default;
+    virtual ~TIFunctionalBase() = default;
+};
+
+struct TIDynamicInputs : public TIFunctionalBase {
+    std::shared_ptr<ov::Model> create_function(const std::vector<reference_tests::Tensor>& ti_inputs,
+                                           const std::vector<reference_tests::Tensor>& results) override {
+        auto X = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto Y = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto M = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+
+        // Set up the cell body, a function from (Xi, Yi) -> (Zo)
+        // Body parameters
+        auto Xi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto Yi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto M_body = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+        auto body_condition = std::make_shared<ov::opset8::Constant>(ov::element::boolean, ov::Shape{1}, true);
+
+        auto trip_count = std::make_shared<ov::opset8::Constant>(ngraph::element::i64, ov::Shape{1}, 3);
+        auto exec_condition = std::make_shared<ov::opset8::Constant>(ngraph::element::boolean, ov::Shape{1}, true);
+        // Body
+        auto sum = std::make_shared<ov::opset8::Add>(Xi, Yi);
+        auto Zo = std::make_shared<ov::opset8::Multiply>(sum, M_body);
+        auto body = std::make_shared<ov::Model>(ov::OutputVector{body_condition, Zo},
+                                                ov::ParameterVector{Xi, Yi, M_body});
+
+        auto tensor_iterator = std::make_shared<ov::opset8::TensorIterator>();
+        tensor_iterator->set_function(body);
+
+        tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 1);
+        tensor_iterator->set_sliced_input(Yi, Y, 0, 1, 1, -1, 0);
+        tensor_iterator->set_merged_input(M_body, M, Zo);
+
+        // Output 0 is last Zo
+        auto out1 = tensor_iterator->get_iter_value(Zo, -1);
+        return std::make_shared<ov::Model>(ov::OutputVector{out1}, ov::ParameterVector{X, Y, M});
+    }
+};
+
+struct TensorIteratorParams {
+    TensorIteratorParams(const std::shared_ptr<TIFunctionalBase>& functional,
+             const std::vector<reference_tests::Tensor>& ti_inputs,
+             const std::vector<reference_tests::Tensor>& expected_results,
+             const std::string& test_case_name)
+            : function(functional),
+              inputs(ti_inputs),
+              expected_results(expected_results),
+              test_case_name(test_case_name) {}
+
+    std::shared_ptr<TIFunctionalBase> function;
+    std::vector<reference_tests::Tensor> inputs;
+    std::vector<reference_tests::Tensor> expected_results;
+    std::string test_case_name;
+};
+
+class ReferenceTILayerTest : public testing::TestWithParam<TensorIteratorParams>,
+        public reference_tests::CommonReferenceTest {
+public:
+    void SetUp() override {
+        SKIP_IF_CURRENT_TEST_IS_DISABLED()
+        auto params = GetParam();
+        function = params.function->create_function(params.inputs, params.expected_results);
+        inputData.reserve(params.inputs.size());
+        refOutData.reserve(params.expected_results.size());
+        for (auto& input_tensor : params.inputs) {
+            inputData.push_back(input_tensor.data);
+        }
+        for (auto& expected_tensor : params.expected_results) {
+            refOutData.push_back(expected_tensor.data);
+        }
+    }
+    static std::string getTestCaseName(const testing::TestParamInfo<TensorIteratorParams>& obj) {
+        auto param = obj.param;
+        return param.test_case_name;
+    }
+};
+
+TEST_P(ReferenceTILayerTest, TensorIteratorWithHardcodedRefs) {
+    Exec();
+}
+
+INSTANTIATE_TEST_SUITE_P(
+        smoke_TensorIterator_With_Hardcoded_Refs,
+        ReferenceTILayerTest,
+        ::testing::Values(
+                TensorIteratorParams(
+                        std::make_shared<TIDynamicInputs>(),
+                        std::vector<reference_tests::Tensor>{
+                            reference_tests::Tensor(ov::element::f32, ov::Shape{1, 2}, std::vector<float>{2, 3}),
+                            reference_tests::Tensor(ov::element::f32, ov::Shape{2, 1}, std::vector<float>{4, 5}),
+                            reference_tests::Tensor(ov::element::f32, ov::Shape{1, 1}, std::vector<float>{5})},
+                        std::vector<reference_tests::Tensor>{
+                            reference_tests::Tensor(ov::element::f32, ov::Shape{1, 1}, std::vector<float>{240})},
+                        "tensor_iterator_dynamic_inputs")),
+        ReferenceTILayerTest::getTestCaseName);

src/core/reference/src/runtime/reference/function.cpp

@@ -15,106 +15,17 @@
 namespace ngraph {
 namespace runtime {
 namespace reference {
-static bool call(const HostTensorVector& func_outputs,
-                 const HostTensorVector& func_inputs,
-                 const std::shared_ptr<ngraph::Function>& function) {
-    // map function params -> HostTensor
-    std::unordered_map<descriptor::Tensor*, std::shared_ptr<HostTensor>> tensor_map;
-    size_t input_count = 0;
-    for (const auto& param : function->get_parameters()) {
-        for (size_t i = 0; i < param->get_output_size(); ++i) {
-            descriptor::Tensor* tensor = &param->output(i).get_tensor();
-            tensor_map.insert({tensor, func_inputs[input_count++]});
-        }
-    }
-
-    std::unordered_map<std::shared_ptr<ngraph::Node>, size_t> results_map;
-    // map function outputs -> HostTensor
-    for (size_t output_count = 0; output_count < function->get_results().size(); ++output_count) {
-        auto output = function->get_results()[output_count];
-        results_map[output] = output_count;
-    }
-
-    // for each ordered op in the graph
-    for (const auto& op : function->get_ordered_ops()) {
-        if (op::is_parameter(op)) {
-            continue;
-        }
-
-        // get op inputs from map
-        std::vector<std::shared_ptr<HostTensor>> op_inputs;
-        for (auto input : op->inputs()) {
-            descriptor::Tensor* tensor = &input.get_tensor();
-            op_inputs.push_back(tensor_map.at(tensor));
-        }
-
-        // get op outputs from map or create
-        std::vector<std::shared_ptr<HostTensor>> op_outputs;
-        for (size_t i = 0; i < op->get_output_size(); ++i) {
-            descriptor::Tensor* tensor = &op->output(i).get_tensor();
-            std::shared_ptr<HostTensor> host_tensor;
-            auto it = tensor_map.find(tensor);
-            if (op::is_output(op)) {
-                host_tensor = func_outputs[results_map[op]];
-            } else if (it == tensor_map.end()) {
-                host_tensor = std::make_shared<HostTensor>(op->output(i));
-                tensor_map.insert({tensor, host_tensor});
-            } else {
-                host_tensor = it->second;
-            }
-            op_outputs.push_back(host_tensor);
-        }
-        op->validate_and_infer_types();
-        OPENVINO_SUPPRESS_DEPRECATED_START
-        if (!op->evaluate(op_outputs, op_inputs)) {
-            throw ngraph_error("Evaluate function is not implemented.");
-        }
-        OPENVINO_SUPPRESS_DEPRECATED_END
-    }
-    return true;
-}
-
 void function(const std::shared_ptr<ngraph::Function>& function,
               const HostTensorVector& inputs,
               HostTensorVector& outputs) {
-    const auto& parameters = function->get_parameters();
-    const auto& parametersNumber = parameters.size();
-    const auto& inputsNumber = inputs.size();
-    NGRAPH_CHECK(parametersNumber == inputsNumber,
-                 "Got function (",
-                 function->get_friendly_name(),
-                 ") with ",
-                 parametersNumber,
-                 " parameters, but ",
-                 inputsNumber,
-                 " input blobs");
-
-    for (const auto& parameter : parameters) {
-        const auto& parameterIndex = function->get_parameter_index(parameter);
-        const auto& parameterShape = parameter->get_shape();
-        const auto& parameterType = parameter->get_element_type();
-        const auto& parameterSize = shape_size(parameterShape) * parameterType.size();
-
-        const auto& input = inputs[parameterIndex];
-        const auto& inputSize = input->get_size_in_bytes();
-        NGRAPH_CHECK(parameterSize == inputSize,
-                     "Got parameter (",
-                     parameter->get_friendly_name(),
-                     ") of size ",
-                     parameterSize,
-                     " bytes, but corresponding input with index ",
-                     parameterIndex,
-                     " has ",
-                     inputSize,
-                     " bytes");
-    }
-
+    OPENVINO_SUPPRESS_DEPRECATED_START
     const auto& results = function->get_results();
     outputs.reserve(results.size());
     for (size_t i = 0; i < results.size(); ++i) {
         outputs.push_back(std::make_shared<HostTensor>());
     }
-    call(outputs, inputs, function);
+    function->evaluate(outputs, inputs);
+    OPENVINO_SUPPRESS_DEPRECATED_END
 }
 }  // namespace reference
 }  // namespace runtime

src/core/reference/src/runtime/reference/loop.cpp

@@ -171,8 +171,9 @@ void loop(const std::shared_ptr<Function>& func,
 
         for (const auto& desc : out_descs) {
             if (const auto& body_desc = std::dynamic_pointer_cast<opset5::Loop::BodyOutputDescription>(desc)) {
-                out[body_desc->m_output_index]->write(body_outputs[body_desc->m_body_value_index]->get_data_ptr(),
-                                                      body_outputs[body_desc->m_body_value_index]->get_size_in_bytes());
+                const auto& res = body_outputs[body_desc->m_body_value_index];
+                out[body_desc->m_output_index]->set_shape(res->get_shape());
+                out[body_desc->m_output_index]->write(res->get_data_ptr(), res->get_size_in_bytes());
             }
         }
 

src/core/reference/src/runtime/reference/tensor_iterator.cpp

@@ -106,8 +106,9 @@ void tensor_iterator(uint64_t num_iterations,
     for (const auto& desc : out_descs) {
         if (const auto& body_desc = std::dynamic_pointer_cast<opset5::TensorIterator::BodyOutputDescription>(desc)) {
             // Copy output values from the last iteration
-            out[body_desc->m_output_index]->write(body_outputs[body_desc->m_body_value_index]->get_data_ptr(),
-                                                  body_outputs[body_desc->m_body_value_index]->get_size_in_bytes());
+            const auto& res = body_outputs[body_desc->m_body_value_index];
+            out[body_desc->m_output_index]->set_shape(res->get_shape());
+            out[body_desc->m_output_index]->write(res->get_data_ptr(), res->get_size_in_bytes());
         }
     }
 

src/core/tests/CMakeLists.txt

@@ -476,6 +476,7 @@ set(OP_EVAL_TEST_SRC
     op_eval/hsigmoid.cpp
     op_eval/hswish.cpp
     op_eval/interpolate.cpp
+    op_eval/loop.cpp
     op_eval/matmul.cpp
     op_eval/memory.cpp
     op_eval/mish.cpp

src/core/tests/op_eval/loop.cpp

@@ -0,0 +1,59 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "engines_util/execute_tools.hpp"
+#include "gtest/gtest.h"
+#include "openvino/opsets/opset8.hpp"
+#include "util/all_close_f.hpp"
+
+OPENVINO_SUPPRESS_DEPRECATED_START
+
+TEST(op_eval, loop_dynamic_shapes) {
+    // That which we iterate over
+    auto X = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+    auto Y = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+    auto M = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+
+    // Set up the cell body, a function from (Xi, Yi) -> (Zo)
+    // Body parameters
+    auto Xi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+    auto Yi = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+    auto M_body = std::make_shared<ov::opset8::Parameter>(ov::element::f32, ov::PartialShape::dynamic());
+    auto body_condition = std::make_shared<ov::opset8::Constant>(ngraph::element::boolean, ngraph::Shape{1}, true);
+
+    auto trip_count = std::make_shared<ov::opset8::Constant>(ngraph::element::i64, ngraph::Shape{1}, 3);
+    auto exec_condition = std::make_shared<ov::opset8::Constant>(ngraph::element::boolean, ngraph::Shape{1}, true);
+    // Body
+    auto sum = std::make_shared<ov::opset8::Add>(Xi, Yi);
+    auto Zo = std::make_shared<ov::opset8::Multiply>(sum, M_body);
+    auto body = std::make_shared<ov::Model>(ov::OutputVector{body_condition, Zo}, ov::ParameterVector{Xi, Yi, M_body});
+
+    auto loop = std::make_shared<ov::opset8::Loop>(trip_count, exec_condition);
+    loop->set_function(body);
+
+    loop->set_invariant_input(Xi, X);
+    loop->set_invariant_input(Yi, Y);
+    loop->set_merged_input(M_body, M, Zo);
+
+    loop->set_special_body_ports(ov::opset8::Loop::SpecialBodyPorts{-1, 0});
+
+    // Output is last Zo
+    auto result = std::make_shared<ov::opset8::Result>(loop->get_iter_value(Zo, -1));
+    auto f = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{X, Y, M});
+
+    std::vector<float> inputX{0, 1, 2, 3}, inputY{1, 2, 3, 4}, inputM{5, 4, 3, 2};
+    std::vector<float> expected_result{5, 108, 375, 686};
+    std::vector<float> actual_result(ov::shape_size(ov::Shape{2, 2}), 2);
+
+    auto r0 = std::make_shared<ov::HostTensor>();
+    using namespace ngraph;
+    ASSERT_TRUE(f->evaluate({r0},
+                            {make_host_tensor<ngraph::element::Type_t::f32>(ov::Shape{2, 2}, inputX),
+                             make_host_tensor<ngraph::element::Type_t::f32>(ov::Shape{2, 2}, inputY),
+                             make_host_tensor<ngraph::element::Type_t::f32>(ov::Shape{2, 2}, inputM)}));
+
+    EXPECT_EQ(r0->get_shape(), (ov::Shape{2, 2}));
+    memcpy(actual_result.data(), r0->get_data_ptr<float>(), ov::shape_size(ov::Shape{2, 2}) * sizeof(float));
+    EXPECT_TRUE(ngraph::test::all_close_f(expected_result, actual_result));
+}