[Flang][OpenMP] Add Semantics support for Nested OpenMPLoopConstructs (#145917)

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

@@ -267,6 +267,7 @@ struct AccEndCombinedDirective;
 struct OpenACCDeclarativeConstruct;
 struct OpenACCRoutineConstruct;
 struct OpenMPConstruct;
+struct OpenMPLoopConstruct;
 struct OpenMPDeclarativeConstruct;
 struct OmpEndLoopDirective;
 struct OmpMemoryOrderClause;
@@ -5021,11 +5022,13 @@ struct OpenMPBlockConstruct {
 };
 
 // OpenMP directives enclosing do loop
+using NestedConstruct =
+    std::variant<DoConstruct, common::Indirection<OpenMPLoopConstruct>>;
 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<NestedConstruct>,
       std::optional<OmpEndLoopDirective>>
       t;
 };

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -3769,6 +3769,16 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
   mlir::Location currentLocation =
       converter.genLocation(beginLoopDirective.source);
 
+  auto &optLoopCons =
+      std::get<std::optional<parser::NestedConstruct>>(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,8 @@ 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
@@ -125,6 +126,16 @@ class CanonicalizationOfOmp {
     parser::Block::iterator nextIt;
     auto &beginDir{std::get<parser::OmpBeginLoopDirective>(x.t)};
     auto &dir{std::get<parser::OmpLoopDirective>(beginDir.t)};
+    auto missingDoConstruct = [](auto &dir, auto &messages) {
+      messages.Say(dir.source,
+          "A DO loop must follow the %s directive"_err_en_US,
+          parser::ToUpperCaseLetters(dir.source.ToString()));
+    };
+    auto tileUnrollError = [](auto &dir, auto &messages) {
+      messages.Say(dir.source,
+          "If a loop construct has been fully unrolled, it cannot then be tiled"_err_en_US,
+          parser::ToUpperCaseLetters(dir.source.ToString()));
+    };
 
     nextIt = it;
     while (++nextIt != block.end()) {
@@ -135,31 +146,95 @@ 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 &nestedBeginDirective =
+            std::get<parser::OmpBeginLoopDirective>(ompLoopCons->t);
+        auto &nestedBeginLoopDirective =
+            std::get<parser::OmpLoopDirective>(nestedBeginDirective.t);
+        if ((nestedBeginLoopDirective.v == llvm::omp::Directive::OMPD_unroll ||
+                nestedBeginLoopDirective.v ==
+                    llvm::omp::Directive::OMPD_tile) &&
+            !(nestedBeginLoopDirective.v == llvm::omp::Directive::OMPD_unroll &&
+                dir.v == llvm::omp::Directive::OMPD_tile)) {
+          // 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;
+          }
+          RewriteOpenMPLoopConstruct(*ompLoopCons, block, nextIt);
+          auto &ompLoop = std::get<std::optional<parser::NestedConstruct>>(x.t);
+          ompLoop =
+              std::optional<parser::NestedConstruct>{parser::NestedConstruct{
+                  common::Indirection{std::move(*ompLoopCons)}}};
+          nextIt = block.erase(nextIt);
+        } else if (nestedBeginLoopDirective.v ==
+                llvm::omp::Directive::OMPD_unroll &&
+            dir.v == llvm::omp::Directive::OMPD_tile) {
+          // if a loop has been unrolled, the user can not then tile that loop
+          // as it has been unrolled
+          parser::OmpClauseList &unrollClauseList{
+              std::get<parser::OmpClauseList>(nestedBeginDirective.t)};
+          if (unrollClauseList.v.empty()) {
+            // if the clause list is empty for an unroll construct, we assume
+            // the loop is being fully unrolled
+            tileUnrollError(dir, messages_);
+          } else {
+            // parse the clauses for the unroll directive to find the full
+            // clause
+            for (auto clause{unrollClauseList.v.begin()};
+                clause != unrollClauseList.v.end(); ++clause) {
+              if (clause->Id() == llvm::omp::OMPC_full) {
+                tileUnrollError(dir, messages_);
+              }
+            }
+          }
+        } else {
+          messages_.Say(nestedBeginLoopDirective.source,
+              "Only Loop Transformation Constructs or Loop Nests can be nested within Loop Constructs"_err_en_US,
+              parser::ToUpperCaseLetters(
+                  nestedBeginLoopDirective.source.ToString()));
+        }
       } else {
-        messages_.Say(dir.source,
-            "A DO loop must follow the %s directive"_err_en_US,
-            parser::ToUpperCaseLetters(dir.source.ToString()));
+        missingDoConstruct(dir, messages_);
       }
       // If we get here, we either found a loop, or issued an error message.
       return;
     }
