[flang][OpenMP] Handle private/firstprivate clauses on sections construct

[NimishMishra]

Summary

This patch adds private/firstprivate support for sections construct. For a source like the following:

!$omp sections private(x) firstprivate(y)
   !$omp section
       <block of code>
   !$omp section
        <block of code>
!$omp end sections

...privatization proceeds to privatize x and y accordingly inside each of the generated omp.section operations.

[NimishMishra]

flang/lib/Lower/OpenMP.cpp

1193

One way to do this would be to reuse collectSymbolSet from https://reviews.llvm.org/D123930

[peixin]

flang/lib/Lower/OpenMP.cpp

1193

The TODO or collectSymbol is no necessary here since the evaluation has the nested information of the parser node.

1197

You can move the code in Bridge.cpp here by getting the clause through the parser node of sections construct using the following code:

auto sectionsConstruct = eval.parentConstruct->getIf<Fortran::parser::SectionsConstruct>();

Then, there is no need to add one argument to genOpenMPConstruct, which is really unexpected.

[NimishMishra]

flang/lib/Lower/OpenMP.cpp

1193

I get your getIf<Fortran::parser::SectionsConstruct>(). I was not aware of this functionality. This is way cleaner than the current implementation.

However, I was thinking of using collectSymbols for an orthogonal problem. Consider the following:

!$omp sections private(x, y)
   !$omp section
       x = x + 1
   !$omp section
       y = y + 1
!$omp end sections

If you check the current generated FIR, both x and y are privatised for both omp.section operations. I wanted to avoid that. I was thinking that since y is not used in the first operation, why privatise it there. For this, I thought to use collectSymbols on omp.section operation, and privatise only the symbols which are present in both private() list and in the section operation.

Let me know what you think of the problem.

[peixin]

flang/lib/Lower/OpenMP.cpp

1193

You are right. It's true that privatization is sometimes redundant. But it seems using collectSymbols will fail for the following case:

program main
  integer :: x
  !$omp sections private(x)
   !$omp section
       call sub()
  !$omp end sections
contains
  subroutine sub()
    x = 1
  end
end

It's OK to have unnecessary privatization since DCE optimization in MLIR or LLVM middle-end will remove it.

[NimishMishra]

Addressed comments.

[NimishMishra]

flang/lib/Lower/OpenMP.cpp

170–171

@peixin I looked into this. Currently sections does not have lastprivate support. Fixing this data race problem shall come once that support is present. Can I take up adding support for lastprivate on sections?

In any case, I guess it would be better to do that in a separate patch.

[peixin]

flang/lib/Lower/OpenMP.cpp

170–171

Remove the brace here. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

170–171

Can I take up adding support for lastprivate on sections?

Sure. Feel free to take up the unhandled clauses on sections.

In any case, I guess it would be better to do that in a separate patch.

Cannot agree more. The single-function patch is preferred for review.

1192

Add one comment here to describe that you only handle the private/firstprivate clause here and the privitization should be performed in section region. Not sure if there should be one FIXME for other clauses. Maybe it's better to do it in line 1335-1348 that add TODO for unhandled clauses? If so, you can add one comment there the privitization should not be done in creating SectionsOp.

1194

Nit:

1198

Nit: highlight the construct name so that users can find the error location easily is something goes wrong.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

641

There are plans performing privitization in MLIR. You can keep it for now. If the final decision is to doing privitization in lowering. The TODOs for all constructs can be removed.

[NimishMishra]

Addressed comments.

[NimishMishra]

flang/lib/Lower/OpenMP.cpp

1192

I checked this. Maybe there's no explicit need to add more TODOs in 1335-1348. Reason being that out of private/firstprivate/lastprivate/reduction/allocate/nowait, the first two are being handled in this patch and the last 3 have some handling there already. For lastprivate, privatizeVars has a TODO already mentioning that no support is present except for wsloop.

1192

Removed this as there is a TODO already inside the function. I only had added this TODO long back. The current TODO for reduction is much better.

[NimishMishra]

Fixed formatting issues

[peixin]

LGTM. The CI failed. Please make sure the CI passed before landing.

flang/lib/Lower/OpenMP.cpp

1192

OK. Makes sense to me.

[NimishMishra]

Re-triggered build. CI passes now.

