[Flang][OpenMP] Add Semantics support for Nested OpenMPLoopConstructs

In OpenMP Version 5.1, the tile and unroll directives were added. When
using these directives, it is possible to nest them within other OpenMP
Loop Constructs. This patch enables the semantics to allow for this
behaviour on these specific directives. Any nested loops will be stored
within the initial Loop Construct until reaching the DoConstruct itself.

Relevant tests have been added, and previous behaviour has been retained
with no changes.

See also, #110008

Author: Stylie777

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

@@ -5025,7 +5025,7 @@ struct OpenMPLoopConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenMPLoopConstruct);
   OpenMPLoopConstruct(OmpBeginLoopDirective &&a)
       : t({std::move(a), std::nullopt, std::nullopt}) {}
-  std::tuple<OmpBeginLoopDirective, std::optional<DoConstruct>,
+  std::tuple<OmpBeginLoopDirective, std::optional<std::variant<DoConstruct, common::Indirection<OpenMPLoopConstruct>>>,
       std::optional<OmpEndLoopDirective>>
       t;
 };

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -4231,6 +4231,12 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
   mlir::Location currentLocation =
       converter.genLocation(beginLoopDirective.source);
 
+  auto &optLoopCons = std::get<1>(loopConstruct.t);
+  if(optLoopCons.has_value())
+    if(auto *ompNestedLoopCons{std::get_if<common::Indirection<parser::OpenMPLoopConstruct>>(&*optLoopCons)}) {
+      genOMP(converter, symTable, semaCtx, eval, ompNestedLoopCons->value());
+    }
+
   llvm::omp::Directive directive =
       std::get<parser::OmpLoopDirective>(beginLoopDirective.t).v;
   const parser::CharBlock &source =

flang/lib/Parser/unparse.cpp

