[Flang][OpenMP] Push genEval calls to individual operations, NFC #77758

kparzysz · 2024-01-11T12:48:45Z

Introduce genNestedEvaluations that will lower all evaluations nested in the given, accouting for a potential COLLAPSE directive.

Recursive lowering [2/5]

llvmbot · 2024-01-11T12:49:12Z

@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-flang-fir-hlfir

Author: Krzysztof Parzyszek (kparzysz)

Changes

Introduce genNestedEvaluations that will lower all evaluations nested in the given, accouting for a potential COLLAPSE directive.

Recursive lowering [2/5]

Full diff: https://github.com/llvm/llvm-project/pull/77758.diff

1 Files Affected:

(modified) flang/lib/Lower/OpenMP.cpp (+66-62)

diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 99690b03eca1d3..496b4ba27a0533 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -110,6 +110,32 @@ static void gatherFuncAndVarSyms(
   }
 }
 
+static Fortran::lower::pft::Evaluation *
+getEvalPastCollapse(Fortran::lower::pft::Evaluation &eval, int collapseValue) {
+  if (collapseValue == 0)
+    return &eval;
+
+  Fortran::lower::pft::Evaluation *curEval = &eval.getFirstNestedEvaluation();
+  for (int i = 1; i < collapseValue; i++) {
+    // The nested evaluations should be DoConstructs (i.e. they should form
+    // a loop nest). Each DoConstruct is a tuple <NonLabelDoStmt, Block,
+    // EndDoStmt>.
+    assert(curEval->isA<Fortran::parser::DoConstruct>());
+    curEval = &*std::next(curEval->getNestedEvaluations().begin());
+  }
+  return curEval;
+}
+
+static void genNestedEvaluations(Fortran::lower::AbstractConverter &converter,
+                                 Fortran::lower::pft::Evaluation &eval,
+                                 int collapseValue = 0) {
+  Fortran::lower::pft::Evaluation *curEval =
+      getEvalPastCollapse(eval, collapseValue);
+
+  for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
+    converter.genEval(e);
+}
+
 //===----------------------------------------------------------------------===//
 // DataSharingProcessor
 //===----------------------------------------------------------------------===//
