[flang][OpenMP][WIP] Separate clause translation from directive translation

[shraiysh]

Summary

This patch separates clause translation from direction translation while converting F90 to OpenMP Dialect. This clause handler can be reused for other operations like wsloop, sections etc.

[clementval]

flang/lib/Lower/OpenMP.cpp

130

The name is a bit misleading. Something with OpenMP in the name would make more sense IMO.

[shraiysh]

Thanks for the review @clementval. Addressed comment.

[peixin]

flang/lib/Lower/OpenMP.cpp

130

How about removing MLIR? OpenMPClauses or OmpClauses?

134

We used someClauseOperand(s) for other constructs in OpenMP.cpp in fir-dev. Can we keep consistent with others?

137

handleProcBindClause -> genProcBindClause to be consistent with other function names in this file.

[shraiysh]

Thanks @peixin for the comments. I will rework on this once upstreaming of translations has been done. My plan is to remove individual clause handling from genOMP functions and put it in a collect() function so all the constructs can use it. I will ping this thread once the upstreaming is done.

[peixin]

flang/test/Lower/OpenMP/parallel.f90

2 ↗

(On Diff #416169)

I am thinking if flang_fc1 -fopenmp -emit-fir should be added. Check the comment in https://reviews.llvm.org/D122595. Also you can check the discussions in https://reviews.llvm.org/D121171#3370137. If new passes is added in flang-new, I am not sure if it will crash the test cases.

16 ↗

(On Diff #416169)

Please add CHECK-LABEL for each subroutine.

[shraiysh]

Added handling for block constructs, sections constructs and critical construct.

[peixin]

flang/lib/Lower/OpenMP.cpp

207

As I mentioned in D124138, there is one drawback collecting clauses in one place for all block constructs. For those unsupported clauses, one hard TODO is changed into one unknown result, maybe one segfault. It is easy to omit some features support and get users into trouble if there is segfault/ICE.

Can we let each construct have their own collect? Here is my recomendation:

struct OpenMPMLIRClauses {
    void genNumThreadsClause(
      const Fortran::parser::OmpClause &clause,
      mlir::Value &mlirNumThreadsExpr)
     void genIFClause(...)
  ...
}

if (blockDirective.v == llvm::omp::OMPD_parallel) {
  for (const Fortran::parser::OmpClause &clause : clauseList.v) {
    genNumThreadsClause(...); ! if lowering to LLVM IR has been supported.
    genIfClause(...); ! if lowering to LLVM IR has been supported.
    TODO(... clause unsupported for parallel construct);
  }
  ...
} else if (blockDirective.v == llvm::omp::OMPD_master) {

As we pursue the OpenMP standard, it is easily know what are not supported yet. And this refector patch won't need to wait for other patches.
@shraiysh @kiranchandramohan @clementval What do you think?

[peixin]

@shraiysh Can you restart this work? There are two small patches left to upstream and it won't be affected too much.

To refactor the code, there are several things that I can think of to handle:

1. Explicit clauses for each construct. One unsupported clause should be given one TODO.
2. The clauses which are not allowed in directives should be given one assert?
3. The combined constructs are not carefully organized.

flang/lib/Lower/OpenMP.cpp

207

I take this back. This cause the code duplication.

[shraiysh]

Sure @peixin, thank you for pinging here. I will restart this work soon and try to incorporate your ideas and concerns.

[skatrak]

Hi @shraiysh, @peixin. I just realized about this patch, thanks to @kiranchandramohan, after I spent time working on a patch with basically the same goal of refactoring the clause processing in the PFT to MLIR lowering. The other patch is here: D155981.

It follows a slightly different approach where there is no collect() method or state to be kept, so the responsibility of deciding which clauses to generate code for is down to the caller. I am interested to see if you have any comments to improve that approach before landing it. It does other nonfunctional changes apart from extracting the clause-processing logic into a class, so the diff can be a bit noisy.

[shraiysh]

Thank you for working on this @skatrak. I will abandon this patch now.

flang/lib/Lower/OpenMP.cpp

createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
   firOpBuilder.create<mlir::omp::TerminatorOp>(loc);
   // Reset the insertion point to the start of the first block.
   firOpBuilder.setInsertionPointToStart(&block);
   // Handle privatization. Do not privatize if this is the outer operation.
   if (clauses && !outerCombined)
     privatizeVars(converter, *clauses);
 }
 
+struct OpenMPMLIRClauses {
+  mlir::Value ifExpr, numThreadsExpr;
+  mlir::omp::ClauseProcBindKindAttr procBindKind;
+  SmallVector<mlir::Value> allocatorOperands, allocateOperands;
+  mlir::UnitAttr nowait;
+  uint64_t hint;
+
+  OpenMPMLIRClauses(Fortran::lower::AbstractConverter &converter)
+      : converter(converter) {}
+
+  void genIfClause(const Fortran::parser::OmpClause::If &ifExpr,
+                   mlir::Value &mlirIfExpr) {
+    auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifExpr.v.t);
+    mlirIfExpr = fir::getBase(
+        converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
+  }
+
+  void genNumThreadsClause(
+      const Fortran::parser::OmpClause::NumThreads &numThreadsExpr,
+      mlir::Value &mlirNumThreadsExpr) {
+    mlirNumThreadsExpr = fir::getBase(converter.genExprValue(
+        *Fortran::semantics::GetExpr(numThreadsExpr.v), stmtCtx));
+  }
+
+  void
+  genProcBindClause(const Fortran::parser::OmpClause::ProcBind &procBindKind,
+                    mlir::omp::ClauseProcBindKindAttr &mlirProcBindKind) {
+    omp::ClauseProcBindKind pbKind;
+    switch (procBindKind.v.v) {
+    case Fortran::parser::OmpProcBindClause::Type::Master:
+      pbKind = omp::ClauseProcBindKind::Master;
+      break;
+    case Fortran::parser::OmpProcBindClause::Type::Close:
+      pbKind = omp::ClauseProcBindKind::Close;
+      break;
+    case Fortran::parser::OmpProcBindClause::Type::Spread:
+      pbKind = omp::ClauseProcBindKind::Spread;
+      break;
+    case Fortran::parser::OmpProcBindClause::Type::Primary:
+      pbKind = omp::ClauseProcBindKind::Primary;
+      break;
+    }
+    mlirProcBindKind = omp::ClauseProcBindKindAttr::get(
+        converter.getFirOpBuilder().getContext(), pbKind);
+  }
+
+  void genAllocateClause(
+      const Fortran::parser::OmpClause::Allocate &ompAllocateClause,
+      SmallVector<Value> &mlirAllocatorOperands,
+      SmallVector<Value> &mlirAllocateOperands) {
+    auto &firOpBuilder = converter.getFirOpBuilder();
+    auto currentLocation = converter.getCurrentLocation();
+    Fortran::lower::StatementContext stmtCtx;
+
+    mlir::Value allocatorOperand;
+    const Fortran::parser::OmpObjectList &ompObjectList =
+        std::get<Fortran::parser::OmpObjectList>(ompAllocateClause.v.t);
+    const auto &allocatorValue =
+        std::get<std::optional<Fortran::parser::OmpAllocateClause::Allocator>>(
+            ompAllocateClause.v.t);
+    // Check if allocate clause has allocator specified. If so, add it
+    // to list of allocators, otherwise, add default allocator to
+    // list of allocators.
+    if (allocatorValue) {
+      allocatorOperand = fir::getBase(converter.genExprValue(
+          *Fortran::semantics::GetExpr(allocatorValue->v), stmtCtx));
+      mlirAllocatorOperands.insert(mlirAllocatorOperands.end(),
+                                   ompObjectList.v.size(), allocatorOperand);
+    } else {
+      allocatorOperand = firOpBuilder.createIntegerConstant(
+          currentLocation, firOpBuilder.getI32Type(), 1);
+      mlirAllocatorOperands.insert(mlirAllocatorOperands.end(),
+                                   ompObjectList.v.size(), allocatorOperand);
+    }
+    genObjectList(ompObjectList, converter, mlirAllocateOperands);
+  }
+
+  void collect(const Fortran::parser::OmpClauseList &clauseList) {
+    for (const Fortran::parser::OmpClause &clause : clauseList.v) {
+      std::visit(
+          Fortran::common::visitors{
+              [this](const Fortran::parser::OmpClause::If &ifExpr) {
+                genIfClause(ifExpr, this->ifExpr);
+              },
+              [this](const Fortran::parser::OmpClause::NumThreads
+                         &numThreadsExpr) {
+                genNumThreadsClause(numThreadsExpr, this->numThreadsExpr);
+              },
+              [this](const Fortran::parser::OmpClause::ProcBind &procBindKind) {
+                genProcBindClause(procBindKind, this->procBindKind);
+              },
+              [this](
+                  const Fortran::parser::OmpClause::Allocate &allocateClause) {
+                genAllocateClause(allocateClause, this->allocatorOperands,
+                                  this->allocateOperands);
+              },
+              [this](const Fortran::parser::OmpClause::Nowait &) {
+                this->nowait = converter.getFirOpBuilder().getUnitAttr();
+              },
+              [this](const Fortran::parser::OmpClause::Hint &hint) {
+                const auto *expr = Fortran::semantics::GetExpr(hint.v);
+                this->hint = *Fortran::evaluate::ToInt64(*expr);
+              },
+              [](const Fortran::parser::OmpClause::Private &) {},
+              [](const Fortran::parser::OmpClause::Firstprivate &) {},
+              [](const Fortran::parser::OmpClause::Threads &) {},
+              [this](const auto &) {
+                TODO(converter.getCurrentLocation(), "Unhandled clause");
+              }},
+          clause.u);
+    }
+  }
+
+private:
+  Fortran::lower::StatementContext stmtCtx;
+  Fortran::lower::AbstractConverter &converter;
+};
+
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPSimpleStandaloneConstruct
                        &simpleStandaloneConstruct) {
   const auto &directive =
       std::get<Fortran::parser::OmpSimpleStandaloneDirective>(
           simpleStandaloneConstruct.t);
   switch (directive.v) {
static void genOMP(Fortran::lower::AbstractConverter &converter,
   case llvm::omp::Directive::OMPD_target_update:
     TODO(converter.getCurrentLocation(), "OMPD_target_update");
   case llvm::omp::Directive::OMPD_ordered:
     TODO(converter.getCurrentLocation(), "OMPD_ordered");
   }
 }
 
 static void
-genAllocateClause(Fortran::lower::AbstractConverter &converter,
-                  const Fortran::parser::OmpAllocateClause &ompAllocateClause,
-                  SmallVector<Value> &allocatorOperands,
-                  SmallVector<Value> &allocateOperands) {
-  auto &firOpBuilder = converter.getFirOpBuilder();
-  auto currentLocation = converter.getCurrentLocation();
-  Fortran::lower::StatementContext stmtCtx;
-
-  mlir::Value allocatorOperand;
-  const Fortran::parser::OmpObjectList &ompObjectList =
-      std::get<Fortran::parser::OmpObjectList>(ompAllocateClause.t);
-  const auto &allocatorValue =
-      std::get<std::optional<Fortran::parser::OmpAllocateClause::Allocator>>(
-          ompAllocateClause.t);
-  // Check if allocate clause has allocator specified. If so, add it
-  // to list of allocators, otherwise, add default allocator to
-  // list of allocators.
-  if (allocatorValue) {
-    allocatorOperand = fir::getBase(converter.genExprValue(
-        *Fortran::semantics::GetExpr(allocatorValue->v), stmtCtx));
-    allocatorOperands.insert(allocatorOperands.end(), ompObjectList.v.size(),
-                             allocatorOperand);
-  } else {
-    allocatorOperand = firOpBuilder.createIntegerConstant(
-        currentLocation, firOpBuilder.getI32Type(), 1);
-    allocatorOperands.insert(allocatorOperands.end(), ompObjectList.v.size(),
-                             allocatorOperand);
-  }
-  genObjectList(ompObjectList, converter, allocateOperands);
-}
-
-static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPStandaloneConstruct &standaloneConstruct) {
   std::visit(
       Fortran::common::visitors{
           [&](const Fortran::parser::OpenMPSimpleStandaloneConstruct
                   &simpleStandaloneConstruct) {
             genOMP(converter, eval, simpleStandaloneConstruct);
const auto &blockDirective =
       std::get<Fortran::parser::OmpBlockDirective>(beginBlockDirective.t);
   const auto &endBlockDirective =
       std::get<Fortran::parser::OmpEndBlockDirective>(blockConstruct.t);
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
 
   Fortran::lower::StatementContext stmtCtx;
   llvm::ArrayRef<mlir::Type> argTy;
-  mlir::Value ifClauseOperand, numThreadsClauseOperand;
-  mlir::omp::ClauseProcBindKindAttr procBindKindAttr;
-  SmallVector<Value> allocateOperands, allocatorOperands;
-  mlir::UnitAttr nowaitAttr;
 
   const auto &opClauseList =
       std::get<Fortran::parser::OmpClauseList>(beginBlockDirective.t);
-  for (const auto &clause : opClauseList.v) {
-    if (const auto &ifClause =
-            std::get_if<Fortran::parser::OmpClause::If>(&clause.u)) {
-      auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
-      ifClauseOperand = fir::getBase(
-          converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
-    } else if (const auto &numThreadsClause =
-                   std::get_if<Fortran::parser::OmpClause::NumThreads>(
-                       &clause.u)) {
-      // OMPIRBuilder expects `NUM_THREAD` clause as a `Value`.
-      numThreadsClauseOperand = fir::getBase(converter.genExprValue(
-          *Fortran::semantics::GetExpr(numThreadsClause->v), stmtCtx));
-    } else if (const auto &procBindClause =
-                   std::get_if<Fortran::parser::OmpClause::ProcBind>(
-                       &clause.u)) {
-      omp::ClauseProcBindKind pbKind;
-      switch (procBindClause->v.v) {
-      case Fortran::parser::OmpProcBindClause::Type::Master:
-        pbKind = omp::ClauseProcBindKind::Master;
-        break;
-      case Fortran::parser::OmpProcBindClause::Type::Close:
-        pbKind = omp::ClauseProcBindKind::Close;
-        break;
-      case Fortran::parser::OmpProcBindClause::Type::Spread:
-        pbKind = omp::ClauseProcBindKind::Spread;
-        break;
-      case Fortran::parser::OmpProcBindClause::Type::Primary:
-        pbKind = omp::ClauseProcBindKind::Primary;
-        break;
-      }
-      procBindKindAttr =
-          omp::ClauseProcBindKindAttr::get(firOpBuilder.getContext(), pbKind);
-    } else if (const auto &allocateClause =
-                   std::get_if<Fortran::parser::OmpClause::Allocate>(
-                       &clause.u)) {
-      genAllocateClause(converter, allocateClause->v, allocatorOperands,
-                        allocateOperands);
-    } else if (std::get_if<Fortran::parser::OmpClause::Private>(&clause.u) ||
-               std::get_if<Fortran::parser::OmpClause::Firstprivate>(
-                   &clause.u)) {
-      // Privatisation clauses are handled elsewhere.
-      continue;
-    } else if (std::get_if<Fortran::parser::OmpClause::Threads>(&clause.u)) {
-      // Nothing needs to be done for threads clause.
-      continue;
-    } else {
-      TODO(currentLocation, "OpenMP Block construct clauses");
-    }
-  }
 
-  for (const auto &clause :
-       std::get<Fortran::parser::OmpClauseList>(endBlockDirective.t).v) {
-    if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
-      nowaitAttr = firOpBuilder.getUnitAttr();
-  }
+  OpenMPMLIRClauses clauses(converter);
+  clauses.collect(opClauseList);
+  clauses.collect(
+      std::get<Fortran::parser::OmpClauseList>(endBlockDirective.t));
 
   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);
+        currentLocation, argTy, clauses.ifExpr, clauses.numThreadsExpr,
+        clauses.allocateOperands, clauses.allocatorOperands,
+        /*reduction_vars=*/ValueRange(),
+        /*reductions=*/nullptr, clauses.procBindKind);
     createBodyOfOp<omp::ParallelOp>(parallelOp, converter, currentLocation,
                                     &opClauseList, /*isCombined=*/false);
   } else if (blockDirective.v == llvm::omp::OMPD_master) {
     auto masterOp =
         firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
     createBodyOfOp<omp::MasterOp>(masterOp, converter, currentLocation);
   } else if (blockDirective.v == llvm::omp::OMPD_single) {
     auto singleOp = firOpBuilder.create<mlir::omp::SingleOp>(
-        currentLocation, allocateOperands, allocatorOperands, nowaitAttr);
+        currentLocation, clauses.allocateOperands, clauses.allocatorOperands,
+        clauses.nowait);
     createBodyOfOp<omp::SingleOp>(singleOp, converter, currentLocation);
   } else if (blockDirective.v == llvm::omp::OMPD_ordered) {
     auto orderedOp = firOpBuilder.create<mlir::omp::OrderedRegionOp>(
         currentLocation, /*simd=*/nullptr);
     createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   std::string name;
   const Fortran::parser::OmpCriticalDirective &cd =
       std::get<Fortran::parser::OmpCriticalDirective>(criticalConstruct.t);
   if (std::get<std::optional<Fortran::parser::Name>>(cd.t).has_value()) {
     name =
         std::get<std::optional<Fortran::parser::Name>>(cd.t).value().ToString();
   }
 
-  uint64_t hint = 0;
   const auto &clauseList = std::get<Fortran::parser::OmpClauseList>(cd.t);
-  for (const Fortran::parser::OmpClause &clause : clauseList.v)
-    if (auto hintClause =
-            std::get_if<Fortran::parser::OmpClause::Hint>(&clause.u)) {
-      const auto *expr = Fortran::semantics::GetExpr(hintClause->v);
-      hint = *Fortran::evaluate::ToInt64(*expr);
-      break;
-    }
+  OpenMPMLIRClauses clauses(converter);
+  clauses.collect(clauseList);
 
   mlir::omp::CriticalOp criticalOp = [&]() {
     if (name.empty()) {
       return firOpBuilder.create<mlir::omp::CriticalOp>(currentLocation,
                                                         FlatSymbolRefAttr());
     } else {
       mlir::ModuleOp module = firOpBuilder.getModule();
       mlir::OpBuilder modBuilder(module.getBodyRegion());
       auto global = module.lookupSymbol<mlir::omp::CriticalDeclareOp>(name);
       if (!global)
         global = modBuilder.create<mlir::omp::CriticalDeclareOp>(
-            currentLocation, name, hint);
+            currentLocation, name, clauses.hint);
       return firOpBuilder.create<mlir::omp::CriticalOp>(
           currentLocation, mlir::FlatSymbolRefAttr::get(
                                firOpBuilder.getContext(), global.sym_name()));
     }
   }();
   createBodyOfOp<omp::CriticalOp>(criticalOp, converter, currentLocation);
 }
 

 
 // 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>(
       std::get<Fortran::parser::OmpBeginSectionsDirective>(sectionsConstruct.t)
           .t);
