[Flang][OpenMP] Upstream the lowering of the parallel do combined construct

[kiranchandramohan]

Summary

When parallel is used in a combined construct, then use a separate
function to create the parallel operation. It handles the parallel
specific clauses and leaves the rest for handling at the inner
operations.

Co-authored-by: Sourabh Singh Tomar <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: Valentin Clement <[e-mail]>
Co-authored-by: Nimish Mishra <[e-mail]>

[peixin]

LGTM except fow a few nits.

flang/lib/Lower/OpenMP.cpp

270

467

Nit: SIMD/Taskloop/Loop and combined constructs will be also implemented here.

flang/test/Lower/OpenMP/omp-parallel-wsloop.f90

3

One more regression test case for private and firstprivate with one simple data type support?

[shraiysh]

How would this affect parallel sections? Should that be changed to use the new function createCombinedParallelOp too? If not, can we change the comment to state that this is for "loop constructs" only?

[kiranchandramohan]

Changes to address review comments.
-> Call the createCombinedParallelOp for Parallel sections too.
-> Add tests for privatisation clauses.
-> Modified the TODO

[kiranchandramohan]

In D125465#3512551, @shraiysh wrote:

How would this affect parallel sections? Should that be changed to use the new function createCombinedParallelOp too? If not, can we change the comment to state that this is for "loop constructs" only?

Modified the parallel sections code to use createCombinedParallelOp.

flang/lib/Lower/OpenMP.cpp

270

These two clauses are also present in inner constructs like worksharing loop, sections. Hence I left them out.

[peixin]

LGTM

flang/lib/Lower/OpenMP.cpp

270

Got it. Thanks.

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

82 ↗