flang/include/flang/Lower/OpenMP.h

 class AbstractConverter;
 
 namespace pft {
 struct Evaluation;
 struct Variable;
 } // namespace pft
 
 void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
-                        const parser::OpenMPConstruct &);
+                        const parser::OpenMPConstruct &,
+                        const Fortran::parser::OmpClauseList
+                            *ompParentConstructClauseList = nullptr);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
                                    const parser::OpenMPDeclarativeConstruct &);
 int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
 void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
 void genOpenMPReduction(AbstractConverter &,
                         const Fortran::parser::OmpClauseList &clauseList);
 
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENMP_H

flang/lib/Lower/Bridge.cpp

void genFIR(const Fortran::parser::OpenACCDeclarativeConstruct &accDecl) {
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
       genFIR(e);
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenMPConstruct &omp) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
+
     genOpenMPConstruct(*this, getEval(), omp);
 
     const Fortran::parser::OpenMPLoopConstruct *ompLoop =
         std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u);
 
     // If loop is part of an OpenMP Construct then the OpenMP dialect
     // workshare loop operation has already been created. Only the
     // body needs to be created here and the do_loop can be skipped.
if (ompLoop) {
           Fortran::lower::getCollapseValue(*loopOpClauseList);
 
       curEval = &curEval->getFirstNestedEvaluation();
       for (int64_t i = 1; i < collapseValue; i++) {
         curEval = &*std::next(curEval->getNestedEvaluations().begin());
       }
     }
 
+    const Fortran::parser::OmpClauseList *ompSectionsClauseList = nullptr;
+    if (const auto ompSections =
+            std::get_if<Fortran::parser::OpenMPSectionsConstruct>(&omp.u)) {
+      const Fortran::parser::OmpBeginSectionsDirective
+          &ompbeginSectionsDirective =
+              std::get<Fortran::parser::OmpBeginSectionsDirective>(
+                  ompSections->t);
+      const Fortran::parser::OmpClauseList *ompSectionsClauseList =
+          &std::get<Fortran::parser::OmpClauseList>(
+              ompbeginSectionsDirective.t);
+      const Fortran::parser::OmpSectionBlocks &ompSectionBlocks =
+          std::get<Fortran::parser::OmpSectionBlocks>(ompSections->t);
+      auto ompSectionBlocksIter = ompSectionBlocks.v.begin();
+      for (Fortran::lower::pft::Evaluation &e :
+           curEval->getNestedEvaluations()) {
+        mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
+        localSymbols.pushScope();
+        genOpenMPConstruct(*this, e, *ompSectionBlocksIter,
+                           ompSectionsClauseList);
+        for (Fortran::lower::pft::Evaluation &sectionBlockEvaluation :
+             e.getNestedEvaluations())
+          genFIR(sectionBlockEvaluation);
+        localSymbols.popScope();
+        builder->restoreInsertionPoint(insertPt);
+        ompSectionBlocksIter++;
+      }
+    } else
       for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
         genFIR(e);
 
     if (ompLoop)
       genOpenMPReduction(*this, *loopOpClauseList);
 
     localSymbols.popScope();
     builder->restoreInsertionPoint(insertPt);
   }
 

flang/lib/Lower/OpenMP.cpp

static void createPrivateVarSyms(Fortran::lower::AbstractConverter &converter,
   }
 }
 
 template <typename Op>
 static bool privatizeVars(Op &op, Fortran::lower::AbstractConverter &converter,
                           const Fortran::parser::OmpClauseList &opClauseList) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   auto insPt = firOpBuilder.saveInsertionPoint();
+  if (mlir::isa<mlir::omp::SectionOp>(op))
+    firOpBuilder.setInsertionPointToStart(&op.getRegion().back());
+  else
     firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
   bool hasFirstPrivateOp = false;
   bool hasLastPrivateOp = false;
   // We need just one ICmpOp for multiple LastPrivate clauses.
   mlir::arith::CmpIOp cmpOp;
