[flang][OpenMP] Added allocate clause translation for OpenMP block constructs

[shraiysh]

Summary

This patch adds translation for allocate clause for parallel and single
constructs.

Also added tests for block constructs.

This patch also adds tests for parallel construct which were not added earlier.

Co-authored-by: Sourabh Singh Tomar <[e-mail]>

[shraiysh]

Handle pre-merge check failure.

[shraiysh]

This comment was removed by shraiysh.

[NimishMishra]

LGTM. Thanks for this. A couple of small nits, please let me know what you think.

flang/lib/Lower/OpenMP.cpp

192

Nit: ProcBindKind has 6 possibilities in total. Although I strongly believe checks would be there, but can you confirm if semantic checks are in place that guard against the other three (Primary, Default, and Unknown) from trickling down till here?

flang/test/Lower/OpenMP/parallel.f90

67

Nit: You also add tests for num_threads, if, proc_bind. Can you edit the patch description to this effect? I think that would make it complete, since we aren't just looking at tests related to allocate clause here.

[shraiysh]

Addressed comments.

[shraiysh]

I will merge this tomorrow if there are no further comments.

[peixin]

One more nit for the title: Added allocate clause translation for block constructs -> Added allocate clause translation for OpenMP block constructs since fortran also has one block construct.

Sorry for the late review.

flang/lib/Lower/OpenMP.cpp

188

This if clause comment should be in line192.

195

This num_threads clause comment should be in 200. Or don't need it since the code is prettry clear.

233

Add one else branch with hard TODO for these clauses.

242

This comment should be in line 243 or no need it. I think the code is prettry clear.

251–252

Remove lines 255 and 256. The same for others.

[peixin]

flang/test/Lower/OpenMP/parallel.f90

3

The same for flang_fc1 -emit-fir -fopenmp.

14

Nit: Please add CHECK-LABEL for each subroutine.

[shraiysh]

Addressed comments.

[shraiysh]

flang/test/Lower/OpenMP/parallel.f90

3

I didn't understand this. Can you please elaborate it?

[peixin]

flang/test/Lower/OpenMP/parallel.f90

3

Please forget about this. Kiran has explained to me. There is no much difference using flang-new or bbc for us to test openmp code.

[shraiysh]

flang/test/Lower/OpenMP/parallel.f90

3

Okay, thank you. I will land this.

flang/include/flang/Parser/parse-tree.h

 };
 
 WRAPPER_CLASS(PauseStmt, std::optional<StopCode>);
 
 // Parse tree nodes for OpenMP 4.5 directives and clauses
 
 // 2.5 proc-bind-clause -> PROC_BIND (MASTER | CLOSE | SPREAD)
 struct OmpProcBindClause {
-  ENUM_CLASS(Type, Close, Master, Spread)
+  ENUM_CLASS(Type, Close, Master, Spread, Primary)
   WRAPPER_CLASS_BOILERPLATE(OmpProcBindClause, Type);
 };
 
 // 2.15.3.1 default-clause -> DEFAULT (PRIVATE | FIRSTPRIVATE | SHARED | NONE)
 struct OmpDefaultClause {
   ENUM_CLASS(Type, Private, Firstprivate, Shared, None)
   WRAPPER_CLASS_BOILERPLATE(OmpDefaultClause, Type);
 };

flang/lib/Lower/OpenMP.cpp

const auto &blockDirective =
       std::get<Fortran::parser::OmpBlockDirective>(beginBlockDirective.t);
   const auto &endBlockDirective =
       std::get<Fortran::parser::OmpEndBlockDirective>(blockConstruct.t);
 
   auto &firOpBuilder = converter.getFirOpBuilder();
   auto currentLocation = converter.getCurrentLocation();
   Fortran::lower::StatementContext stmtCtx;
   llvm::ArrayRef<mlir::Type> argTy;
