[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

1321–1322

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

[peixin]

flang/lib/Lower/OpenMP.cpp

1321–1322

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

1325

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

1321–1322

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

1321–1322

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

223

@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

203–206

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

223

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.

1306

Nit:

1310

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

1320

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.

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

641 ↗

(On Diff #454272)

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

1305

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.

1320

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.

[NimishMishra]

Fixed formatting issues

[peixin]

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

flang/lib/Lower/OpenMP.cpp

1320

OK. Makes sense to me.

[NimishMishra]

Re-triggered build. CI passes now.

flang/lib/Lower/OpenMP.cpp

if (const auto &defaultClause =
         collectSymbols(Fortran::semantics::Symbol::Flag::OmpPrivate);
       else if (defaultClause->v.v ==
                Fortran::parser::OmpDefaultClause::Type::Firstprivate)
         collectSymbols(Fortran::semantics::Symbol::Flag::OmpFirstPrivate);
     }
   }
 
   bool needBarrier = false;
+  if (mlir::isa<mlir::omp::SectionOp>(op))
+    firOpBuilder.setInsertionPointToStart(&op.getRegion().back());
+  else
     firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
   for (auto sym : privatizedSymbols) {
     privatizeSymbol(converter, sym, lastPrivBlock);
     if (sym->test(Fortran::semantics::Symbol::Flag::OmpFirstPrivate) &&
         sym->test(Fortran::semantics::Symbol::Flag::OmpLastPrivate))
       needBarrier = true;
   }
 
   for (auto sym : defaultSymbols)
     if (!symbolsInNestedRegions.contains(sym) &&
         !symbolsInParentRegions.contains(sym) &&
         !privatizedSymbols.contains(sym))
       privatizeSymbol(converter, sym);
 
   // Emit implicit barrier to synchronize threads and avoid data races on
   // initialization of firstprivate variables and post-update of lastprivate
   // variables.
   // FIXME: Emit barrier for lastprivate clause when 'sections' directive has
   // 'nowait' clause. Otherwise, emit barrier when 'sections' directive has
   // both firstprivate and lastprivate clause.
   // Emit implicit barrier for linear clause. Maybe on somewhere else.
   if (needBarrier)
     firOpBuilder.create<mlir::omp::BarrierOp>(converter.getCurrentLocation());
 
   firOpBuilder.restoreInsertionPoint(insPt);
   return hasLastPrivateOp;

 
 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();
+  const Fortran::parser::OpenMPConstruct *parentOmpConstruct =
+      eval.parentConstruct->getIf<Fortran::parser::OpenMPConstruct>();
+  assert(parentOmpConstruct &&
+         "No enclosing parent OpenMPConstruct on SECTION construct");
+  const Fortran::parser::OpenMPSectionsConstruct *sectionsConstruct =
+      std::get_if<Fortran::parser::OpenMPSectionsConstruct>(
+          &parentOmpConstruct->u);
+  assert(sectionsConstruct && "SECTION construct must have parent"
+          "SECTIONS construct");
+  const Fortran::parser::OmpClauseList &sectionsClauseList =
+      std::get<Fortran::parser::OmpClauseList>(
+          std::get<Fortran::parser::OmpBeginSectionsDirective>(
+              sectionsConstruct->t)
+              .t);
+  // Currently only private/firstprivate clause is handled, and
+  // all privatization is done within `omp.section` operations.
   mlir::omp::SectionOp sectionOp =
       firOpBuilder.create<mlir::omp::SectionOp>(currentLocation);
-  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation, eval);
+  createBodyOfOp<omp::SectionOp>(sectionOp, converter, currentLocation, eval,
+                                 &sectionsClauseList);
 }
 
-// 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;
   const auto &sectionsClauseList = std::get<Fortran::parser::OmpClauseList>(

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.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