@@ -2944,7 +2970,7 @@ genOmpFlush(Fortran::lower::AbstractConverter &converter,
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        Fortran::semantics::SemanticsContext &semanticsContext,
        const Fortran::parser::OpenMPStandaloneConstruct &standaloneConstruct) {
   std::visit(
@@ -3025,15 +3051,17 @@ createSimdLoop(Fortran::lower::AbstractConverter &converter,
   createBodyOfOp<mlir::omp::SimdLoopOp>(simdLoopOp, converter, loc, eval,
                                         &loopOpClauseList, iv,
                                         /*outer=*/false, &dsp);
+
+  genNestedEvaluations(converter, eval,
+                       Fortran::lower::getCollapseValue(loopOpClauseList));
 }
 
-static void
-createWsLoop(Fortran::lower::AbstractConverter &converter,
-             Fortran::lower::pft::Evaluation &eval,
-             llvm::omp::Directive ompDirective,
-             const Fortran::parser::OmpClauseList &beginClauseList,
-             const Fortran::parser::OmpClauseList *endClauseList,
-             mlir::Location loc) {
+static void createWsLoop(Fortran::lower::AbstractConverter &converter,
+                         Fortran::lower::pft::Evaluation &eval,
+                         llvm::omp::Directive ompDirective,
+                         const Fortran::parser::OmpClauseList &beginClauseList,
+                         const Fortran::parser::OmpClauseList *endClauseList,
+                         mlir::Location loc) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   DataSharingProcessor dsp(converter, beginClauseList, eval);
   dsp.processStep1();
@@ -3107,9 +3135,13 @@ createWsLoop(Fortran::lower::AbstractConverter &converter,
   createBodyOfOp<mlir::omp::WsLoopOp>(wsLoopOp, converter, loc, eval,
                                       &beginClauseList, iv,
                                       /*outer=*/false, &dsp);
+
+  genNestedEvaluations(converter, eval,
+                       Fortran::lower::getCollapseValue(beginClauseList));
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
+                   Fortran::lower::SymMap &symTable,
                    Fortran::lower::pft::Evaluation &eval,
                    Fortran::semantics::SemanticsContext &semanticsContext,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
@@ -3179,11 +3211,12 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
     createWsLoop(converter, eval, ompDirective, loopOpClauseList, endClauseList,
                  currentLocation);
   }
+  genOpenMPReduction(converter, loopOpClauseList);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        Fortran::semantics::SemanticsContext &semanticsContext,
        const Fortran::parser::OpenMPBlockConstruct &blockConstruct) {
   const auto &beginBlockDirective =
@@ -3298,11 +3331,14 @@ genOMP(Fortran::lower::AbstractConverter &converter,
     break;
   }
   }
+
+  genNestedEvaluations(converter, eval);
+  genOpenMPReduction(converter, beginClauseList);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
@@ -3336,11 +3372,12 @@ genOMP(Fortran::lower::AbstractConverter &converter,
   }();
   createBodyOfOp<mlir::omp::CriticalOp>(criticalOp, converter, currentLocation,
                                         eval);
+  genNestedEvaluations(converter, eval);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
   mlir::Location currentLocation = converter.getCurrentLocation();
   const Fortran::parser::OpenMPConstruct *parentOmpConstruct =
@@ -3359,14 +3396,17 @@ genOMP(Fortran::lower::AbstractConverter &converter,
               .t);
   // Currently only private/firstprivate clause is handled, and
   // all privatization is done within `omp.section` operations.
+  symTable.pushScope();
   genOpWithBody<mlir::omp::SectionOp>(converter, eval, currentLocation,
                                       /*outerCombined=*/false,
                                       &sectionsClauseList);
+  genNestedEvaluations(converter, eval);
+  symTable.popScope();
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPSectionsConstruct &sectionsConstruct) {
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands;
@@ -3406,11 +3446,13 @@ genOMP(Fortran::lower::AbstractConverter &converter,
                                        /*reduction_vars=*/mlir::ValueRange(),
                                        /*reductions=*/nullptr, allocateOperands,
                                        allocatorOperands, nowaitClauseOperand);
+
+  genNestedEvaluations(converter, eval);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
-       Fortran::lower::pft::Evaluation &eval,
+       Fortran::lower::SymMap &symTable, Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPAtomicConstruct &atomicConstruct) {
   std::visit(
       Fortran::common::visitors{
@@ -3504,6 +3546,7 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
+                   Fortran::lower::SymMap &symTable,
                    Fortran::semantics::SemanticsContext &semanticsContext,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPConstruct &ompConstruct) {
@@ -3511,17 +3554,18 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
       Fortran::common::visitors{
           [&](const Fortran::parser::OpenMPStandaloneConstruct
                   &standaloneConstruct) {
-            genOMP(converter, eval, semanticsContext, standaloneConstruct);
+            genOMP(converter, symTable, eval, semanticsContext,
+                   standaloneConstruct);
           },
           [&](const Fortran::parser::OpenMPSectionsConstruct
                   &sectionsConstruct) {
-            genOMP(converter, eval, sectionsConstruct);
+            genOMP(converter, symTable, eval, sectionsConstruct);
           },
           [&](const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
-            genOMP(converter, eval, sectionConstruct);
+            genOMP(converter, symTable, eval, sectionConstruct);
           },
           [&](const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
-            genOMP(converter, eval, semanticsContext, loopConstruct);
+            genOMP(converter, symTable, eval, semanticsContext, loopConstruct);
           },
           [&](const Fortran::parser::OpenMPDeclarativeAllocate
                   &execAllocConstruct) {
@@ -3536,14 +3580,14 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
             TODO(converter.getCurrentLocation(), "OpenMPAllocatorsConstruct");
           },
           [&](const Fortran::parser::OpenMPBlockConstruct &blockConstruct) {
-            genOMP(converter, eval, semanticsContext, blockConstruct);
+            genOMP(converter, symTable, eval, semanticsContext, blockConstruct);
           },
           [&](const Fortran::parser::OpenMPAtomicConstruct &atomicConstruct) {
-            genOMP(converter, eval, atomicConstruct);
+            genOMP(converter, symTable, eval, atomicConstruct);
           },
           [&](const Fortran::parser::OpenMPCriticalConstruct
                   &criticalConstruct) {
-            genOMP(converter, eval, criticalConstruct);
+            genOMP(converter, symTable, eval, criticalConstruct);
           },
       },
       ompConstruct.u);
@@ -3607,47 +3651,8 @@ void Fortran::lower::genOpenMPConstruct(
     Fortran::semantics::SemanticsContext &semanticsContext,
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenMPConstruct &omp) {
-
   symTable.pushScope();
-  genOMP(converter, semanticsContext, eval, omp);
-
-  const Fortran::parser::OpenMPLoopConstruct *ompLoop =
-      std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u);
-  const Fortran::parser::OpenMPBlockConstruct *ompBlock =
-      std::get_if<Fortran::parser::OpenMPBlockConstruct>(&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.
-  // Skip the number of collapsed loops, which is 1 when there is a
-  // no collapse requested.
-
-  Fortran::lower::pft::Evaluation *curEval = &eval;
-  const Fortran::parser::OmpClauseList *loopOpClauseList = nullptr;
-  if (ompLoop) {
-    loopOpClauseList = &std::get<Fortran::parser::OmpClauseList>(
-        std::get<Fortran::parser::OmpBeginLoopDirective>(ompLoop->t).t);
-    int64_t collapseValue = Fortran::lower::getCollapseValue(*loopOpClauseList);
-
-    curEval = &curEval->getFirstNestedEvaluation();
-    for (int64_t i = 1; i < collapseValue; i++) {
-      curEval = &*std::next(curEval->getNestedEvaluations().begin());
-    }
-  }
-
-  for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
-    converter.genEval(e);
-
-  if (ompLoop) {
-    genOpenMPReduction(converter, *loopOpClauseList);
-  } else if (ompBlock) {
-    const auto &blockStart =
-        std::get<Fortran::parser::OmpBeginBlockDirective>(ompBlock->t);
-    const auto &blockClauses =
-        std::get<Fortran::parser::OmpClauseList>(blockStart.t);
-    genOpenMPReduction(converter, blockClauses);
-  }
-
+  genOMP(converter, symTable, semanticsContext, eval, omp);
   symTable.popScope();
 }
 
@@ -3656,8 +3661,7 @@ void Fortran::lower::genOpenMPDeclarativeConstruct(
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenMPDeclarativeConstruct &omp) {
   genOMP(converter, eval, omp);
-  for (Fortran::lower::pft::Evaluation &e : eval.getNestedEvaluations())
-    converter.genEval(e);
+  genNestedEvaluations(converter, eval);
 }
 
 void Fortran::lower::genOpenMPSymbolProperties(

These two constructs were both handled in `genOMP` for loop constructs. There is some shared code between the two, but there are also enough differences to separate these two cases into individual functions. The shared code may be placed into a helper function later if needed. Recursive lowering [1/5]

Introduce `genNestedEvaluations` that will lower all evaluations nested in the given, accouting for a potential COLLAPSE directive. Recursive lowering [2/5]

skatrak

Thanks Krzysztof! Just some small comments, but I like the approach.

skatrak · 2024-01-12T12:38:27Z

flang/lib/Lower/OpenMP.cpp

@@ -110,6 +110,32 @@ static void gatherFuncAndVarSyms(
  }
 }

+static Fortran::lower::pft::Evaluation *
+getEvalPastCollapse(Fortran::lower::pft::Evaluation &eval, int collapseValue) {


Nit: I tend to read this function name as meaning "return the next eval after a collapsed nest of loops", when it actually looks to be returning the eval for the innermost collapsed loop. Maybe the name can be changed to something a bit more descriptive.

flang/lib/Lower/OpenMP.cpp

skatrak

Thanks for addressing my comments, LGTM.

skatrak · 2024-01-15T11:07:10Z

flang/lib/Lower/OpenMP.cpp

@@ -110,6 +110,34 @@ static void gatherFuncAndVarSyms(
  }
 }

+static Fortran::lower::pft::Evaluation *
+getCollapsedEval(Fortran::lower::pft::Evaluation &eval, int collapseValue) {


Nit: Maybe getCollapsedLoopEval would be slightly more self-explanatory?

I'll make this change in another commit.

kiranchandramohan

LG.

Please add a pointer to the discussion where it was agreed to pass a reference to the localSymbolTable.

kiranchandramohan · 2024-01-15T11:48:17Z

flang/lib/Lower/OpenMP.cpp

+  // Return the Evaluation of the innermost collapsed loop, or the current
+  // evaluation, if there is nothing to collapse.
+  if (collapseValue == 0)
+    return &eval;


Nit: Is it better to convert this to an assert (for > 0) and move this code to the parent function?

I'll make this change in another commit.

kparzysz · 2024-01-15T13:54:43Z

Here's link to my comment about symbol table in the thread, followed by response from Valentin:
https://discourse.llvm.org/t/openmp-lowering-from-pft-to-fir/75263/49

kparzysz · 2024-01-15T14:25:35Z

Suggested changed made in 705d927.

The buildkite times are really long now (7h on Friday), and I didn't want to delay merging of this PR by that much, plus the changes were really minor.

…m#77758) Introduce `genNestedEvaluations` that will lower all evaluations nested in the given, accouting for a potential COLLAPSE directive. Recursive lowering [2/5]

kparzysz requested review from clementval, skatrak, kiranchandramohan and jeanPerier January 11, 2024 12:48

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir flang:openmp labels Jan 11, 2024

kparzysz added 2 commits January 11, 2024 06:56

[Flang][OpenMP] Push genEval calls to individual operations, NFC

841ae5f

Introduce `genNestedEvaluations` that will lower all evaluations nested in the given, accouting for a potential COLLAPSE directive. Recursive lowering [2/5]

kparzysz force-pushed the users/kparzysz/spr/a02-geneval-individual branch from fd51d9b to 841ae5f Compare January 11, 2024 13:19

kparzysz force-pushed the users/kparzysz/spr/a01-separate-ws-simd branch from dfdee7e to 62f3165 Compare January 11, 2024 13:19

skatrak reviewed Jan 12, 2024

View reviewed changes

Base automatically changed from users/kparzysz/spr/a01-separate-ws-simd to main January 12, 2024 22:46

kparzysz added 4 commits January 12, 2024 18:26

Rename getEvalPastCollapse to getCollapsedEval

8733efc

Add symbol table parameter to remaining genOMP functions

ffda6b2

Merge branch 'main' into users/kparzysz/spr/a02-geneval-individual

28d4f7d

Make all genOMP functions have the same interface

0d9755d

skatrak approved these changes Jan 15, 2024

View reviewed changes

kiranchandramohan approved these changes Jan 15, 2024

View reviewed changes

kparzysz merged commit c5a9e35 into main Jan 15, 2024
4 checks passed

kparzysz deleted the users/kparzysz/spr/a02-geneval-individual branch January 15, 2024 14:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Flang][OpenMP] Push genEval calls to individual operations, NFC #77758

[Flang][OpenMP] Push genEval calls to individual operations, NFC #77758

kparzysz commented Jan 11, 2024

llvmbot commented Jan 11, 2024 •

edited

skatrak left a comment

skatrak Jan 12, 2024

skatrak left a comment

skatrak Jan 15, 2024

kparzysz Jan 15, 2024

kiranchandramohan left a comment

kiranchandramohan Jan 15, 2024

kparzysz Jan 15, 2024

kparzysz commented Jan 15, 2024

kparzysz commented Jan 15, 2024

[Flang][OpenMP] Push genEval calls to individual operations, NFC #77758

[Flang][OpenMP] Push genEval calls to individual operations, NFC #77758

Conversation

kparzysz commented Jan 11, 2024

llvmbot commented Jan 11, 2024 • edited

skatrak left a comment

Choose a reason for hiding this comment

skatrak Jan 12, 2024

Choose a reason for hiding this comment

skatrak left a comment

Choose a reason for hiding this comment

skatrak Jan 15, 2024

Choose a reason for hiding this comment

kparzysz Jan 15, 2024

Choose a reason for hiding this comment

kiranchandramohan left a comment

Choose a reason for hiding this comment

kiranchandramohan Jan 15, 2024

Choose a reason for hiding this comment

kparzysz Jan 15, 2024

Choose a reason for hiding this comment

kparzysz commented Jan 15, 2024

kparzysz commented Jan 15, 2024

llvmbot commented Jan 11, 2024 •

edited