(On Diff #430319)

flang/test/Lower/OpenMP/omp-parallel-wsloop.f90

13

Usually, I think it's better to put these filecheck before the subroutine (line 6). These filecheck and omp directives are both grey color. It's kind of hard to read. But it's also OK to me to not change this.

[shraiysh]

LGTM thank you. There is a bit of code duplication with procbind clause, can we maybe extract that in a genProcBindAttr(Fortran::parser::OmpClause::ProcBind &, mlir::omp::ClauseProcBindKindAttr&) function?

[kiranchandramohan]

Refactor genProcBindKindAttr.

[shraiysh]

flang/lib/Lower/OpenMP.cpp

315

allocateOperands and allocatorOperands are not set. They should probably use genAllocateClause. Or is it not required for parallel constructs?

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

315

They are left for handling at the inner construct level.

[kiranchandramohan]

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

82 ↗

(On Diff #430319)

This is the input, so no patterns here. Hence I chose not to address this comment.

flang/test/Lower/OpenMP/omp-parallel-wsloop.f90

13

At some point we can standardise the way we write. Chose to skip this for now.

[peixin]

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

82 ↗

(On Diff #430319)

Sorry. I misread this.

flang/test/Lower/OpenMP/omp-parallel-wsloop.f90

13

Sure. I like to push forward the upstream work. Thanks for landing this.

flang/lib/Lower/OpenMP.cpp

 //===-- OpenMP.cpp -- Open MP directive lowering --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
std::visit(
           [&](const Fortran::parser::OpenMPCancellationPointConstruct
                   &cancellationPointConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPCancelConstruct");
           },
       },
       standaloneConstruct.u);
 }
 
+/* When parallel is used in a combined construct, then use this function to
+ * create the parallel operation. It handles the parallel specific clauses
+ * and leaves the rest for handling at the inner operations.
+ * TODO: Refactor clause handling
+ */
+template <typename Directive>
+static void
+createCombinedParallelOp(Fortran::lower::AbstractConverter &converter,
+                         Fortran::lower::pft::Evaluation &eval,
+                         const Directive &directive) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+  Fortran::lower::StatementContext stmtCtx;
+  llvm::ArrayRef<mlir::Type> argTy;
+  mlir::Value ifClauseOperand, numThreadsClauseOperand;
+  SmallVector<Value> allocatorOperands, allocateOperands;
+  mlir::omp::ClauseProcBindKindAttr procBindKindAttr;
+  const auto &opClauseList =
+      std::get<Fortran::parser::OmpClauseList>(directive.t);
+  // TODO: Handle the following clauses
+  // 1. default
+  // 2. copyin
+  // Note: rest of the clauses are handled when the inner operation is created
+  for (const Fortran::parser::OmpClause &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);
+      mlir::Value ifVal = fir::getBase(
+          converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
+      ifClauseOperand = firOpBuilder.createConvert(
+          currentLocation, firOpBuilder.getI1Type(), ifVal);
+    } else if (const auto &numThreadsClause =
+                   std::get_if<Fortran::parser::OmpClause::NumThreads>(
+                       &clause.u)) {
+      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);
+    }
+  }
+  // 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, /*iv=*/nullptr,
+                                  /*isCombined=*/true);
+}
+
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPBlockConstruct &blockConstruct) {
   const auto &beginBlockDirective =
       std::get<Fortran::parser::OmpBeginBlockDirective>(blockConstruct.t);
   const auto &blockDirective =
       std::get<Fortran::parser::OmpBlockDirective>(beginBlockDirective.t);
static void genOMP(Fortran::lower::AbstractConverter &converter,
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
   mlir::Value scheduleChunkClauseOperand;
   mlir::Attribute scheduleClauseOperand, collapseClauseOperand,
       noWaitClauseOperand, orderedClauseOperand, orderClauseOperand;
-  const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
-      std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
-  if (llvm::omp::OMPD_do !=
+
+  const auto ompDirective =
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
-          .v) {
+          .v;
+  if (llvm::omp::OMPD_parallel_do == ompDirective) {
+    createCombinedParallelOp<Fortran::parser::OmpBeginLoopDirective>(
+        converter, eval,
+        std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t));
+  } else if (llvm::omp::OMPD_do != ompDirective) {
     TODO(converter.getCurrentLocation(), "Combined worksharing loop construct");
   }
 
   Fortran::lower::pft::Evaluation *doConstructEval =
       &eval.getFirstNestedEvaluation();
 
   Fortran::lower::pft::Evaluation *doLoop =
       &doConstructEval->getFirstNestedEvaluation();
   auto *doStmt = doLoop->getIf<Fortran::parser::NonLabelDoStmt>();
auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       scheduleChunkClauseOperand, /*schedule_modifiers=*/nullptr,
       /*simd_modifier=*/nullptr,
       collapseClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       noWaitClauseOperand.dyn_cast_or_null<UnitAttr>(),
       orderedClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
       /*inclusive=*/firOpBuilder.getUnitAttr());
 
+  const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
+      std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
   // Handle attribute based clauses.
   for (const Fortran::parser::OmpClause &clause : wsLoopOpClauseList.v) {
     if (const auto &scheduleClause =
             std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
       mlir::MLIRContext *context = firOpBuilder.getContext();
       const auto &scheduleType = scheduleClause->v;
       const auto &scheduleKind =
           std::get<Fortran::parser::OmpScheduleClause::ScheduleType>(

flang/test/Lower/OpenMP/omp-parallel-wsloop.f90

+! This test checks lowering of OpenMP DO Directive (Worksharing).
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+! CHECK-LABEL: func @_QPsimple_parallel_do()
+subroutine simple_parallel_do
+  integer :: i
+  ! CHECK:  omp.parallel
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+  ! CHECK:     omp.wsloop for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP PARALLEL DO
+  do i=1, 9
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL DO
+end subroutine
+
+! CHECK-LABEL: func @_QPparallel_do_with_parallel_clauses
+! CHECK-SAME: %[[COND_REF:.*]]: !fir.ref<!fir.logical<4>> {fir.bindc_name = "cond"}, %[[NT_REF:.*]]: !fir.ref<i32> {fir.bindc_name = "nt"}
+subroutine parallel_do_with_parallel_clauses(cond, nt)
+  logical :: cond
+  integer :: nt
+  integer :: i
+  ! CHECK:  %[[COND:.*]] = fir.load %[[COND_REF]] : !fir.ref<!fir.logical<4>>
+  ! CHECK:  %[[COND_CVT:.*]] = fir.convert %[[COND]] : (!fir.logical<4>) -> i1
+  ! CHECK:  %[[NT:.*]] = fir.load %[[NT_REF]] : !fir.ref<i32>
+  ! CHECK:  omp.parallel if(%[[COND_CVT]] : i1) num_threads(%[[NT]] : i32) proc_bind(close)
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+  ! CHECK:     omp.wsloop for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP PARALLEL DO IF(cond) NUM_THREADS(nt) PROC_BIND(close)
+  do i=1, 9
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL DO
+end subroutine
+
+! CHECK-LABEL: func @_QPparallel_do_with_clauses
+! CHECK-SAME: %[[NT_REF:.*]]: !fir.ref<i32> {fir.bindc_name = "nt"}
+subroutine parallel_do_with_clauses(nt)
+  integer :: nt
+  integer :: i
+  ! CHECK:  %[[NT:.*]] = fir.load %[[NT_REF]] : !fir.ref<i32>
+  ! CHECK:  omp.parallel num_threads(%[[NT]] : i32)
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+  ! CHECK:     omp.wsloop schedule(dynamic) for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP PARALLEL DO NUM_THREADS(nt) SCHEDULE(dynamic)
+  do i=1, 9
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL DO
+end subroutine