[Flang][OpenMP] Fixes for unstructured OpenMP code

Since the FIR operations are mostly structured, it is only the functions
that could contain multiple blocks inside an operation. This changes
with OpenMP since OpenMP regions can contain multiple blocks. For
unstructured code, the blocks are created in advance and belong to the
top-level function. This caused code in OpenMP region to be placed under
the function level.

In this fix, if the OpenMP region is unstructured then new blocks are
created inside it.

Note1: This is part of upstreaming from the fir-dev branch of
https://github.com/flang-compiler/f18-llvm-project. The code in this patch is a
subset of the changes in https://github.com/flang-compiler/f18-llvm-project/pull/1178.

Reviewed By: vdonaldson

Differential Revision: https://reviews.llvm.org/D126293

Co-authored-by: Val Donaldson <vdonaldson@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>
Co-authored-by: Valentin Clement <clementval@gmail.com>

flang/lib/Lower/OpenMP.cpp

@@ -138,11 +138,35 @@ static mlir::Type getLoopVarType(Fortran::lower::AbstractConverter &converter,
   return converter.getFirOpBuilder().getIntegerType(loopVarTypeSize);
 }
 
+/// Create empty blocks for the current region.
+/// These blocks replace blocks parented to an enclosing region.
+void createEmptyRegionBlocks(
+    fir::FirOpBuilder &firOpBuilder,
+    std::list<Fortran::lower::pft::Evaluation> &evaluationList) {
+  auto *region = &firOpBuilder.getRegion();
+  for (auto &eval : evaluationList) {
+    if (eval.block) {
+      if (eval.block->empty()) {
+        eval.block->erase();
+        eval.block = firOpBuilder.createBlock(region);
+      } else {
+        [[maybe_unused]] auto &terminatorOp = eval.block->back();
+        assert((mlir::isa<mlir::omp::TerminatorOp>(terminatorOp) ||
+                mlir::isa<mlir::omp::YieldOp>(terminatorOp)) &&
+               "expected terminator op");
+      }
+    }
+    if (eval.hasNestedEvaluations())
+      createEmptyRegionBlocks(firOpBuilder, eval.getNestedEvaluations());
+  }
+}
+
 /// Create the body (block) for an OpenMP Operation.
 ///
 /// \param [in]    op - the operation the body belongs to.
 /// \param [inout] converter - converter to use for the clauses.
 /// \param [in]    loc - location in source code.
+/// \param [in]    eval - current PFT node/evaluation.
 /// \oaran [in]    clauses - list of clauses to process.
 /// \param [in]    args - block arguments (induction variable[s]) for the
 ////                      region.