-  if (blockDirective.v == llvm::omp::OMPD_parallel) {
-
-    mlir::Value ifClauseOperand, numThreadsClauseOperand;
-    Attribute procBindClauseOperand;
+  mlir::omp::ClauseProcBindKindAttr procBindKindAttr;
+  SmallVector<Value> allocateOperands, allocatorOperands;
+  mlir::UnitAttr nowaitAttr;
 
-    const auto &parallelOpClauseList =
-        std::get<Fortran::parser::OmpClauseList>(beginBlockDirective.t);
-    for (const auto &clause : parallelOpClauseList.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));
-      }
-      // TODO: Handle private, firstprivate, shared and copyin
-    }
-    // Create and insert the operation.
-    auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
-        currentLocation, argTy, ifClauseOperand, numThreadsClauseOperand,
-        /*allocate_vars=*/ValueRange(), /*allocators_vars=*/ValueRange(),
-        /*reduction_vars=*/ValueRange(), /*reductions=*/nullptr,
-        procBindClauseOperand.dyn_cast_or_null<omp::ClauseProcBindKindAttr>());
-    // Handle attribute based clauses.
-    for (const auto &clause : parallelOpClauseList.v) {
-      // TODO: Handle default clause
-      if (const auto &procBindClause =
-              std::get_if<Fortran::parser::OmpClause::ProcBind>(&clause.u)) {
-        const auto &ompProcBindClause{procBindClause->v};
-        omp::ClauseProcBindKind pbKind;
-        switch (ompProcBindClause.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;
-        }
-        parallelOp.proc_bind_valAttr(omp::ClauseProcBindKindAttr::get(
-            firOpBuilder.getContext(), pbKind));
-      }
-    }
-    createBodyOfOp<omp::ParallelOp>(parallelOp, firOpBuilder, currentLocation);
-  } else if (blockDirective.v == llvm::omp::OMPD_master) {
-    auto masterOp =
-        firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
-    createBodyOfOp<omp::MasterOp>(masterOp, firOpBuilder, currentLocation);
+  for (const auto &clause :
+       std::get<Fortran::parser::OmpClauseList>(beginBlockDirective.t).v) {
+      auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
+      ifClauseOperand = fir::getBase(
+          converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
+    } else if (const auto &procBindClause =
+                   std::get_if<Fortran::parser::OmpClause::ProcBind>(
+                       &clause.u)) {
+      switch (procBindClause->v.v) {
+      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 (const auto &privateClause =
+                   std::get_if<Fortran::parser::OmpClause::Private>(
+                       &clause.u)) {
+      // TODO: Handle private. This cannot be a hard TODO because testing for
+      // allocate clause requires private variables.
+    } else {
+      TODO(currentLocation, "OpenMP Block construct clauses");
 
-    // Single Construct
-  } else if (blockDirective.v == llvm::omp::OMPD_single) {
-    mlir::UnitAttr nowaitAttr;
-    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();
-      // TODO: Handle allocate clause (D122302)
-    }
+
+  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);
+    createBodyOfOp<omp::ParallelOp>(parallelOp, firOpBuilder, currentLocation);
+  } else if (blockDirective.v == llvm::omp::OMPD_master) {
+    auto masterOp =
+        firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
+    createBodyOfOp<omp::MasterOp>(masterOp, firOpBuilder, currentLocation);
+  } else if (blockDirective.v == llvm::omp::OMPD_single) {
     auto singleOp = firOpBuilder.create<mlir::omp::SingleOp>(
-        currentLocation, /*allocate_vars=*/ValueRange(),
-        /*allocators_vars=*/ValueRange(), nowaitAttr);
+        currentLocation, allocateOperands, allocatorOperands, nowaitAttr);
     createBodyOfOp(singleOp, firOpBuilder, currentLocation);
   }
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {

flang/test/Lower/OpenMP/parallel.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s --check-prefixes="FIRDialect,OMPDialect"
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="OMPDialect"
+
+!FIRDialect-LABEL: func @_QPparallel_simple
+subroutine parallel_simple()
+   !OMPDialect: omp.parallel
+!$omp parallel
+   !FIRDialect: fir.call
+   call f1()
+!$omp end parallel
+end subroutine parallel_simple
+
+!===============================================================================
+! `if` clause
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPparallel_if
+subroutine parallel_if(alpha)
+   integer, intent(in) :: alpha
+
+   !OMPDialect: omp.parallel if(%{{.*}} : i1) {
+   !$omp parallel if(alpha .le. 0)
+   !FIRDialect: fir.call
+   call f1()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel if(%{{.*}} : i1) {
+   !$omp parallel if(.false.)
+   !FIRDialect: fir.call
+   call f2()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel if(%{{.*}} : i1) {
+   !$omp parallel if(alpha .ge. 0)
+   !FIRDialect: fir.call
+   call f3()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel if(%{{.*}} : i1) {
+   !$omp parallel if(.true.)
+   !FIRDialect: fir.call
+   call f4()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+end subroutine parallel_if
+
+!===============================================================================
+! `num_threads` clause
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPparallel_numthreads
+subroutine parallel_numthreads(num_threads)
+   integer, intent(inout) :: num_threads
+
+   !OMPDialect: omp.parallel num_threads(%{{.*}}: i32) {
+   !$omp parallel num_threads(16)
+   !FIRDialect: fir.call
+   call f1()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   num_threads = 4
+
+   !OMPDialect: omp.parallel num_threads(%{{.*}} : i32) {
+   !$omp parallel num_threads(num_threads)
+   !FIRDialect: fir.call
+   call f2()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+end subroutine parallel_numthreads
+
+!===============================================================================
+! `proc_bind` clause
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPparallel_proc_bind
+subroutine parallel_proc_bind()
+
+   !OMPDialect: omp.parallel proc_bind(master) {
+   !$omp parallel proc_bind(master)
+   !FIRDialect: fir.call
+   call f1()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel proc_bind(close) {
+   !$omp parallel proc_bind(close)
+   !FIRDialect: fir.call
+   call f2()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel proc_bind(spread) {
+   !$omp parallel proc_bind(spread)
+   !FIRDialect: fir.call
+   call f3()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+end subroutine parallel_proc_bind
+
+!===============================================================================
+! `allocate` clause
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPparallel_allocate
+subroutine parallel_allocate()
+   use omp_lib
+   integer :: x
+   !OMPDialect: omp.parallel allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+   !$omp parallel allocate(omp_high_bw_mem_alloc: x) private(x)
+   !FIRDialect: arith.addi
+   x = x + 12
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+end subroutine parallel_allocate
+
+!===============================================================================
+! multiple clauses
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPparallel_multiple_clauses
+subroutine parallel_multiple_clauses(alpha, num_threads)
+   use omp_lib
+   integer, intent(inout) :: alpha
+   integer, intent(in) :: num_threads
+
+   !OMPDialect: omp.parallel if({{.*}} : i1) proc_bind(master) {
+   !$omp parallel if(alpha .le. 0) proc_bind(master)
+   !FIRDialect: fir.call
+   call f1()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel num_threads({{.*}} : i32) proc_bind(close) {
+   !$omp parallel proc_bind(close) num_threads(num_threads)
+   !FIRDialect: fir.call
+   call f2()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel if({{.*}} : i1) num_threads({{.*}} : i32) {
+   !$omp parallel num_threads(num_threads) if(alpha .le. 0)
+   !FIRDialect: fir.call
+   call f3()
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+   !OMPDialect: omp.parallel if({{.*}} : i1) num_threads({{.*}} : i32) allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+   !$omp parallel num_threads(num_threads) if(alpha .le. 0) allocate(omp_high_bw_mem_alloc: alpha) private(alpha)
+   !FIRDialect: fir.call
+   call f3()
+   !FIRDialect: arith.addi
+   alpha = alpha + 12
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+
+end subroutine parallel_multiple_clauses

flang/test/Lower/OpenMP/single.f90

!$omp single
     !FIRDialect: %[[res:.*]] = arith.addi %[[xval]], %{{.*}} : i32
     !FIRDialect: fir.store %[[res]] to %[[x]] : !fir.ref<i32>
     x = x + 12
   !OMPDialect: omp.terminator
   !$omp end single nowait
   !OMPDialect: omp.terminator
   !$omp end parallel
 end subroutine omp_single_nowait
+
+!===============================================================================
+! Single construct with allocate
+!===============================================================================
+
+!FIRDialect-LABEL: func @_QPsingle_allocate
+subroutine single_allocate()
+  use omp_lib
+  integer :: x
+  !OMPDialect: omp.parallel {
+  !$omp parallel
+  !OMPDialect: omp.single allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+  !$omp single allocate(omp_high_bw_mem_alloc: x) private(x)
+  !FIRDialect: arith.addi
+  x = x + 12
+  !OMPDialect: omp.terminator
+  !$omp end single
+  !OMPDialect: omp.terminator
+  !$omp end parallel
+end subroutine single_allocate