[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

220

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

189

This if clause comment should be in line192.

196

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

223

Add one else branch with hard TODO for these clauses.

232

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

246

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/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
-    }
+  for (const auto &clause :
+       std::get<Fortran::parser::OmpClauseList>(beginBlockDirective.t).v) {
+    // `if` clause
+      auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
+      ifClauseOperand = fir::getBase(
+          converter.genExprValue(*Fortran::semantics::GetExpr(expr), stmtCtx));
+
+      // `num_threads` clause
+
+      // `proc_bind` Clause
+    } else if (const auto &procBindClause =
+                   std::get_if<Fortran::parser::OmpClause::ProcBind>(
+                       &clause.u)) {
+      const auto &ompProcBindClause{procBindClause->v};
+      auto pbKind = mlir::omp::symbolizeClauseProcBindKind(
+          Fortran::parser::OmpProcBindClause::EnumToString(
+              procBindClause->v.v));
+      assert(pbKind.hasValue() && "unable to translate proc_bind kind from PFT to MLIR");
+      procBindKindAttr =
+          omp::ClauseProcBindKindAttr::get(firOpBuilder.getContext(), *pbKind);
+
+      // `allocate` clause
+    } else if (const auto &allocateClause =
+                   std::get_if<Fortran::parser::OmpClause::Allocate>(
+                       &clause.u)) {
+      genAllocateClause(converter, allocateClause->v, allocatorOperands,
+                        allocateOperands);
+
+  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();
+  }
+
+  // Parallel Construct
+  if (blockDirective.v == llvm::omp::OMPD_parallel) {
     // Create and insert the operation.
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
         currentLocation, argTy, ifClauseOperand, numThreadsClauseOperand,
-        ValueRange(), ValueRange(),
-        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));
-      }
-    }
+        allocateOperands, allocatorOperands, procBindKindAttr);
     createBodyOfOp<omp::ParallelOp>(parallelOp, firOpBuilder, currentLocation);
+
+    // Master Construct
   } else if (blockDirective.v == llvm::omp::OMPD_master) {
     auto masterOp =
         firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
     createBodyOfOp<omp::MasterOp>(masterOp, firOpBuilder, currentLocation);
 
     // 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)
-    }
     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 --cfg-conversion | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="OMPDialect"
+
+subroutine parallel_simple()
+   !OMPDialect: omp.parallel
+!$omp parallel
+   !FIRDialect: fir.call
+   call f1()
+!$omp end parallel
+end subroutine parallel_simple
+
+!===============================================================================
+! `if` clause
+!===============================================================================
+
+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
+!===============================================================================
+
+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
+!===============================================================================
+
+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
+!===============================================================================
+
+subroutine parallel_allocate()
+   integer :: x
+   !OMPDialect: omp.parallel allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+   !$omp parallel allocate(omp_high_bw_mem_alloc: x)
+   !FIRDialect: arith.addi
+   x = x + 12
+   !OMPDialect: omp.terminator
+   !$omp end parallel
+end subroutine parallel_allocate
+
+!===============================================================================
+! multiple clauses
+!===============================================================================
+
+subroutine parallel_multiple_clauses(alpha, num_threads)
+   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)
+   !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
+!===============================================================================
+subroutine single_allocate()
+  integer :: x
+  !OMPDialect: omp.parallel {
+  !$omp parallel
+  !OMPDialect: omp.single allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+  !$omp single allocate(omp_high_bw_mem_alloc: x)
+  !FIRDialect: arith.addi
+  x = x + 12
+  !OMPDialect: omp.terminator
+  !$omp end single
+  !OMPDialect: omp.terminator
+  !$omp end parallel
+end subroutine single_allocate