+    if (nextIt == block.end()) {
+      missingDoConstruct(dir, messages_);
+    }
   }
 
   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<std::optional<parser::NestedConstruct>>(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);
@@ -786,12 +789,15 @@ const parser::Name OmpStructureChecker::GetLoopIndex(
   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<std::optional<parser::NestedConstruct>>(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,27 +863,30 @@ 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<std::optional<parser::NestedConstruct>>(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()) {
-        const parser::Name &itrVal{GetLoopIndex(loop)};
-        if (itrVal.symbol) {
-          const auto *type{itrVal.symbol->GetType()};
-          if (!type->IsNumeric(TypeCategory::Integer)) {
-            context_.Say(itrVal.source,
-                "The DO loop iteration"
-                " variable must be of the type integer."_err_en_US,
-                itrVal.ToString());
+      for (const parser::DoConstruct *loop{&*loopConstruct}; loop;) {
+        if (loop->IsDoNormal()) {
+          const parser::Name &itrVal{GetLoopIndex(loop)};
+          if (itrVal.symbol) {
+            const auto *type{itrVal.symbol->GetType()};
+            if (!type->IsNumeric(TypeCategory::Integer)) {
+              context_.Say(itrVal.source,
+                  "The DO loop iteration"
+                  " variable must be of the type integer."_err_en_US,
+                  itrVal.ToString());
+            }
           }
         }
+        // Get the next DoConstruct if block is not empty.
+        const auto &block{std::get<parser::Block>(loop->t)};
+        const auto it{block.begin()};
+        loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
+                                 : nullptr;
       }
-      // Get the next DoConstruct if block is not empty.
-      const auto &block{std::get<parser::Block>(loop->t)};
-      const auto it{block.begin()};
-      loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
-                               : nullptr;
     }
   }
 }
@@ -1077,25 +1086,28 @@ void OmpStructureChecker::CheckDistLinear(
 
     // Match the loop index variables with the collected symbols from linear
     // clauses.
-    if (const auto &loopConstruct{
-            std::get<std::optional<parser::DoConstruct>>(x.t)}) {
-      for (const parser::DoConstruct *loop{&*loopConstruct}; loop;) {
-        if (loop->IsDoNormal()) {
-          const parser::Name &itrVal{GetLoopIndex(loop)};
-          if (itrVal.symbol) {
-            // Remove the symbol from the collected set
-            indexVars.erase(&itrVal.symbol->GetUltimate());
-          }
-          collapseVal--;
-          if (collapseVal == 0) {
-            break;
+    auto &optLoopCons = std::get<std::optional<parser::NestedConstruct>>(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()) {
+            const parser::Name &itrVal{GetLoopIndex(loop)};
+            if (itrVal.symbol) {
+              // Remove the symbol from the collected set
+              indexVars.erase(&itrVal.symbol->GetUltimate());
+            }
+            collapseVal--;
+            if (collapseVal == 0) {
+              break;
+            }
           }
+          // Get the next DoConstruct if block is not empty.
+          const auto &block{std::get<parser::Block>(loop->t)};
+          const auto it{block.begin()};
+          loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
+                                   : nullptr;
         }
-        // Get the next DoConstruct if block is not empty.
-        const auto &block{std::get<parser::Block>(loop->t)};
-        const auto it{block.begin()};
-        loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
-                                 : nullptr;
       }
     }