[GPU] Fix constant dimension in case of node from if-op internal body (…

…openvinotoolkit#21214)

* [GPU] Fix constant dimension in case of node from if-op internal body

* Add a condition to restrict to case where allow_new_shape_infer is false

* Add a condition to restrict to case where allow_new_shape_infer is false

* Add functional test

* Add a condition to retrict to case where the program is inner program

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

@@ -137,6 +137,7 @@ class ProgramBuilder final {
 
     bool use_new_shape_infer() const { return allow_new_shape_infer; }
     bool requires_new_shape_infer(const std::shared_ptr<ov::Node>& op) const;
+    bool is_inner_program() const { return m_is_inner_program; }
 
     std::shared_ptr<ov::threading::IStreamsExecutor> get_task_executor() const { return m_task_executor; }
     std::shared_ptr<cldnn::ICompilationContext> get_compilation_context() const { return m_compilation_context; }
@@ -159,6 +160,8 @@ class ProgramBuilder final {
     std::shared_ptr<ov::threading::IStreamsExecutor> m_task_executor;
     std::shared_ptr<cldnn::ICompilationContext> m_compilation_context;
 
+    bool m_is_inner_program = false;
+
     void EnableQueryMode() { queryMode = true; }
     void DisableQueryMode() { queryMode = false; }
 

src/plugins/intel_gpu/src/plugin/ops/constant.cpp

@@ -219,6 +219,12 @@ static void CreateConstantOp(ProgramBuilder& p, const std::shared_ptr<ov::op::v0
             // To pass check_memory_to_set in input_layout::set_data for this case, Set constDims to [N, 1, 1, 1]
             // when constDims is one dim and user op is Loop or TensorIterator.
             consts[op].needsBatchInterpretation = constDims.size() == 1;
+        } else if (ov::is_type<ov::op::v0::Result>(outOp) && !p.use_new_shape_infer() && p.is_inner_program()) {
+            // When IF-operation generates branch-true and branch-false,
+            // simple nodes for both can be created such as Parameter->Result, Constant->Result
+            // And each layout will be like Parameter->Result [N, 1, 1, 1], Constant->Result [1, N, 1, 1], that produces layout mismatch error.
+            // For that case, Constant->Result needs to be [N, 1, 1, 1]
+            consts[op].needsBatchInterpretation = constDims.size() == 1;
         }
     }
 

src/plugins/intel_gpu/src/plugin/program_builder.cpp

@@ -64,7 +64,8 @@ ProgramBuilder::ProgramBuilder(std::shared_ptr<ov::Model> model, cldnn::engine&
     , m_engine(engine)
     , queryMode(false)
     , m_task_executor(task_executor)
-    , m_compilation_context(compilation_context) {
+    , m_compilation_context(compilation_context)
+    , m_is_inner_program(is_inner_program) {
     if (m_task_executor == nullptr)
         m_task_executor = cldnn::program::make_task_executor(m_config);
 

src/plugins/intel_gpu/tests/functional/subgraph_tests/condition.cpp

@@ -51,7 +51,15 @@ enum InnerBodyType {
     /**
      * Inner body with single constant with zero dimensions
      */
-    Type06 = 6
+    Type06 = 6,
+    /**
+     * Inner body with single constant
+    */
+    Type07 = 7,
+    /**
+     * Inner body with single parameter
+    */
+    Type08 = 8
 };
 
 public:
@@ -288,6 +296,39 @@ class InnerBodyType06 : public InnerBodyGenerator {
     }
 };
 
+class InnerBodyType07 : public InnerBodyGenerator {
+protected:
+    std::shared_ptr<ngraph::Function> generate(ov::PartialShape& input_shape, ngraph::element::Type prc) override {
+        auto constant   = ngraph::opset9::Constant::create(prc, input_shape.to_shape(), {2.0f});
+        constant->set_friendly_name("body7_constant");
+        auto result     = std::make_shared<ngraph::opset1::Result>(constant);
+        auto o_layout = result->get_layout();
+        result->set_friendly_name("body7_result");
+        auto body       = std::make_shared<ngraph::Function>(
+            ngraph::OutputVector {result},
+            ngraph::ParameterVector{},
+            "constant_to_result");
+        return body;
+    }
+};
+
+class InnerBodyType08 : public InnerBodyGenerator {
+protected:
+    std::shared_ptr<ngraph::Function> generate(ov::PartialShape& input_shape, ngraph::element::Type prc) override {
+        auto constant   = std::make_shared<ngraph::opset9::Constant>(prc, ngraph::Shape{}, 10.0f);
+        constant->set_friendly_name("body8_const");
+        auto data       = std::make_shared<ngraph::opset9::Parameter>(prc, input_shape);
+        data->set_friendly_name("body8_data");
+        auto result     = std::make_shared<ngraph::opset1::Result>(data);
+        result->set_friendly_name("body8_result");
+        auto body       = std::make_shared<ngraph::Function>(
+            ngraph::OutputVector {result},
+            ngraph::ParameterVector{data},
+            "parameter_to_result");
+        return body;
+    }
+};
+
 static std::shared_ptr<InnerBodyGenerator> get_inner_body_generator(InnerBodyGenerator::InnerBodyType type) {
     std::shared_ptr<InnerBodyGenerator> generator_ptr;
     switch (type) {
@@ -315,6 +356,14 @@ static std::shared_ptr<InnerBodyGenerator> get_inner_body_generator(InnerBodyGen
         {
             return std::make_shared<InnerBodyType06>();
         }
+        case InnerBodyGenerator::InnerBodyType::Type07:
+        {
+            return std::make_shared<InnerBodyType07>();
+        }
+        case InnerBodyGenerator::InnerBodyType::Type08:
+        {
+            return std::make_shared<InnerBodyType08>();
+        }
         default:
         {
             OPENVINO_ASSERT(false, "Not supported type");
@@ -453,6 +502,16 @@ static std::ostream& operator<<(std::ostream& os, const InnerBodyGenerator::Inne
             os << "Type06";
             break;
         }
+        case InnerBodyGenerator::InnerBodyType::Type07:
+        {
+            os << "Type07";
+            break;
+        }
+        case InnerBodyGenerator::InnerBodyType::Type08:
+        {
+            os << "Type08";
+            break;
+        }
         default:
         {
             os << "NONE";
@@ -579,6 +638,92 @@ INSTANTIATE_TEST_SUITE_P(smoke_ConditionGPUTest_static, StaticConditionLayerGPUT
                 StaticConditionLayerGPUTest::getTestCaseName);
 
 
+/// Static shape single layer test
+class StaticConditionSingleLayerGPUTest : public testing::WithParamInterface<ConditionParams>,
+                        virtual public LayerTestsUtils::LayerTestsCommon {
+public:
+    static std::string getTestCaseName(const testing::TestParamInfo<ConditionParams>& obj) {
+        InferenceEngine::SizeVector data_shape;
+        InferenceEngine::Precision data_prc;
+        TestModelGenerator::PredicateTypes pred;
+        std::string targetDevice;
+
+        std::tie(data_shape, data_prc, pred, targetDevice) = obj.param;
+        std::ostringstream result;
+        result << "IS=" << ov::test::utils::vec2str(data_shape) << "_";
+        result << "netPRC=" << std::to_string(data_prc) << "_";
+        result << "ifCond=" << pred << "_";
+        result << "targetDevice=" << targetDevice << "_";
+        auto res_str = result.str();
+        std::replace(res_str.begin(), res_str.end(), '-', '_');
+        return res_str;
+    }
+
+protected:
+    void SetUp() override {
+        targetDevice = ov::test::utils::DEVICE_GPU;
+        TestModelGenerator::PredicateTypes pred;
+        std::tie(data_shape, data_prc, pred, targetDevice) = GetParam();
+        const auto ngShape = ov::PartialShape{data_shape};
+        const auto prc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(data_prc);
+        TestModelGenerator model_generator(InnerBodyGenerator::InnerBodyType::Type07,
+                                            InnerBodyGenerator::InnerBodyType::Type08,
+                                            pred,
+                                            prc,
+                                            ngShape);
+        function = model_generator.get_function();
+    }
+
+    InferenceEngine::Blob::Ptr GenerateInput(const InferenceEngine::InputInfo &info) const override {
+        auto tensor_desc = info.getTensorDesc();
+        auto blob = make_blob_with_precision(tensor_desc);
+        blob->allocate();
+
+        if (tensor_desc.getLayout() == InferenceEngine::SCALAR) {
+            auto prc = tensor_desc.getPrecision();
+            auto scalar_1d = ov::test::utils::make_reshape_view(blob, {1});
+            if (prc == InferenceEngine::Precision::BOOL) {
+                auto mem_blob = dynamic_cast<InferenceEngine::MemoryBlob*>(blob.get());
+                auto mem = mem_blob->rwmap();
+                auto data_ptr = mem.as<bool*>();
+                *data_ptr = false;
+            } else {
+                ov::test::utils::fill_data_with_broadcast(scalar_1d, 0, {20.f});
+            }
+        } else {
+            ov::test::utils::fill_data_with_broadcast(blob, 0, {20.f});
+        }
+        return blob;
+    }
+
+    InferenceEngine::SizeVector data_shape;
+    InferenceEngine::Precision data_prc;
+};
+
+TEST_P(StaticConditionSingleLayerGPUTest, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED();
+    Run();
+}
+
+std::vector<InferenceEngine::Precision> netPrecisions_static_single = {
+    InferenceEngine::Precision::FP32,
+    InferenceEngine::Precision::FP16,
+    InferenceEngine::Precision::I8
+};
+
+std::vector<InferenceEngine::SizeVector> inputs_shape_single = {
+    {64}
+};
+
+INSTANTIATE_TEST_SUITE_P(smoke_ConditionGPUTest_static, StaticConditionSingleLayerGPUTest,
+                testing::Combine(
+                    testing::ValuesIn(inputs_shape_single),
+                    testing::ValuesIn(netPrecisions_static_single),
+                    testing::ValuesIn(if_cond_types),
+                    testing::Values<std::string>(ov::test::utils::DEVICE_GPU)),
+                StaticConditionLayerGPUTest::getTestCaseName);
+
+
 /// Dynamic shape test
 struct InnerBodyTypeParams {
     InnerBodyGenerator::InnerBodyType then_body_type;