@@ -151,14 +175,14 @@ static mlir::Type getLoopVarType(Fortran::lower::AbstractConverter &converter,
 template <typename Op>
 static void
 createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
-               mlir::Location &loc,
+               mlir::Location &loc, Fortran::lower::pft::Evaluation &eval,
                const Fortran::parser::OmpClauseList *clauses = nullptr,
                const SmallVector<const Fortran::semantics::Symbol *> &args = {},
                bool outerCombined = false) {
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  auto &firOpBuilder = converter.getFirOpBuilder();
-  // If arguments for the region are provided then create the block with those
+  // If an argument for the region is provided then create the block with that
-  // arguments. Also update the symbol's address with the mlir argument values.
+  // argument. Also update the symbol's address with the mlir argument value.
-  // e.g. For loops the arguments are the induction variable. And all further
+  // e.g. For loops the argument is the induction variable. And all further
   // uses of the induction variable should use this mlir value.
   if (args.size()) {
     std::size_t loopVarTypeSize = 0;
@@ -184,7 +208,10 @@ createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
   auto &block = op.getRegion().back();
   firOpBuilder.setInsertionPointToStart(&block);
 
-  // Insert the terminator.
+  if (eval.lowerAsUnstructured())
+    createEmptyRegionBlocks(firOpBuilder, eval.getNestedEvaluations());
+
+  // Ensure the block is well-formed by inserting terminators.
   if constexpr (std::is_same_v<Op, omp::WsLoopOp>) {
     mlir::ValueRange results;
     firOpBuilder.create<mlir::omp::YieldOp>(loc, results);
@@ -369,7 +396,7 @@ createCombinedParallelOp(Fortran::lower::AbstractConverter &converter,
       allocateOperands, allocatorOperands, /*reduction_vars=*/ValueRange(),
       /*reductions=*/nullptr, procBindKindAttr);
 
-  createBodyOfOp<omp::ParallelOp>(parallelOp, converter, currentLocation,
+  createBodyOfOp<omp::ParallelOp>(parallelOp, converter, currentLocation, eval,
                                   &opClauseList, /*iv=*/{},
                                   /*isCombined=*/true);
 }
@@ -461,26 +488,27 @@ genOMP(Fortran::lower::AbstractConverter &converter,
         allocateOperands, allocatorOperands, /*reduction_vars=*/ValueRange(),
         /*reductions=*/nullptr, procBindKindAttr);
     createBodyOfOp<omp::ParallelOp>(parallelOp, converter, currentLocation,
-                                    &opClauseList);
+                                    eval, &opClauseList);
   } else if (blockDirective.v == llvm::omp::OMPD_master) {
     auto masterOp =
         firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
-    createBodyOfOp<omp::MasterOp>(masterOp, converter, currentLocation);
+    createBodyOfOp<omp::MasterOp>(masterOp, converter, currentLocation, eval);
   } else if (blockDirective.v == llvm::omp::OMPD_single) {
     auto singleOp = firOpBuilder.create<mlir::omp::SingleOp>(
         currentLocation, allocateOperands, allocatorOperands, nowaitAttr);
-    createBodyOfOp<omp::SingleOp>(singleOp, converter, currentLocation);
+    createBodyOfOp<omp::SingleOp>(singleOp, converter, currentLocation, eval);
   } else if (blockDirective.v == llvm::omp::OMPD_ordered) {
     auto orderedOp = firOpBuilder.create<mlir::omp::OrderedRegionOp>(
         currentLocation, /*simd=*/nullptr);
-    createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation);
+    createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation,
+                                         eval);
   } else if (blockDirective.v == llvm::omp::OMPD_task) {
     auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
         currentLocation, ifClauseOperand, finalClauseOperand, untiedAttr,
         mergeableAttr, /*in_reduction_vars=*/ValueRange(),
         /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
         allocatorOperands);
-    createBodyOfOp(taskOp, converter, currentLocation, &opClauseList);
+    createBodyOfOp(taskOp, converter, currentLocation, eval, &opClauseList);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
@@ -644,7 +672,7 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
         wsLoopOp.nowaitAttr(firOpBuilder.getUnitAttr());
   }
 
-  createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation,
+  createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation, eval,
                                 &wsLoopOpClauseList, iv);
 }
 
@@ -688,7 +716,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
                                firOpBuilder.getContext(), global.sym_name()));
     }
   }();
-  createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation);
+  createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation, eval);
 }
 
 static void
@@ -700,7 +728,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
   auto currentLocation = converter.getCurrentLocation();
   mlir::omp::SectionOp sectionOp =
       firOpBuilder.create<mlir::omp::SectionOp>(currentLocation);
-  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation);
+  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation, eval);
 }
 
 // TODO: Add support for reduction
@@ -757,14 +785,15 @@ genOMP(Fortran::lower::AbstractConverter &converter,
         currentLocation, /*reduction_vars*/ ValueRange(),
         /*reductions=*/nullptr, allocateOperands, allocatorOperands,
         /*nowait=*/nullptr);
-    createBodyOfOp(sectionsOp, converter, currentLocation);
+    createBodyOfOp(sectionsOp, converter, currentLocation, eval);
 
     // Sections Construct
   } else if (dir == llvm::omp::Directive::OMPD_sections) {
     auto sectionsOp = firOpBuilder.create<mlir::omp::SectionsOp>(
         currentLocation, reductionVars, /*reductions = */ nullptr,
         allocateOperands, allocatorOperands, noWaitClauseOperand);
-    createBodyOfOp<omp::SectionsOp>(sectionsOp, converter, currentLocation);
+    createBodyOfOp<omp::SectionsOp>(sectionsOp, converter, currentLocation,
+                                    eval);
   }
 }
 

flang/test/Lower/OpenMP/omp-unstructured.f90

