[flang][OpenMP] Set isNewBlock directly on OpenMP constructs

When the PFT builder decides that an evaluation needs a new block
it checks if the evaluation has nested evaluations. In such case it
sets the flag on the first nested evaluation. This works under the
assuption that such an evaluation only serves as a container, and
does not, by itself, generate any code.

This fails for OpenMP constructs that contain nested evaluations
because the top-level evaluation does generate code that wraps the
code from the nested evaluations. In such cases, the code for the
top-level evaluation may be emitted in a wrong place.

When setting the `isNewBlock` flag, recognize OpenMP directives,
and treat them accordingly.

This fixes #139071

Author: kparzysz

flang/include/flang/Lower/PFTBuilder.h

@@ -184,6 +184,11 @@ static constexpr bool isExecutableDirective{common::HasMember<
     A, std::tuple<parser::CompilerDirective, parser::OpenACCConstruct,
                   parser::OpenMPConstruct, parser::CUFKernelDoConstruct>>};
 
+template <typename A>
+static constexpr bool isOpenMPDirective{
+    common::HasMember<A, std::tuple<parser::OpenMPConstruct,
+                                    parser::OpenMPDeclarativeConstruct>>};
+
 template <typename A>
 static constexpr bool isFunctionLike{common::HasMember<
     A, std::tuple<parser::MainProgram, parser::FunctionSubprogram,
@@ -267,6 +272,11 @@ struct Evaluation : EvaluationVariant {
       return pft::isExecutableDirective<std::decay_t<decltype(r)>>;
     }});
   }
+  constexpr bool isOpenMPDirective() const {
+    return visit(common::visitors{[](auto &r) {
+      return pft::isOpenMPDirective<std::decay_t<decltype(r)>>;
+    }});
+  }
 
   /// Return the predicate:  "This is a non-initial, non-terminal construct
   /// statement."  For an IfConstruct, this is ElseIfStmt and ElseStmt.

flang/lib/Lower/PFTBuilder.cpp

@@ -1096,7 +1096,9 @@ class PFTBuilder {
 
     // The first executable statement in the subprogram is preceded by a
     // branch to the entry point, so it starts a new block.
-    if (initialEval->hasNestedEvaluations())
+    // OpenMP directives can generate code around the nested evaluations.
+    if (initialEval->hasNestedEvaluations() &&
+        !initialEval->isOpenMPDirective())
       initialEval = &initialEval->getFirstNestedEvaluation();
     else if (initialEval->isA<Fortran::parser::EntryStmt>())
       initialEval = initialEval->lexicalSuccessor;

flang/test/Lower/OpenMP/multiple-entry-points.f90

@@ -0,0 +1,46 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+! Check the first entry point
+!CHECK: func.func @_QPprocess_a
+!CHECK: omp.parallel
+!CHECK: omp.wsloop
+!CHECK: %[[V0:[0-9]+]] = fir.load %{{[0-9]+}} : !fir.ref<f32>
+!CHECK: %[[V1:[a-z_0-9]+]] = arith.constant 2.000000e+00 : f32
+!CHECK:   = arith.mulf %[[V0]], %[[V1]] fastmath<contract> : f32
+!CHECK: omp.terminator
+!CHECK-NOT: omp
+!CHECK: return
+
+! Check the second entry point
+!CHECK: func.func @_QPprocess_b
+!CHECK: omp.parallel
+!CHECK: fir.do_loop
+!CHECK: %[[V3:[0-9]+]] = fir.load %[[V2:[0-9]+]]#0 : !fir.ref<i32>
+!CHECK: %[[V4:[0-9]+]] = fir.load %[[V2]]#0 : !fir.ref<i32>
+!CHECK:   = arith.muli %[[V3]], %[[V4]] : i32
+!CHECK: omp.terminator
+!CHECK-NOT: omp
+!CHECK: return
+
+subroutine process_a(n, a)
+  integer, intent(in) :: n
+  real, intent(inout) :: a(n)
+  integer :: i
+
+  !$omp parallel do
+  do i = 1, n
+    a(i) = a(i) * 2.0
+  end do
+  !$omp end parallel do
+
+  return
+
+  entry process_b(n, b)
+    
+  !$omp parallel
+  do i = 1, n
+    a(i) = i * i
+  end do
+  !$omp end parallel
+
+end subroutine process_a