[CPU] Fixed dynamic Loop with explicit Concat Result (#14041)

src/plugins/intel_cpu/src/nodes/tensoriterator.cpp

@@ -482,6 +482,17 @@ bool TensorIterator::needPrepareParams() const {
             return true;
     }
 
+    // If sliced input shapes of node and body input shapes aren't equal, we should reshape body
+    if (checkForInputAndBodyShapesInequality()) {
+        return true;
+    }
+
+    // In cases, when sliced input shapes of node and body input shapes are equal,
+    // original input shapes of node may be not equal to previous input shapes and count of iterations will be another
+    // For example, TensorIterator with single sliced input by axis 1:
+    //    Input shape of node: [10, 8, 10]  ->  Sliced input shape: [10, 1, 10]  ->  Body input shape:  [10, 1, 10]  -> 8 iterations
+    //    Input shape of node: [10, 4, 10]  ->  Sliced input shape: [10, 1, 10]  ->  Body input shape:  [10, 1, 10]  -> 4 iterations
+    // Thus, sliced input shapes and body input shapes are equal but iteration counts are different. So we should update trip count
     return Node::needPrepareParams();
 }
 
@@ -667,16 +678,22 @@ void TensorIterator::prepareTripCount() {
 
 /* *==============* *==============* *==============* *==============* *==============* */
 
+inline SizeVector sliced_input_dims(const MemoryPtr& mem, const int axis, const int stride) {
+    auto dims = mem->getStaticDims();
+    if (axis != -1)
+        dims[axis] = abs(stride);
+    return dims;
+}
+
 void TensorIterator::reshapeSubgraphInput() {
     for (auto map_rule : inputPortMap) {
-        auto &from_mem = getParentEdgesAtPort(map_rule.from)[0]->getMemoryPtr();
+        auto new_dims = sliced_input_dims(getParentEdgesAtPort(map_rule.from)[0]->getMemoryPtr(), map_rule.axis, map_rule.stride);
         auto &to_mems = input_mems[map_rule.to];
-        auto new_dims = from_mem->getStaticDims();
-        if (map_rule.axis != -1)
-            new_dims[map_rule.axis] = abs(map_rule.stride);
-
-        const auto desc = std::make_shared<CpuBlockedMemoryDesc>(to_mems.front()->getDesc().getPrecision(), Shape(new_dims));
-        redefineToMemories(to_mems, desc);
+        const auto& body_inshape = to_mems.front()->GetShape();
+        if (body_inshape.isDynamic() || body_inshape.getDims() != new_dims) {
+            const auto desc = std::make_shared<CpuBlockedMemoryDesc>(to_mems.front()->getDesc().getPrecision(), Shape(new_dims));
+            redefineToMemories(to_mems, desc);
+        }
     }
 }
 
@@ -708,6 +725,19 @@ void TensorIterator::reshapeAndFillOutput(dnnl::stream strm) {
     }
 }
 
