[Flang][OpenMP] Fixes for unnstructured OpenMP code

[kiranchandramohan]

Summary

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 by @vdonaldson.

Co-authored-by: Val Donaldson <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: Valentin Clement <[e-mail]>

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

157–165

Remove commented code.

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

22

The code for this particular test (ss1) was previously upstreamed. Can push this separately if required.

[vdonaldson]

Looks ok to me. It probably is a good idea to remove the comment block in createEmptyRegionBlocks.

[kiranchandramohan]

Remove commented code block.

[shraiysh]

One small question @kiranchandramohan

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

53

Don't we want an omp.terminator or some form of FIR Operation translation for this exit here? Right now, it just feels like this is missed in the !CHECK: lines. Is that intended?

[kiranchandramohan]

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

53

There is an MLIR cond_br for the exit that branches out of the loop to bb4 if the condition is true.

[shraiysh]

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

53

Right, okay. I missed that. Thanks for clarifying.

flang/lib/Lower/OpenMP.cpp

mlir::emitWarning(converter.getCurrentLocation(),
                       "bits size and will be narrowed into 64 bits.");
   }
   assert((loopVarTypeSize == 32 || loopVarTypeSize == 64) &&
          "OpenMP loop iteration variable size must be transformed into 32-bit "
          "or 64-bit");
   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.
 /// \param [in]    outerCombined - is this an outer operation - prevents
 ///                                privatization.
 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();
-  // If arguments for the region are provided then create the block with those
-  // arguments. Also update the symbol's address with the mlir argument values.
-  // e.g. For loops the arguments are the induction variable. And all further
+  auto &firOpBuilder = converter.getFirOpBuilder();
+  // If an argument for the region is provided then create the block with that
+  // argument. Also update the symbol's address with the mlir argument value.
+  // 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;
     for (const Fortran::semantics::Symbol *arg : args)
       loopVarTypeSize = std::max(loopVarTypeSize, arg->GetUltimate().size());
     mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
     SmallVector<Type> tiv;
     SmallVector<Location> locs;
for (const Fortran::semantics::Symbol *arg : args) {
       argIndex++;
     }
   } else {
     firOpBuilder.createBlock(&op.getRegion());
   }
   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);
   } else {
     firOpBuilder.create<mlir::omp::TerminatorOp>(loc);
   }
 
   // Reset the insertion point to the start of the first block.
for (const Fortran::parser::OmpClause &clause : opClauseList.v) {
     }
   }
   // Create and insert the operation.
   auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
       currentLocation, argTy, ifClauseOperand, numThreadsClauseOperand,
       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);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPBlockConstruct &blockConstruct) {
genOMP(Fortran::lower::AbstractConverter &converter,
 
   if (blockDirective.v == llvm::omp::OMPD_parallel) {
     // Create and insert the operation.
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
         currentLocation, argTy, ifClauseOperand, numThreadsClauseOperand,
         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");
   }
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
if (const auto &endClauseList =
               loopConstruct.t)) {
     const auto &clauseList =
         std::get<Fortran::parser::OmpClauseList>((*endClauseList).t);
     for (const Fortran::parser::OmpClause &clause : clauseList.v)
       if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
         wsLoopOp.nowaitAttr(firOpBuilder.getUnitAttr());
   }
 
-  createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation,
+  createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation, eval,
                                 &wsLoopOpClauseList, iv);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
if (name.empty()) {
       if (!global)
         global = modBuilder.create<mlir::omp::CriticalDeclareOp>(
             currentLocation, name, hint);
       return firOpBuilder.create<mlir::omp::CriticalOp>(
           currentLocation, mlir::FlatSymbolRefAttr::get(
                                firOpBuilder.getContext(), global.sym_name()));
     }
   }();
-  createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation);
+  createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation, eval);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
 
   auto &firOpBuilder = converter.getFirOpBuilder();
   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
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPSectionsConstruct &sectionsConstruct) {
   auto &firOpBuilder = converter.getFirOpBuilder();
if (dir == llvm::omp::Directive::OMPD_parallel_sections) {
     createCombinedParallelOp<Fortran::parser::OmpBeginSectionsDirective>(
         converter, eval,
         std::get<Fortran::parser::OmpBeginSectionsDirective>(
             sectionsConstruct.t));
     auto sectionsOp = firOpBuilder.create<mlir::omp::SectionsOp>(
         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);
   }
 }
 
 static void genOmpAtomicHintAndMemoryOrderClauses(
     Fortran::lower::AbstractConverter &converter,
     const Fortran::parser::OmpAtomicClauseList &clauseList,
     mlir::IntegerAttr &hint,
     mlir::omp::ClauseMemoryOrderKindAttr &memory_order) {

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

+! 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