-  for (const Fortran::parser::OmpClause &clause : sectionsClauseList.v) {
-
-    // Reduction Clause
-    if (std::get_if<Fortran::parser::OmpClause::Reduction>(&clause.u)) {
-      TODO(currentLocation, "OMPC_Reduction");
-
-      // Allocate clause
-    } else if (const auto &allocateClause =
-                   std::get_if<Fortran::parser::OmpClause::Allocate>(
-                       &clause.u)) {
-      genAllocateClause(converter, allocateClause->v, allocatorOperands,
-                        allocateOperands);
-    }
-  }
-  const auto &endSectionsClauseList =
-      std::get<Fortran::parser::OmpEndSectionsDirective>(sectionsConstruct.t);
-  const auto &clauseList =
-      std::get<Fortran::parser::OmpClauseList>(endSectionsClauseList.t);
-  for (const auto &clause : clauseList.v) {
-    // Nowait clause
-    if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u)) {
-      noWaitClauseOperand = firOpBuilder.getUnitAttr();
-    }
-  }
+  const auto &endSectionsClauseList = std::get<Fortran::parser::OmpClauseList>(
+      std::get<Fortran::parser::OmpEndSectionsDirective>(sectionsConstruct.t)
+          .t);
+  OpenMPMLIRClauses clauses(converter);
+  clauses.collect(sectionsClauseList);
+  clauses.collect(endSectionsClauseList);
 
   llvm::omp::Directive dir =
       std::get<Fortran::parser::OmpSectionsDirective>(
           std::get<Fortran::parser::OmpBeginSectionsDirective>(
               sectionsConstruct.t)
               .t)
           .v;
 
   // Parallel Sections Construct
   if (dir == llvm::omp::Directive::OMPD_parallel_sections) {
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
         currentLocation, /*if_expr_var*/ nullptr, /*num_threads_var*/ nullptr,
-        allocateOperands, allocatorOperands, /*reduction_vars=*/ValueRange(),
+        clauses.allocateOperands, clauses.allocatorOperands,
+        /*reduction_vars=*/ValueRange(),
         /*reductions=*/nullptr, /*proc_bind_val*/ nullptr);
     createBodyOfOp(parallelOp, converter, currentLocation);
     auto sectionsOp = firOpBuilder.create<mlir::omp::SectionsOp>(
         currentLocation, /*reduction_vars*/ ValueRange(),
         /*reductions=*/nullptr, /*allocate_vars*/ ValueRange(),
         /*allocators_vars*/ ValueRange(), /*nowait=*/nullptr);
     createBodyOfOp(sectionsOp, converter, currentLocation);
 
     // Sections Construct
   } else if (dir == llvm::omp::Directive::OMPD_sections) {
     auto sectionsOp = firOpBuilder.create<mlir::omp::SectionsOp>(
-        currentLocation, reductionVars, /*reductions = */ nullptr,
-        allocateOperands, allocatorOperands, noWaitClauseOperand);
+        currentLocation, ValueRange(), /*reductions = */ nullptr,
+        clauses.allocateOperands, clauses.allocatorOperands, clauses.nowait);
     createBodyOfOp<omp::SectionsOp>(sectionsOp, converter, currentLocation);
   }
 }
 
 void Fortran::lower::genOpenMPConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenMPConstruct &ompConstruct) {