+bool TensorIterator::checkForInputAndBodyShapesInequality() const {
+    for (auto map_rule : inputPortMap) {
+        auto original_dims = sliced_input_dims(getParentEdgesAtPort(map_rule.from)[0]->getMemoryPtr(), map_rule.axis, map_rule.stride);
+        auto &to_mems = input_mems[map_rule.to];
+        const auto& body_inshape = to_mems.front()->GetShape();
+        if (body_inshape.isDynamic() || body_inshape.getDims() != original_dims) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 int TensorIterator::getNumIteration(const std::vector<PortMap>& inputPortMap, const std::vector<PortMap>& outputPortMap) const {
     const auto isIterable = [](const PortMap& rule) {
         return rule.axis != -1;

src/plugins/intel_cpu/src/nodes/tensoriterator.h

@@ -131,6 +131,7 @@ private:
     /* Dynamic support */
     void reshapeSubgraphInput();
     void reshapeAndFillOutput(dnnl::stream strm);
+    bool checkForInputAndBodyShapesInequality() const;
     int getNumIteration(const std::vector<PortMap>& inputPortMap, const std::vector<PortMap>& outputPortMap) const;
 
     ExtensionManager::Ptr ext_mng;

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

@@ -153,6 +153,8 @@ std::vector<std::string> disabledTestPatterns() {
         R"(.*LoopLayerCPUTest.*trip_count=0.*)",
         R"(.*LoopForDiffShapesLayerCPUTest.*exec_cond=0.*)",
         R"(.*LoopForDiffShapesLayerCPUTest.*trip_count=0.*)",
+        R"(.*LoopForConcatLayerCPUTest.*exec_cond=0.*)",
+        R"(.*LoopForConcatLayerCPUTest.*trip_count=0.*)",
         // [ INFO ] Can't compile network without cache for ..  with precision ..
         R"(.*CompileModelCacheTestBase.*CompareWithRefImpl.*KSOFunction.*)",
         R"(.*CompileModelCacheTestBase.*CompareWithRefImpl.*NonMaxSuppression.*)",

src/plugins/intel_cpu/tests/functional/single_layer_tests/loop.cpp

@@ -299,6 +299,68 @@ protected:
     }
 };
 
+class LoopForConcatLayerCPUTest : public LoopLayerCPUTest {
+    // for 10:
+    //   x = y + 10;
+    //   y = concat(y, x)
+    //   return y
+
+protected:
+    void SetUp() override {
+        InputLayerType trip_count_type;
+        int64_t trip_count;
+        bool exec_cond;
+        std::vector<InputShape> shapes;
+        std::vector<LOOP_IN_TYPE> types;
+        std::tie(trip_count_type, trip_count, exec_cond, shapes, types, inType) = this->GetParam();
+
+        targetDevice = CommonTestUtils::DEVICE_CPU;
+        init_input_shapes(shapes);
+
+        auto params = ngraph::builder::makeDynamicParams(inType, inputDynamicShapes);
+
+        // Body parameters
+        const std::vector<ngraph::PartialShape> body_params_shapes(shapes.size(), ngraph::PartialShape::dynamic());
+        auto body_params = ngraph::builder::makeDynamicParams(inType, inputDynamicShapes);
+
+        auto body_condition_const = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{1}, true);
+        auto exec_condition = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{1}, exec_cond);
+        std::shared_ptr<ngraph::Node> trip_count_input;
+        int shift = 0;
+        if (trip_count_type == InputLayerType::PARAMETER) {
+            for (auto& target : targetStaticShapes)
+                target.insert(target.begin(), ngraph::Shape{});
+            trip_count_input = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, ngraph::Shape{1});
+            trip_count_input->set_friendly_name("trip_count");
+            params.insert(params.begin(), ov::as_type_ptr<ngraph::opset5::Parameter>(trip_count_input));
+            shift++;
+        } else {
+            trip_count_input = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, ngraph::Shape{1}, trip_count);
+        }
+
+        // Body
+        auto constant = ngraph::builder::makeConstant(inType, std::vector<size_t>{1}, std::vector<float>{10});
+        auto add = std::make_shared<ngraph::opset5::Add>(body_params[0], constant);
+        auto concat = ngraph::builder::makeConcat({body_params[1], add}, 0);
+
+        auto body = std::make_shared<ov::Model>(ngraph::OutputVector{body_condition_const, concat}, body_params);
+
+        auto loop = std::make_shared<ngraph::opset5::Loop>(trip_count_input, exec_condition);
+        loop->set_function(body);
+        loop->set_special_body_ports(ngraph::opset5::Loop::SpecialBodyPorts{-1, 0});
+
+        loop->set_invariant_input(body_params[0], params[shift]);
+        loop->set_merged_input(body_params[1], params[shift + 1], concat);
+
+        auto out0 = loop->get_iter_value(body_condition_const, -1);
+        auto out1 = loop->get_iter_value(concat, -1);
+
+        auto result0 = std::make_shared<ngraph::opset5::Result>(out0);
+        auto result1 = std::make_shared<ngraph::opset5::Result>(out1);
+        function = std::make_shared<ov::Model>(ngraph::ResultVector{result0, result1}, params, "loop");
+    }
+};
+
 TEST_P(LoopLayerCPUTest, CompareWithRefs) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED()
 
@@ -317,6 +379,12 @@ TEST_P(LoopForDiffShapesLayerCPUTest, CompareWithRefs) {
     run();
 }
 
+TEST_P(LoopForConcatLayerCPUTest, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED()
+
+    run();
+}
+
 namespace {
 
 const std::vector<ElementType> inputPrecisions = {
@@ -477,5 +545,58 @@ INSTANTIATE_TEST_SUITE_P(smoke_LoopForDiffShapesConcat, LoopForDiffShapesLayerCP
                                  ::testing::ValuesIn(inputPrecisions)),
                          LoopLayerCPUTest::getTestCaseName);
 
+std::vector<std::vector<InputShape>> inputs_4 = {
+        {  // first test suit
+            {  // first input
+                {-1, 10, 10},
+                { // target static shapes
+                    {10, 10, 10},
+                    {5, 10, 10},
+                    {5, 10, 10},
+                    {8, 10, 10},
+                }
+            },
+            {  // second input
+                {-1, 10, 10},
+                { // target static shapes
+                    {0, 10, 10},
+                    {0, 10, 10},
+                    {0, 10, 10},
+                    {0, 10, 10},
+                }
+            },
+        },
+        {  // second test suit
+            {  // first input
+                {{0, 10}, 10, 10},
+                { // target static shapes
+                    {10, 10, 10},
+                    {5, 10, 10},
+                    {5, 10, 10},
+                    {8, 10, 10},
+                }
+            },
+            {  // second input
+                {-1, 10, 10},
+                { // target static shapes
+                    {0, 10, 10},
+                    {0, 10, 10},
+                    {0, 10, 10},
+                    {0, 10, 10},
+                }
+            },
+        },
+};
+
+INSTANTIATE_TEST_SUITE_P(smoke_LoopForConcat, LoopForConcatLayerCPUTest,
+                         ::testing::Combine(
+                                 ::testing::ValuesIn(trip_count_type),
+                                 ::testing::ValuesIn(trip_count),
+                                 ::testing::ValuesIn(exec_cond),
+                                 ::testing::ValuesIn(inputs_4),
+                                 ::testing::Values(std::vector<LOOP_IN_TYPE>{}),
+                                 ::testing::ValuesIn(inputPrecisions)),
+                         LoopLayerCPUTest::getTestCaseName);
+
 }  // namespace
 } // namespace CPULayerTestsDefinitions