[Flang] Syntax support for OMP Allocators Construct

OpenMP 5.2 introduces a Fortran specific construct that aims to replace the executable allocate directive.  This patch seeks to add parser support for the directive as well as the allocator clause with the [[ https://www.openmp.org/wp-content/uploads/OpenMP-API-Specification-5-2.pdf#section.6.6 | extended align/complex ]] modifier.

Reviewed By: kiranchandramohan

Differential Revision: https://reviews.llvm.org/D149727

8 files changed, 218 insertions, 14 deletions

diff --git a/flang/examples/FeatureList/FeatureList.cpp b/flang/examples/FeatureList/FeatureList.cpp
index 5efda806856f..a93c1b19ea68 100644
--- a/flang/examples/FeatureList/FeatureList.cpp
+++ b/flang/examples/FeatureList/FeatureList.cpp
@@ -474,6 +474,7 @@ public:
   READ_FEATURE(OmpDependenceType::Type)
   READ_FEATURE(OmpDependSinkVec)
   READ_FEATURE(OmpDependSinkVecLength)
+  READ_FEATURE(OmpEndAllocators)
   READ_FEATURE(OmpEndAtomic)
   READ_FEATURE(OmpEndBlockDirective)
   READ_FEATURE(OmpEndCriticalDirective)
@@ -506,7 +507,10 @@ public:
   READ_FEATURE(OmpReductionInitializerClause)
   READ_FEATURE(OmpReductionOperator)
   READ_FEATURE(OmpAllocateClause)
-  READ_FEATURE(OmpAllocateClause::Allocator)
+  READ_FEATURE(OmpAllocateClause::AllocateModifier)
+  READ_FEATURE(OmpAllocateClause::AllocateModifier::Allocator)
+  READ_FEATURE(OmpAllocateClause::AllocateModifier::ComplexModifier)
+  READ_FEATURE(OmpAllocateClause::AllocateModifier::Align)
   READ_FEATURE(OmpScheduleClause)
   READ_FEATURE(OmpScheduleClause::ScheduleType)
   READ_FEATURE(OmpDeviceClause)
@@ -553,6 +557,7 @@ public:
   READ_FEATURE(OpenMPFlushConstruct)
   READ_FEATURE(OpenMPLoopConstruct)
   READ_FEATURE(OpenMPExecutableAllocate)
+  READ_FEATURE(OpenMPAllocatorsConstruct)
   READ_FEATURE(OpenMPRequiresConstruct)
   READ_FEATURE(OpenMPSimpleStandaloneConstruct)
   READ_FEATURE(OpenMPStandaloneConstruct)
diff --git a/flang/examples/FlangOmpReport/FlangOmpReportVisitor.cpp b/flang/examples/FlangOmpReport/FlangOmpReportVisitor.cpp
index d78f898d499e..7382e99de1ab 100644
--- a/flang/examples/FlangOmpReport/FlangOmpReportVisitor.cpp
+++ b/flang/examples/FlangOmpReport/FlangOmpReportVisitor.cpp
@@ -128,6 +128,10 @@ std::string OpenMPCounterVisitor::getName(const OpenMPConstruct &c) {
             const CharBlock &source{std::get<0>(c.t).source};
             return normalize_construct_name(source.ToString());
           },
+          [&](const OpenMPAllocatorsConstruct &c) -> std::string {
+            const CharBlock &source{std::get<0>(c.t).source};
+            return normalize_construct_name(source.ToString());
+          },
           [&](const OpenMPAtomicConstruct &c) -> std::string {
             return std::visit(
                 [&](const auto &c) {
diff --git a/flang/include/flang/Parser/dump-parse-tree.h b/flang/include/flang/Parser/dump-parse-tree.h
index 7bc1d08c1d5a..5b680ea3403d 100644
--- a/flang/include/flang/Parser/dump-parse-tree.h
+++ b/flang/include/flang/Parser/dump-parse-tree.h
@@ -490,6 +490,7 @@ public:
   NODE_ENUM(OmpDependenceType, Type)
   NODE(parser, OmpDependSinkVec)
   NODE(parser, OmpDependSinkVecLength)
+  NODE(parser, OmpEndAllocators)
   NODE(parser, OmpEndAtomic)
   NODE(parser, OmpEndBlockDirective)
   NODE(parser, OmpEndCriticalDirective)
@@ -527,7 +528,10 @@ public:
   NODE(parser, OmpReductionInitializerClause)
   NODE(parser, OmpReductionOperator)
   NODE(parser, OmpAllocateClause)
-  NODE(OmpAllocateClause, Allocator)
+  NODE(OmpAllocateClause, AllocateModifier)
+  NODE(OmpAllocateClause::AllocateModifier, Allocator)
+  NODE(OmpAllocateClause::AllocateModifier, ComplexModifier)
+  NODE(OmpAllocateClause::AllocateModifier, Align)
   NODE(parser, OmpScheduleClause)
   NODE_ENUM(OmpScheduleClause, ScheduleType)
   NODE(parser, OmpDeviceClause)
@@ -574,6 +578,7 @@ public:
   NODE(parser, OpenMPFlushConstruct)
   NODE(parser, OpenMPLoopConstruct)
   NODE(parser, OpenMPExecutableAllocate)
+  NODE(parser, OpenMPAllocatorsConstruct)
   NODE(parser, OpenMPRequiresConstruct)
   NODE(parser, OpenMPSimpleStandaloneConstruct)
   NODE(parser, OpenMPStandaloneConstruct)
diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h
index d729f444ef95..9059cfbe7f8c 100644
--- a/flang/include/flang/Parser/parse-tree.h
+++ b/flang/include/flang/Parser/parse-tree.h
@@ -3482,10 +3482,23 @@ struct OmpInReductionClause {
 };
 
 // OMP 5.0 2.11.4 allocate-clause -> ALLOCATE ([allocator:] variable-name-list)
+// OMP 5.2 2.13.4 allocate-clause -> ALLOCATE ([allocate-modifier [,
+//  	                               allocate-modifier] :]
+//                                   variable-name-list)
+//                allocate-modifier -> allocator | align
 struct OmpAllocateClause {
   TUPLE_CLASS_BOILERPLATE(OmpAllocateClause);
-  WRAPPER_CLASS(Allocator, ScalarIntExpr);
-  std::tuple<std::optional<Allocator>, OmpObjectList> t;
+  struct AllocateModifier {
+    UNION_CLASS_BOILERPLATE(AllocateModifier);
+    WRAPPER_CLASS(Allocator, ScalarIntExpr);
+    WRAPPER_CLASS(Align, ScalarIntExpr);
+    struct ComplexModifier {
+      TUPLE_CLASS_BOILERPLATE(ComplexModifier);
+      std::tuple<Allocator, Align> t;
+    };
+    std::variant<Allocator, ComplexModifier, Align> u;
+  };
+  std::tuple<std::optional<AllocateModifier>, OmpObjectList> t;
 };
 
 // 2.13.9 depend-vec-length -> +/- non-negative-constant
@@ -3703,6 +3716,20 @@ struct OpenMPExecutableAllocate {
       t;
 };
 
+EMPTY_CLASS(OmpEndAllocators);
+
+// 6.7 Allocators construct [OpenMP 5.2]
+//     allocators-construct -> ALLOCATORS [allocate-clause [,]]
+//                                allocate-stmt
+//                             [omp-end-allocators-construct]
+struct OpenMPAllocatorsConstruct {
+  TUPLE_CLASS_BOILERPLATE(OpenMPAllocatorsConstruct);
+  CharBlock source;
+  std::tuple<Verbatim, OmpClauseList, Statement<AllocateStmt>,
+      std::optional<OmpEndAllocators>>
+      t;
+};
+
 // 2.17.7 Atomic construct/2.17.8 Flush construct [OpenMP 5.0]
 //        memory-order-clause -> acq_rel
 //                               release
@@ -3889,7 +3916,8 @@ struct OpenMPConstruct {
   std::variant<OpenMPStandaloneConstruct, OpenMPSectionsConstruct,
       OpenMPSectionConstruct, OpenMPLoopConstruct, OpenMPBlockConstruct,
       OpenMPAtomicConstruct, OpenMPDeclarativeAllocate,
-      OpenMPExecutableAllocate, OpenMPCriticalConstruct>
+      OpenMPExecutableAllocate, OpenMPAllocatorsConstruct,
+      OpenMPCriticalConstruct>
       u;
 };
 
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 8ab9dc40a50b..6b03a92c956c 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -881,15 +881,31 @@ genAllocateClause(Fortran::lower::AbstractConverter &converter,
   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);
+  const auto &allocateModifier = std::get<
+      std::optional<Fortran::parser::OmpAllocateClause::AllocateModifier>>(
+      ompAllocateClause.t);
+
+  // If the allocate modifier is present, check if we only use the allocator
+  // submodifier.  ALIGN in this context is unimplemented
+  const bool onlyAllocator =
+      allocateModifier &&
+      std::holds_alternative<
+          Fortran::parser::OmpAllocateClause::AllocateModifier::Allocator>(
+          allocateModifier->u);
+
+  if (allocateModifier && !onlyAllocator) {
+    TODO(converter.getCurrentLocation(), "OmpAllocateClause ALIGN modifier");
+  }
+
   // 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) {
+  if (onlyAllocator) {
+    const auto &allocatorValue = std::get<
+        Fortran::parser::OmpAllocateClause::AllocateModifier::Allocator>(
+        allocateModifier->u);
     allocatorOperand = fir::getBase(converter.genExprValue(
-        *Fortran::semantics::GetExpr(allocatorValue->v), stmtCtx));
+        *Fortran::semantics::GetExpr(allocatorValue.v), stmtCtx));
     allocatorOperands.insert(allocatorOperands.end(), ompObjectList.v.size(),
                              allocatorOperand);
   } else {
@@ -2380,6 +2396,10 @@ void Fortran::lower::genOpenMPConstruct(
                   &execAllocConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPExecutableAllocate");
           },
+          [&](const Fortran::parser::OpenMPAllocatorsConstruct
+                  &allocsConstruct) {
+            TODO(converter.getCurrentLocation(), "OpenMPAllocatorsConstruct");
+          },
           [&](const Fortran::parser::OpenMPBlockConstruct &blockConstruct) {
             genOMP(converter, eval, blockConstruct);
           },
diff --git a/flang/lib/Parser/openmp-parsers.cpp b/flang/lib/Parser/openmp-parsers.cpp
index cec08e18e44c..13ebb7f7efdc 100644
--- a/flang/lib/Parser/openmp-parsers.cpp
+++ b/flang/lib/Parser/openmp-parsers.cpp
@@ -140,9 +140,35 @@ TYPE_PARSER(construct<OmpReductionClause>(
 TYPE_PARSER(construct<OmpInReductionClause>(
     Parser<OmpReductionOperator>{} / ":", Parser<OmpObjectList>{}))
 
-// OMP 5.0 2.11.4  ALLOCATE ([allocator:] variable-name-list)
+// OMP 5.0 2.11.4 allocate-clause -> ALLOCATE ([allocator:] variable-name-list)
+// OMP 5.2 2.13.4 allocate-clause -> ALLOCATE ([allocate-modifier
+//                                   [, allocate-modifier] :]
+//                                   variable-name-list)
+//                allocate-modifier -> allocator | align
 TYPE_PARSER(construct<OmpAllocateClause>(
-    maybe(construct<OmpAllocateClause::Allocator>(scalarIntExpr) / ":"),
+    maybe(
+        first(
+            construct<OmpAllocateClause::AllocateModifier>("ALLOCATOR" >>
+                construct<OmpAllocateClause::AllocateModifier::ComplexModifier>(
+                    parenthesized(construct<
+                        OmpAllocateClause::AllocateModifier::Allocator>(
+                        scalarIntExpr)) /
+                        ",",
+                    "ALIGN" >> parenthesized(construct<
+                                   OmpAllocateClause::AllocateModifier::Align>(
+                                   scalarIntExpr)))),
+            construct<OmpAllocateClause::AllocateModifier>("ALLOCATOR" >>
+                parenthesized(
+                    construct<OmpAllocateClause::AllocateModifier::Allocator>(
+                        scalarIntExpr))),
+            construct<OmpAllocateClause::AllocateModifier>("ALIGN" >>
+                parenthesized(
+                    construct<OmpAllocateClause::AllocateModifier::Align>(
+                        scalarIntExpr))),
+            construct<OmpAllocateClause::AllocateModifier>(
+                construct<OmpAllocateClause::AllocateModifier::Allocator>(
+                    scalarIntExpr))) /
+        ":"),
     Parser<OmpObjectList>{}))
 
 // 2.13.9 DEPEND (SOURCE | SINK : vec | (IN | OUT | INOUT) : list
@@ -562,6 +588,16 @@ TYPE_PARSER(
         maybe(nonemptyList(Parser<OpenMPDeclarativeAllocate>{})) / endOmpLine,
         statement(allocateStmt))))
 
+// 6.7 Allocators construct [OpenMP 5.2]
+//     allocators-construct -> ALLOCATORS [allocate-clause [,]]
+//                                allocate-stmt
+//                             [omp-end-allocators-construct]
+TYPE_PARSER(sourced(construct<OpenMPAllocatorsConstruct>(
+    verbatim("ALLOCATORS"_tok), Parser<OmpClauseList>{} / endOmpLine,
+    statement(allocateStmt), maybe(Parser<OmpEndAllocators>{} / endOmpLine))))
+
+TYPE_PARSER(construct<OmpEndAllocators>(startOmpLine >> "END ALLOCATORS"_tok))
+
 // 2.8.2 Declare Simd construct
 TYPE_PARSER(
     sourced(construct<OpenMPDeclareSimdConstruct>(verbatim("DECLARE SIMD"_tok),
@@ -638,6 +674,7 @@ TYPE_CONTEXT_PARSER("OpenMP construct"_en_US,
             construct<OpenMPConstruct>(Parser<OpenMPStandaloneConstruct>{}),
             construct<OpenMPConstruct>(Parser<OpenMPAtomicConstruct>{}),
             construct<OpenMPConstruct>(Parser<OpenMPExecutableAllocate>{}),
+            construct<OpenMPConstruct>(Parser<OpenMPAllocatorsConstruct>{}),
             construct<OpenMPConstruct>(Parser<OpenMPDeclarativeAllocate>{}),
             construct<OpenMPConstruct>(Parser<OpenMPCriticalConstruct>{})))
 
diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp
index 3b34c4ec89ce..4490f7a7cc57 100644
--- a/flang/lib/Parser/unparse.cpp
+++ b/flang/lib/Parser/unparse.cpp
@@ -2042,10 +2042,35 @@ public:
     Walk(std::get<OmpObjectList>(x.t));
   }
   void Unparse(const OmpAllocateClause &x) {
-    Walk(std::get<std::optional<OmpAllocateClause::Allocator>>(x.t));
-    Put(":");
+    Walk(
+        std::get<std::optional<OmpAllocateClause::AllocateModifier>>(x.t), ":");
     Walk(std::get<OmpObjectList>(x.t));
   }
+  void Unparse(const OmpAllocateClause::AllocateModifier &x) {
+    common::visit(
+        common::visitors{
+            [&](const OmpAllocateClause::AllocateModifier::Allocator &y) {
+              Walk(y);
+            },
+            [&](const OmpAllocateClause::AllocateModifier::ComplexModifier &y) {
+              Word("ALLOCATOR(");
+              Walk(std::get<OmpAllocateClause::AllocateModifier::Allocator>(
+                  y.t));
+              Put(")");
+              Put(",");
+              Walk(std::get<OmpAllocateClause::AllocateModifier::Align>(y.t));
+            },
+            [&](const OmpAllocateClause::AllocateModifier::Align &y) {
+              Walk(y);
+            },
+        },
+        x.u);
+  }
+  void Unparse(const OmpAllocateClause::AllocateModifier::Align &x) {
+    Word("ALIGN(");
+    Walk(x.v);
+    Put(")");
+  }
   void Unparse(const OmpOrderClause &x) {
     Walk(std::get<std::optional<OmpOrderModifier>>(x.t), ":");
     Walk(std::get<OmpOrderClause::Type>(x.t));
@@ -2351,6 +2376,23 @@ public:
     Put("\n");
     EndOpenMP();
   }
+  void Unparse(const OmpEndAllocators &x) {
+    BeginOpenMP();
+    Word("!$OMP END ALLOCATE");
+    Put("\n");
+    EndOpenMP();
+  }
+  void Unparse(const OpenMPAllocatorsConstruct &x) {
+    BeginOpenMP();
+    Word("!$OMP ALLOCATE");
+    Walk(std::get<OmpClauseList>(x.t));
+    Put("\n");
+    EndOpenMP();
+    Walk(std::get<Statement<AllocateStmt>>(x.t));
+    if (const auto &end = std::get<std::optional<OmpEndAllocators>>(x.t)) {
+      Walk(*end);
+    }
+  }
   void Unparse(const OmpCriticalDirective &x) {
     BeginOpenMP();
     Word("!$OMP CRITICAL");
diff --git a/flang/test/Parser/OpenMP/allocators-unparse.f90 b/flang/test/Parser/OpenMP/allocators-unparse.f90
new file mode 100644
index 000000000000..062a48b02635
--- /dev/null
+++ b/flang/test/Parser/OpenMP/allocators-unparse.f90
@@ -0,0 +1,63 @@
+! RUN: %flang_fc1 -fdebug-unparse-no-sema -fopenmp %s | FileCheck --ignore-case %s
+! RUN: %flang_fc1 -fdebug-dump-parse-tree-no-sema -fopenmp %s | FileCheck --check-prefix="PARSE-TREE" %s
+! Check unparsing of OpenMP 5.2 Allocators construct
+
+subroutine allocate()
+  use omp_lib
+  integer, allocatable :: arr1(:), arr2(:, :)
+
+  !$omp allocators allocate(omp_default_mem_alloc: arr1)
+    allocate(arr1(5))
+
+  !$omp allocators allocate(allocator(omp_default_mem_alloc), align(32): arr1) &
+  !$omp allocate(omp_default_mem_alloc: arr2)
+    allocate(arr1(10), arr2(3, 2))
+
+  !$omp allocators allocate(align(32): arr2)
+    allocate(arr2(5, 3))
+end subroutine allocate
+
+!CHECK: INTEGER, ALLOCATABLE :: arr1(:), arr2(:,:)
+!CHECK-NEXT:!$OMP ALLOCATE ALLOCATE(omp_default_mem_alloc:arr1)
+!CHECK-NEXT: ALLOCATE(arr1(5))
+!CHECK-NEXT:!$OMP ALLOCATE ALLOCATE(ALLOCATOR(omp_default_mem_alloc),ALIGN(32):arr1) ALLOC&
+!CHECK-NEXT:!$OMP&ATE(omp_default_mem_alloc:arr2)
+!CHECK-NEXT: ALLOCATE(arr1(10), arr2(3,2))
+!CHECK-NEXT:!$OMP ALLOCATE ALLOCATE(ALIGN(32):arr2)
+!CHECK-NEXT: ALLOCATE(arr2(5,3))
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPAllocatorsConstruct
+!PARSE-TREE-NEXT: Verbatim
+!PARSE-TREE-NEXT: OmpClauseList -> OmpClause -> Allocate -> OmpAllocateClause
+!PARSE-TREE-NEXT: AllocateModifier -> Allocator -> Scalar -> Integer -> Expr -> Designator -> DataRef -> Name =
+!PARSE-TREE-NEXT: OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'arr1'
+!PARSE-TREE-NEXT: AllocateStmt
+!PARSE-TREE-NEXT: Allocation
+!PARSE-TREE-NEXT: AllocateObject -> Name = 'arr1'
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPAllocatorsConstruct
+!PARSE-TREE-NEXT: Verbatim
+!PARSE-TREE-NEXT: OmpClauseList -> OmpClause -> Allocate -> OmpAllocateClause
+!PARSE-TREE-NEXT: AllocateModifier -> ComplexModifier
+!PARSE-TREE-NEXT: Allocator -> Scalar -> Integer -> Expr -> Designator -> DataRef -> Name =
+!PARSE-TREE-NEXT: Align -> Scalar -> Integer -> Expr -> LiteralConstant -> IntLiteralConstant = '32'
+!PARSE-TREE-NEXT: OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'arr1'
+!PARSE-TREE-NEXT: OmpClause -> Allocate -> OmpAllocateClause
+!PARSE-TREE-NEXT: AllocateModifier -> Allocator -> Scalar -> Integer -> Expr -> Designator -> DataRef -> Name =
+!PARSE-TREE-NEXT: OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'arr2'
+!PARSE-TREE-NEXT: AllocateStmt
+!PARSE-TREE-NEXT: Allocation
+!PARSE-TREE-NEXT: AllocateObject -> Name = 'arr1'
+!PARSE-TREE-NEXT: AllocateShapeSpec
+!PARSE-TREE-NEXT: Scalar -> Integer -> Expr -> LiteralConstant -> IntLiteralConstant = '10'
+!PARSE-TREE-NEXT: Allocation
+!PARSE-TREE-NEXT: AllocateObject -> Name = 'arr2'
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPAllocatorsConstruct
+!PARSE-TREE-NEXT: Verbatim
+!PARSE-TREE-NEXT: OmpClauseList -> OmpClause -> Allocate -> OmpAllocateClause
+!PARSE-TREE-NEXT: AllocateModifier -> Align -> Scalar -> Integer -> Expr -> LiteralConstant -> IntLiteralConstant = '32'
+!PARSE-TREE-NEXT: OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'arr2'
+!PARSE-TREE-NEXT: AllocateStmt
+!PARSE-TREE-NEXT: Allocation
+!PARSE-TREE-NEXT: AllocateObject -> Name = 'arr2'