-
   for (const Fortran::parser::OmpClause &clause : opClauseList.v) {
     if (const auto &privateClause =
             std::get_if<Fortran::parser::OmpClause::Private>(&clause.u)) {
       createPrivateVarSyms(converter, privateClause);
     } else if (const auto &firstPrivateClause =
                    std::get_if<Fortran::parser::OmpClause::Firstprivate>(
                        &clause.u)) {
       createPrivateVarSyms(converter, firstPrivateClause);
if (const auto &privateClause =
         createPrivateVarSyms(converter, lastPrivateClause,
                              &(ifOp.getThenRegion().front()));
       } else {
         TODO(converter.getCurrentLocation(),
              "lastprivate clause in constructs other than work-share loop");
       }
       hasLastPrivateOp = true;
     }
   }
   if (hasFirstPrivateOp)
     firOpBuilder.create<mlir::omp::BarrierOp>(converter.getCurrentLocation());
   firOpBuilder.restoreInsertionPoint(insPt);
   return hasLastPrivateOp;
 }
 
 /// The COMMON block is a global structure. \p commonValue is the base address
 /// of the the COMMON block. As the offset from the symbol \p sym, generate the
 /// COMMON block member value (commonValue + offset) for the symbol.
mlir::omp::CriticalOp criticalOp = [&]() {
     }
   }();
   createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation, eval);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
-       const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
+       const Fortran::parser::OpenMPSectionConstruct &sectionConstruct,
+       const Fortran::parser::OmpClauseList *ompParentConstructClauseList =
+           nullptr) {
 
   auto &firOpBuilder = converter.getFirOpBuilder();
   auto currentLocation = converter.getCurrentLocation();
   mlir::omp::SectionOp sectionOp =
       firOpBuilder.create<mlir::omp::SectionOp>(currentLocation);
-  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation, eval);
+  // TODO: For private/firstprivate symbols defined on `omp.sections`
+  // operations,
+  // // inside each `omp.section` region, privatize only the symbols used in
+  // that `omp.section` region.
+  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation, eval,
+                                 ompParentConstructClauseList);
 }
 
 // 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();
   auto currentLocation = converter.getCurrentLocation();
   SmallVector<Value> reductionVars, allocateOperands, allocatorOperands;
   mlir::UnitAttr noWaitClauseOperand;
std::visit(Fortran::common::visitors{
                  },
              },
              atomicConstruct.u);
 }
 
 void Fortran::lower::genOpenMPConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::lower::pft::Evaluation &eval,