@@ -2926,7 +2926,7 @@ class UnparseVisitor {
     Walk(std::get<OmpBeginLoopDirective>(x.t));
     Put("\n");
     EndOpenMP();
-    Walk(std::get<std::optional<DoConstruct>>(x.t));
+    Walk(std::get<std::optional<std::variant<DoConstruct, common::Indirection<parser::OpenMPLoopConstruct>>>>(x.t));
     Walk(std::get<std::optional<OmpEndLoopDirective>>(x.t));
   }
   void Unparse(const BasedPointer &x) {

flang/lib/Semantics/canonicalize-omp.cpp

@@ -8,6 +8,7 @@
 
 #include "canonicalize-omp.h"
 #include "flang/Parser/parse-tree-visitor.h"
+#include "flang/Parser/parse-tree.h"
 
 // After Loop Canonicalization, rewrite OpenMP parse tree to make OpenMP
 // Constructs more structured which provide explicit scopes for later
@@ -106,6 +107,12 @@ class CanonicalizationOfOmp {
     return nullptr;
   }
 
+  void missingDoConstruct(parser::OmpLoopDirective &dir) {
+    messages_.Say(dir.source,
+      "A DO loop must follow the %s directive"_err_en_US,
+      parser::ToUpperCaseLetters(dir.source.ToString()));
+  }
+
   void RewriteOpenMPLoopConstruct(parser::OpenMPLoopConstruct &x,
       parser::Block &block, parser::Block::iterator it) {
     // Check the sequence of DoConstruct and OmpEndLoopDirective
@@ -135,31 +142,62 @@ class CanonicalizationOfOmp {
       if (auto *doCons{GetConstructIf<parser::DoConstruct>(*nextIt)}) {
         if (doCons->GetLoopControl()) {
           // move DoConstruct
-          std::get<std::optional<parser::DoConstruct>>(x.t) =
+          std::get<std::optional<std::variant<parser::DoConstruct, common::Indirection<parser::OpenMPLoopConstruct>>>>(x.t) =
               std::move(*doCons);
           nextIt = block.erase(nextIt);
           // try to match OmpEndLoopDirective
-          if (nextIt != block.end()) {
-            if (auto *endDir{
-                    GetConstructIf<parser::OmpEndLoopDirective>(*nextIt)}) {
-              std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
-                  std::move(*endDir);
-              block.erase(nextIt);
-            }
+          if (auto *endDir{
+                  GetConstructIf<parser::OmpEndLoopDirective>(*nextIt)}) {
+            std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
+                std::move(*endDir);
+            nextIt = block.erase(nextIt);
           }
         } else {
           messages_.Say(dir.source,
               "DO loop after the %s directive must have loop control"_err_en_US,
               parser::ToUpperCaseLetters(dir.source.ToString()));
         }
+      } else if (auto *ompLoopCons{
+                     GetOmpIf<parser::OpenMPLoopConstruct>(*nextIt)}) {
+        // We should allow UNROLL and TILE constructs to be inserted between an OpenMP Loop Construct and the DO loop itself
+        auto &beginDirective =
+            std::get<parser::OmpBeginLoopDirective>(ompLoopCons->t);
+        auto &beginLoopDirective =
+        std::get<parser::OmpLoopDirective>(beginDirective.t);
+        // iterate through the remaining block items to find the end directive for the unroll/tile directive.
+        parser::Block::iterator endIt;
+        endIt = nextIt;
+        while(endIt != block.end()) {
+          if (auto *endDir{
+            GetConstructIf<parser::OmpEndLoopDirective>(*endIt)}) {
+              auto &endLoopDirective = std::get<parser::OmpLoopDirective>(endDir->t);
+              if(endLoopDirective.v == dir.v) {
+                std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
+                std::move(*endDir);
+                endIt = block.erase(endIt);
+                continue;
+              }
+            }
+            ++endIt;
+          }
+        if ((beginLoopDirective.v == llvm::omp::Directive::OMPD_unroll ||
+            beginLoopDirective.v == llvm::omp::Directive::OMPD_tile)) {
+          RewriteOpenMPLoopConstruct(*ompLoopCons, block, nextIt);
+          auto &ompLoop = std::get<std::optional<std::variant<parser::DoConstruct, common::Indirection<parser::OpenMPLoopConstruct>>>>(x.t);
+          ompLoop = std::optional<std::variant<parser::DoConstruct, common::Indirection<parser::OpenMPLoopConstruct>>>{
+            std::variant<parser::DoConstruct, common::Indirection<parser::OpenMPLoopConstruct>>{
+            common::Indirection{std::move(*ompLoopCons)}}};
+          nextIt = block.erase(nextIt);
+        }
       } else {
-        messages_.Say(dir.source,
-            "A DO loop must follow the %s directive"_err_en_US,
-            parser::ToUpperCaseLetters(dir.source.ToString()));
+        missingDoConstruct(dir);
       }
       // If we get here, we either found a loop, or issued an error message.
       return;
     }
+    if (nextIt == block.end()) {
+      missingDoConstruct(dir);
+    }
   }
 
   void RewriteOmpAllocations(parser::ExecutionPart &body) {

flang/lib/Semantics/check-omp-structure.cpp

@@ -762,10 +762,13 @@ void OmpStructureChecker::Enter(const parser::OpenMPLoopConstruct &x) {
   }
   SetLoopInfo(x);
 
-  if (const auto &doConstruct{
-          std::get<std::optional<parser::DoConstruct>>(x.t)}) {
-    const auto &doBlock{std::get<parser::Block>(doConstruct->t)};
-    CheckNoBranching(doBlock, beginDir.v, beginDir.source);
+  auto &optLoopCons = std::get<1>(x.t);
+  if(optLoopCons.has_value()) {
+    if (const auto &doConstruct{
+            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+      const auto &doBlock{std::get<parser::Block>(doConstruct->t)};
+      CheckNoBranching(doBlock, beginDir.v, beginDir.source);
+    }
   }
   CheckLoopItrVariableIsInt(x);
   CheckAssociatedLoopConstraints(x);
@@ -779,19 +782,28 @@ void OmpStructureChecker::Enter(const parser::OpenMPLoopConstruct &x) {
       (beginDir.v == llvm::omp::Directive::OMPD_distribute_simd)) {
     CheckDistLinear(x);
   }
+  if (beginDir.v == llvm::omp::Directive::OMPD_tile) {
+    const auto &clauses{std::get<parser::OmpClauseList>(beginLoopDir.t)};
+    for (auto &clause : clauses.v) {
+
+    }
+  }
 }
 const parser::Name OmpStructureChecker::GetLoopIndex(
     const parser::DoConstruct *x) {
   using Bounds = parser::LoopControl::Bounds;
   return std::get<Bounds>(x->GetLoopControl()->u).name.thing;
 }
 void OmpStructureChecker::SetLoopInfo(const parser::OpenMPLoopConstruct &x) {
-  if (const auto &loopConstruct{
-          std::get<std::optional<parser::DoConstruct>>(x.t)}) {
-    const parser::DoConstruct *loop{&*loopConstruct};
-    if (loop && loop->IsDoNormal()) {
-      const parser::Name &itrVal{GetLoopIndex(loop)};
-      SetLoopIv(itrVal.symbol);
+  auto &optLoopCons = std::get<1>(x.t);
+  if (optLoopCons.has_value()) {
+    if (const auto &loopConstruct{
+            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+      const parser::DoConstruct *loop{&*loopConstruct};
+      if (loop && loop->IsDoNormal()) {
+        const parser::Name &itrVal{GetLoopIndex(loop)};
+        SetLoopIv(itrVal.symbol);
+      }
     }
   }
 }
@@ -857,8 +869,10 @@ void OmpStructureChecker::CheckIteratorModifier(const parser::OmpIterator &x) {
 
 void OmpStructureChecker::CheckLoopItrVariableIsInt(
     const parser::OpenMPLoopConstruct &x) {
-  if (const auto &loopConstruct{
-          std::get<std::optional<parser::DoConstruct>>(x.t)}) {
+  auto &optLoopCons = std::get<1>(x.t);
+  if (optLoopCons.has_value()) {
+    if (const auto &loopConstruct{
+          std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
 
     for (const parser::DoConstruct *loop{&*loopConstruct}; loop;) {
       if (loop->IsDoNormal()) {
@@ -878,6 +892,7 @@ void OmpStructureChecker::CheckLoopItrVariableIsInt(
       const auto it{block.begin()};
       loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
                                : nullptr;
+      }
     }
   }
 }
@@ -1077,8 +1092,10 @@ void OmpStructureChecker::CheckDistLinear(
 
     // Match the loop index variables with the collected symbols from linear
     // clauses.
+    auto &optLoopCons = std::get<1>(x.t);
+    if (optLoopCons.has_value()) {
     if (const auto &loopConstruct{
-            std::get<std::optional<parser::DoConstruct>>(x.t)}) {
+            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
       for (const parser::DoConstruct *loop{&*loopConstruct}; loop;) {
         if (loop->IsDoNormal()) {
           const parser::Name &itrVal{GetLoopIndex(loop)};
@@ -1096,6 +1113,7 @@ void OmpStructureChecker::CheckDistLinear(
         const auto it{block.begin()};
         loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
                                  : nullptr;
+        }
       }
     }
 

flang/lib/Semantics/resolve-directives.cpp

@@ -1796,10 +1796,13 @@ bool OmpAttributeVisitor::Pre(const parser::OpenMPLoopConstruct &x) {
   SetContextAssociatedLoopLevel(GetAssociatedLoopLevelFromClauses(clauseList));
 
   if (beginDir.v == llvm::omp::Directive::OMPD_do) {
-    if (const auto &doConstruct{
-            std::get<std::optional<parser::DoConstruct>>(x.t)}) {
-      if (doConstruct.value().IsDoWhile()) {
-        return true;
+    auto &optLoopCons = std::get<1>(x.t);
+    if (optLoopCons.has_value()) {
+      if (const auto &doConstruct{
+              std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+        if (doConstruct->IsDoWhile()) {
+          return true;
+        }
       }
     }
   }
@@ -1962,48 +1965,64 @@ void OmpAttributeVisitor::PrivatizeAssociatedLoopIndexAndCheckLoopLevel(
   bool hasCollapseClause{
       clause ? (clause->Id() == llvm::omp::OMPC_collapse) : false};
 
-  const auto &outer{std::get<std::optional<parser::DoConstruct>>(x.t)};
-  if (outer.has_value()) {
-    for (const parser::DoConstruct *loop{&*outer}; loop && level > 0; --level) {
-      if (loop->IsDoConcurrent()) {
-        // DO CONCURRENT is explicitly allowed for the LOOP construct so long as
-        // there isn't a COLLAPSE clause
-        if (isLoopConstruct) {
-          if (hasCollapseClause) {
-            // hasCollapseClause implies clause != nullptr
-            context_.Say(clause->source,
-                "DO CONCURRENT loops cannot be used with the COLLAPSE clause."_err_en_US);
+  auto &optLoopCons = std::get<1>(x.t);
+  if (optLoopCons.has_value()) {
+    if (const auto &outer{std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+      for (const parser::DoConstruct *loop{&*outer}; loop && level > 0; --level) {
+        if (loop->IsDoConcurrent()) {
+          // DO CONCURRENT is explicitly allowed for the LOOP construct so long as
+          // there isn't a COLLAPSE clause
+          if (isLoopConstruct) {
+            if (hasCollapseClause) {
+              // hasCollapseClause implies clause != nullptr
+              context_.Say(clause->source,
+                  "DO CONCURRENT loops cannot be used with the COLLAPSE clause."_err_en_US);
+            }
+          } else {
+            auto &stmt =
+                std::get<parser::Statement<parser::NonLabelDoStmt>>(loop->t);
+            context_.Say(stmt.source,
+                "DO CONCURRENT loops cannot form part of a loop nest."_err_en_US);
           }
-        } else {
-          auto &stmt =
-              std::get<parser::Statement<parser::NonLabelDoStmt>>(loop->t);
-          context_.Say(stmt.source,
-              "DO CONCURRENT loops cannot form part of a loop nest."_err_en_US);
-        }
-      }
-      // go through all the nested do-loops and resolve index variables
-      const parser::Name *iv{GetLoopIndex(*loop)};
-      if (iv) {
-        if (auto *symbol{ResolveOmp(*iv, ivDSA, currScope())}) {
-          SetSymbolDSA(*symbol, {Symbol::Flag::OmpPreDetermined, ivDSA});
-          iv->symbol = symbol; // adjust the symbol within region
-          AddToContextObjectWithDSA(*symbol, ivDSA);
         }
+        // go through all the nested do-loops and resolve index variables
+        const parser::Name *iv{GetLoopIndex(*loop)};
+        if (iv) {
+          if (auto *symbol{ResolveOmp(*iv, ivDSA, currScope())}) {
+            SetSymbolDSA(*symbol, {Symbol::Flag::OmpPreDetermined, ivDSA});
+            iv->symbol = symbol; // adjust the symbol within region
+            AddToContextObjectWithDSA(*symbol, ivDSA);
+          }
 
-        const auto &block{std::get<parser::Block>(loop->t)};
-        const auto it{block.begin()};
-        loop = it != block.end() ? GetDoConstructIf(*it) : nullptr;
+          const auto &block{std::get<parser::Block>(loop->t)};
+          const auto it{block.begin()};
+          loop = it != block.end() ? GetDoConstructIf(*it) : nullptr;
+        }
+      }
+      CheckAssocLoopLevel(level, GetAssociatedClause());
+    } else if (const auto &loop{std::get_if<common::Indirection<parser::OpenMPLoopConstruct>>(&*optLoopCons)}) {
+      auto &beginDirective =
+            std::get<parser::OmpBeginLoopDirective>(loop->value().t);
+      auto &beginLoopDirective =
+          std::get<parser::OmpLoopDirective>(beginDirective.t);
+      if ((beginLoopDirective.v != llvm::omp::Directive::OMPD_unroll &&
+            beginLoopDirective.v != llvm::omp::Directive::OMPD_tile)) {
+        context_.Say(GetContext().directiveSource,
+          "Only UNROLL or TILE constructs are allowed between an OpenMP Loop Construct and a DO construct"_err_en_US,
+          parser::ToUpperCaseLetters(llvm::omp::getOpenMPDirectiveName(GetContext().directive, version).str()));
+      } else {
+        PrivatizeAssociatedLoopIndexAndCheckLoopLevel(loop->value());
       }
+    } else {
+      context_.Say(GetContext().directiveSource,
+          "A DO loop must follow the %s directive"_err_en_US,
+          parser::ToUpperCaseLetters(
+              llvm::omp::getOpenMPDirectiveName(GetContext().directive, version)
+                  .str()));
     }
-    CheckAssocLoopLevel(level, GetAssociatedClause());
-  } else {
-    context_.Say(GetContext().directiveSource,
-        "A DO loop must follow the %s directive"_err_en_US,
-        parser::ToUpperCaseLetters(
-            llvm::omp::getOpenMPDirectiveName(GetContext().directive, version)
-                .str()));
   }
 }
+
 void OmpAttributeVisitor::CheckAssocLoopLevel(
     std::int64_t level, const parser::OmpClause *clause) {
   if (clause && level != 0) {

flang/test/Lower/OpenMP/nested-loop-transformation-construct01.f90

@@ -0,0 +1,20 @@
+! Test to ensure TODO message is emitted for tile OpenMP 5.1 Directives when they are nested.
+
+!RUN: not %flang -fopenmp -fopenmp-version=51 %s 2<&1 | FileCheck %s
+
+subroutine loop_transformation_construct
+  implicit none
+  integer :: I = 10
+  integer :: x
+  integer :: y(I)
+
+  !$omp do
+  !$omp tile
+  do i = 1, I
+    y(i) = y(i) * 5
+  end do
+  !$omp end tile
+  !$omp end do
+end subroutine
+
+!CHECK: not yet implemented: Unhandled loop directive (tile)

flang/test/Lower/OpenMP/nested-loop-transformation-construct02.f90

@@ -0,0 +1,20 @@
+! Test to ensure TODO message is emitted for unroll OpenMP 5.1 Directives when they are nested.
+
+!RUN: not %flang -fopenmp -fopenmp-version=51 %s 2<&1 | FileCheck %s
+
+program loop_transformation_construct
+  implicit none
+  integer, parameter :: I = 10
+  integer :: x
+  integer :: y(I)
+
+  !$omp do
+  !$omp unroll
+  do x = 1, I
+    y(x) = y(x) * 5
+  end do
+  !$omp end unroll
+  !$omp end do
+end program loop_transformation_construct
+
+!CHECK: not yet implemented: Unhandled loop directive (unroll)