@@ -0,0 +1,113 @@
+! Test unstructured code adjacent to and inside OpenMP constructs.
+
+! RUN: bbc %s -fopenmp -o "-" | FileCheck %s
+
+! CHECK-LABEL: func @_QPss1{{.*}} {
+! CHECK:   br ^bb1
+! CHECK: ^bb1:  // 2 preds: ^bb0, ^bb3
+! CHECK:   cond_br %{{[0-9]*}}, ^bb2, ^bb4
+! CHECK: ^bb2:  // pred: ^bb1
+! CHECK:   cond_br %{{[0-9]*}}, ^bb4, ^bb3
+! CHECK: ^bb3:  // pred: ^bb2
+! CHECK:   @_FortranAioBeginExternalListOutput
+! CHECK:   br ^bb1
+! CHECK: ^bb4:  // 2 preds: ^bb1, ^bb2
+! CHECK:   omp.master  {
+! CHECK:     @_FortranAioBeginExternalListOutput
+! CHECK:     omp.terminator
+! CHECK:   }
+! CHECK:   @_FortranAioBeginExternalListOutput
+! CHECK: }
+subroutine ss1(n) ! unstructured code followed by a structured OpenMP construct
+  do i = 1, 3
+    if (i .eq. n) exit
+    print*, 'ss1-A', i
+  enddo
+  !$omp master
+    print*, 'ss1-B', i
+  !$omp end master
+  print*
+end
+
+! CHECK-LABEL: func @_QPss2{{.*}} {
+! CHECK:   omp.master  {
+! CHECK:     @_FortranAioBeginExternalListOutput
+! CHECK:     br ^bb1
+! CHECK:   ^bb1:  // 2 preds: ^bb0, ^bb3
+! CHECK:     cond_br %{{[0-9]*}}, ^bb2, ^bb4
+! CHECK:   ^bb2:  // pred: ^bb1
+! CHECK:     cond_br %{{[0-9]*}}, ^bb4, ^bb3
+! CHECK:   ^bb3:  // pred: ^bb2
+! CHECK:     @_FortranAioBeginExternalListOutput
+! CHECK:     br ^bb1
+! CHECK:   ^bb4:  // 2 preds: ^bb1, ^bb2
+! CHECK:     omp.terminator
+! CHECK:   }
+! CHECK:   @_FortranAioBeginExternalListOutput
+! CHECK:   @_FortranAioBeginExternalListOutput
+! CHECK: }
+subroutine ss2(n) ! unstructured OpenMP construct; loop exit inside construct
+  !$omp master
+    print*, 'ss2-A', n
+    do i = 1, 3
+      if (i .eq. n) exit
+      print*, 'ss2-B', i
+    enddo
+  !$omp end master
+  print*, 'ss2-C', i
+  print*
+end
+
+! CHECK-LABEL: func @_QPss3{{.*}} {
+! CHECK:   omp.parallel  {
+! CHECK:     br ^bb1
+! CHECK:   ^bb1:  // 2 preds: ^bb0, ^bb2
+! CHECK:     cond_br %{{[0-9]*}}, ^bb2, ^bb3
+! CHECK:   ^bb2:  // pred: ^bb1
+! CHECK:     omp.wsloop {{.*}} {
+! CHECK:     @_FortranAioBeginExternalListOutput
+! CHECK:       omp.yield
+! CHECK:     }
+! CHECK:     omp.wsloop {{.*}} {
+! CHECK:       br ^bb1
+! CHECK:     ^bb1:  // 2 preds: ^bb0, ^bb3
+! CHECK:       cond_br %{{[0-9]*}}, ^bb2, ^bb4
+! CHECK:     ^bb2:  // pred: ^bb1
+! CHECK:       cond_br %{{[0-9]*}}, ^bb4, ^bb3
+! CHECK:     ^bb3:  // pred: ^bb2
+! CHECK:       @_FortranAioBeginExternalListOutput
+! CHECK:       br ^bb1
+! CHECK:     ^bb4:  // 2 preds: ^bb1, ^bb2
+! CHECK:       omp.yield
+! CHECK:     }
+! CHECK:     br ^bb1
+! CHECK:   ^bb3:  // pred: ^bb1
+! CHECK:     omp.terminator
+! CHECK:   }
+! CHECK: }
+subroutine ss3(n) ! nested unstructured OpenMP constructs
+  !$omp parallel
+    do i = 1, 3
+      !$omp do
+        do k = 1, 3
+          print*, 'ss3-A', k
+        enddo
+      !$omp end do
+      !$omp do
+        do j = 1, 3
+          do k = 1, 3
+            if (k .eq. n) exit
+            print*, 'ss3-B', k
+          enddo
+        enddo
+      !$omp end do
+    enddo
+  !$omp end parallel
+end
+
+! CHECK-LABEL: func @_QQmain
+program p
+  call ss1(2)
+  call ss2(2)
+  call ss3(2)
+end