-    const Fortran::parser::OpenMPConstruct &ompConstruct) {
+    const Fortran::parser::OpenMPConstruct &ompConstruct,
+    const Fortran::parser::OmpClauseList *ompParentConstructClauseList) {
 
   std::visit(
       common::visitors{
           [&](const Fortran::parser::OpenMPStandaloneConstruct
                   &standaloneConstruct) {
             genOMP(converter, eval, standaloneConstruct);
           },
           [&](const Fortran::parser::OpenMPSectionsConstruct
                   &sectionsConstruct) {
             genOMP(converter, eval, sectionsConstruct);
           },
           [&](const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
-            genOMP(converter, eval, sectionConstruct);
+            genOMP(converter, eval, sectionConstruct,
+                   ompParentConstructClauseList);
           },
           [&](const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
             genOMP(converter, eval, loopConstruct);
           },
           [&](const Fortran::parser::OpenMPDeclarativeAllocate
                   &execAllocConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPDeclarativeAllocate");
           },

flang/test/Lower/OpenMP/sections.f90

 ! This test checks the lowering of OpenMP sections construct with several clauses present
 
 ! RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
 
 !CHECK: func @_QQmain() {
 !CHECK:   %[[COUNT:.*]] = fir.address_of(@_QFEcount) : !fir.ref<i32>
-!CHECK:   %[[DOUBLE_COUNT:.*]] = fir.address_of(@_QFEdouble_count) : !fir.ref<i32>
 !CHECK:   %[[ETA:.*]] = fir.alloca f32 {bindc_name = "eta", uniq_name = "_QFEeta"}
 !CHECK:   %[[CONST_1:.*]] = arith.constant 1 : i32
 !CHECK:   omp.sections allocate(%[[CONST_1]] : i32 -> %0 : !fir.ref<i32>)  {
 !CHECK:     omp.section {
-!CHECK:       {{.*}} = arith.constant 5 : i32
-!CHECK:       fir.store {{.*}} to {{.*}} : !fir.ref<i32>
-!CHECK:       {{.*}} = fir.load %[[COUNT]] : !fir.ref<i32>
-!CHECK:       {{.*}} = fir.load %[[DOUBLE_COUNT]] : !fir.ref<i32>
-!CHECK:       {{.*}} = arith.muli {{.*}}, {{.*}} : i32
-!CHECK:       {{.*}} = fir.convert {{.*}} : (i32) -> f32
-!CHECK:       fir.store {{.*}} to %[[ETA]] : !fir.ref<f32>
+!CHECK:       %[[PRIVATE_ETA:.*]] = fir.alloca f32 {bindc_name = "eta", pinned, uniq_name = "_QFEeta"}
+!CHECK:       %[[PRIVATE_DOUBLE_COUNT:.*]] = fir.alloca i32 {bindc_name = "double_count", pinned, uniq_name = "_QFEdouble_count"} 
+!CHECK:       %[[const:.*]] = arith.constant 5 : i32
+!CHECK:       fir.store %[[const]] to %[[COUNT]] : !fir.ref<i32>
+!CHECK:       %[[temp_count:.*]] = fir.load %[[COUNT]] : !fir.ref<i32>
+!CHECK:       %[[temp_double_count:.*]] = fir.load %[[PRIVATE_DOUBLE_COUNT]] : !fir.ref<i32>
+!CHECK:       %[[result:.*]] = arith.muli %[[temp_count]], %[[temp_double_count]] : i32
+!CHECK:       {{.*}} = fir.convert %[[result]] : (i32) -> f32
+!CHECK:       fir.store {{.*}} to %[[PRIVATE_ETA]] : !fir.ref<f32>
 !CHECK:       omp.terminator
 !CHECK:     }
 !CHECK:     omp.section {
-!CHECK:       {{.*}} = fir.load %[[DOUBLE_COUNT]] : !fir.ref<i32>
-!CHECK:       {{.*}} = arith.constant 1 : i32
-!CHECK:       {{.*}} = arith.addi {{.*}} : i32
-!CHECK:       fir.store {{.*}} to %[[DOUBLE_COUNT]] : !fir.ref<i32>
+!CHECK:       %[[PRIVATE_ETA:.*]] = fir.alloca f32 {bindc_name = "eta", pinned, uniq_name = "_QFEeta"}
+!CHECK:       %[[PRIVATE_DOUBLE_COUNT:.*]] = fir.alloca i32 {bindc_name = "double_count", pinned, uniq_name = "_QFEdouble_count"} 
+!CHECK:       %[[temp:.*]] = fir.load %[[PRIVATE_DOUBLE_COUNT]] : !fir.ref<i32>
+!CHECK:       %[[const:.*]] = arith.constant 1 : i32
+!CHECK:       %[[result:.*]] = arith.addi %[[temp]], %[[const]] : i32
+!CHECK:       fir.store %[[result]] to %[[PRIVATE_DOUBLE_COUNT]] : !fir.ref<i32>
 !CHECK:       omp.terminator
 !CHECK:     }
 !CHECK:     omp.section {
-!CHECK:       {{.*}} = fir.load %[[ETA]] : !fir.ref<f32>
-!CHECK:       {{.*}} = arith.constant 7.000000e+00 : f32
-!CHECK:       {{.*}} = arith.subf {{.*}} : f32
-!CHECK:       fir.store {{.*}} to %[[ETA]] : !fir.ref<f32>
+!CHECK:       %[[PRIVATE_ETA:.*]] = fir.alloca f32 {bindc_name = "eta", pinned, uniq_name = "_QFEeta"}
+!CHECK:       %[[PRIVATE_DOUBLE_COUNT:.*]] = fir.alloca i32 {bindc_name = "double_count", pinned, uniq_name = "_QFEdouble_count"} 
+!CHECK:       %[[temp:.*]] = fir.load %[[PRIVATE_ETA]] : !fir.ref<f32>
+!CHECK:       %[[const:.*]] = arith.constant 7.000000e+00 : f32
+!CHECK:       %[[result:.*]] = arith.subf %[[temp]], %[[const]] : f32
+!CHECK:       fir.store %[[result]] to %[[PRIVATE_ETA]] : !fir.ref<f32>
 !CHECK:       {{.*}} = fir.load %[[COUNT]] : !fir.ref<i32>
-!CHECK:       {{.*}} = fir.convert {{.*}} : (i32) -> f32
-!CHECK:       {{.*}} = fir.load %[[ETA]] : !fir.ref<f32>
-!CHECK:       {{.*}} = arith.mulf {{.*}}, {{.*}} : f32
-!CHECK:       {{.*}} = fir.convert {{.*}} : (f32) -> i32
-!CHECK:       fir.store {{.*}} to %[[COUNT]] : !fir.ref<i32>
+!CHECK:       %[[temp_count:.*]] = fir.convert {{.*}} : (i32) -> f32
+!CHECK:       %[[temp_eta:.*]] = fir.load %[[PRIVATE_ETA]] : !fir.ref<f32>
+!CHECK:       {{.*}} = arith.mulf %[[temp_count]], %[[temp_eta]] : f32
+!CHECK:       %[[result:.*]] = fir.convert {{.*}} : (f32) -> i32
+!CHECK:       fir.store %[[result]] to %[[COUNT]] : !fir.ref<i32>
 !CHECK:       {{.*}} = fir.load %[[COUNT]] : !fir.ref<i32>
-!CHECK:       {{.*}} = fir.convert {{.*}} : (i32) -> f32
-!CHECK:       {{.*}} = fir.load %[[ETA]] : !fir.ref<f32>
-!CHECK:       {{.*}} = arith.subf {{.*}}, {{.*}} : f32
-!CHECK:       {{.*}} = fir.convert {{.*}} : (f32) -> i32
-!CHECK:       fir.store {{.*}} to %[[DOUBLE_COUNT]] : !fir.ref<i32>
+!CHECK:       %[[temp_count:.*]] = fir.convert {{.*}} : (i32) -> f32
+!CHECK:       %[[temp_eta:.*]] = fir.load %[[PRIVATE_ETA]] : !fir.ref<f32>
+!CHECK:       {{.*}} = arith.subf %[[temp_count]], %[[temp_eta]] : f32
+!CHECK:       %[[result:.*]] = fir.convert {{.*}} : (f32) -> i32
+!CHECK:       fir.store %[[result]] to %[[PRIVATE_DOUBLE_COUNT]] : !fir.ref<i32>
 !CHECK:       omp.terminator
 !CHECK:     }
 !CHECK:     omp.terminator
 !CHECK:   }
 !CHECK:   omp.sections nowait {
 !CHECK:     omp.terminator
 !CHECK:   }
 !CHECK:   return
program sample
 
     !$omp sections
     !$omp end sections nowait
 end program sample
 
 !CHECK: func @_QPfirstprivate(%[[ARG:.*]]: !fir.ref<f32> {fir.bindc_name = "alpha"}) {
 !CHECK:   omp.sections {
 !CHECK:     omp.section  {
+!CHECK:         %[[PRIVATE_ALPHA:.*]] = fir.alloca f32 {bindc_name = "alpha", pinned, uniq_name = "_QFfirstprivateEalpha"}
+!CHECK:         %[[temp:.*]] = fir.load %[[ARG]] : !fir.ref<f32>
+!CHECK:         fir.store %[[temp]] to %[[PRIVATE_ALPHA]] : !fir.ref<f32>
+!CHECK:         omp.barrier
 !CHECK:       omp.terminator
 !CHECK:     }
 !CHECK:     omp.terminator
 !CHECK:   }
 !CHECK:   omp.sections {
 !CHECK:     omp.section  {
 !CHECK:       %[[PRIVATE_VAR:.*]] = fir.load %[[ARG]] : !fir.ref<f32>
 !CHECK:       %[[CONSTANT:.*]] = arith.constant 5.000000e+00 : f32

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
   // No sections within omp.sections operation - skip generation. This situation
   // is only possible if there is only a terminator operation inside the
   // sections operation
   if (sectionCBs.empty())
     return success();
 
   assert(isa<omp::SectionOp>(*sectionsOp.region().op_begin()));
 
-  // TODO: Perform appropriate actions according to the data-sharing
-  // attribute (shared, private, firstprivate, ...) of variables.
-  // Currently defaults to shared.
   auto privCB = [&](InsertPointTy, InsertPointTy codeGenIP, llvm::Value &,
                     llvm::Value &vPtr,
                     llvm::Value *&replacementValue) -> InsertPointTy {
     replacementValue = &vPtr;
     return codeGenIP;
   };
 
   // TODO: Perform finalization actions for variables. This has to be