Loop sequences and loop fuse semantics

Author: Ferran Toda

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

@@ -5177,15 +5177,15 @@ using NestedConstruct =
 struct OpenMPLoopConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenMPLoopConstruct);
   OpenMPLoopConstruct(OmpBeginLoopDirective &&a)
-      : t({std::move(a), std::nullopt, std::nullopt}) {}
+      : t({std::move(a), std::list<NestedConstruct>(), std::nullopt}) {}
 
   const OmpBeginLoopDirective &BeginDir() const {
     return std::get<OmpBeginLoopDirective>(t);
   }
   const std::optional<OmpEndLoopDirective> &EndDir() const {
     return std::get<std::optional<OmpEndLoopDirective>>(t);
   }
-  std::tuple<OmpBeginLoopDirective, std::optional<NestedConstruct>,
+  std::tuple<OmpBeginLoopDirective, std::list<NestedConstruct>,
       std::optional<OmpEndLoopDirective>>
       t;
 };

flang/include/flang/Semantics/openmp-directive-sets.h

@@ -275,10 +275,17 @@ static const OmpDirectiveSet loopConstructSet{
     Directive::OMPD_teams_distribute_parallel_do_simd,
     Directive::OMPD_teams_distribute_simd,
     Directive::OMPD_teams_loop,
+    Directive::OMPD_fuse,
     Directive::OMPD_tile,
     Directive::OMPD_unroll,
 };
 
+static const OmpDirectiveSet loopTransformationSet{
+    Directive::OMPD_tile,
+    Directive::OMPD_unroll,
+    Directive::OMPD_fuse,
+};
+
 static const OmpDirectiveSet nonPartialVarSet{
     Directive::OMPD_allocate,
     Directive::OMPD_allocators,

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -3360,6 +3360,7 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
     newOp = genTeamsOp(converter, symTable, stmtCtx, semaCtx, eval, loc, queue,
                        item);
     break;
+  case llvm::omp::Directive::OMPD_fuse:
   case llvm::omp::Directive::OMPD_tile: {
     unsigned version = semaCtx.langOptions().OpenMPVersion;
     if (!semaCtx.langOptions().OpenMPSimd)
@@ -3814,12 +3815,12 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
 
   mlir::Location currentLocation = converter.genLocation(beginSpec.source);
 
-  auto &optLoopCons =
-      std::get<std::optional<parser::NestedConstruct>>(loopConstruct.t);
-  if (optLoopCons.has_value()) {
+  auto &loopConsList =
+      std::get<std::list<parser::NestedConstruct>>(loopConstruct.t);
+  for (auto &loopCons : loopConsList) {
     if (auto *ompNestedLoopCons{
             std::get_if<common::Indirection<parser::OpenMPLoopConstruct>>(
-                &*optLoopCons)}) {
+                &loopCons)}) {
       llvm::omp::Directive nestedDirective =
           parser::omp::GetOmpDirectiveName(*ompNestedLoopCons).v;
       switch (nestedDirective) {

flang/lib/Lower/OpenMP/Utils.cpp

@@ -607,13 +607,13 @@ static void processTileSizesFromOpenMPConstruct(
   if (!ompCons)
     return;
   if (auto *ompLoop{std::get_if<parser::OpenMPLoopConstruct>(&ompCons->u)}) {
-    const auto &nestedOptional =
-        std::get<std::optional<parser::NestedConstruct>>(ompLoop->t);
-    assert(nestedOptional.has_value() &&
+    const auto &loopConsList =
+        std::get<std::list<parser::NestedConstruct>>(ompLoop->t);
+    assert(loopConsList.size() == 1 &&
            "Expected a DoConstruct or OpenMPLoopConstruct");
     const auto *innerConstruct =
         std::get_if<common::Indirection<parser::OpenMPLoopConstruct>>(
-            &(nestedOptional.value()));
+            &(loopConsList.front()));
     if (innerConstruct) {
       const auto &innerLoopDirective = innerConstruct->value();
       const parser::OmpDirectiveSpecification &innerBeginSpec =

flang/lib/Parser/openmp-parsers.cpp

@@ -2019,6 +2019,7 @@ static constexpr DirectiveSet GetLoopDirectives() {
       unsigned(Directive::OMPD_teams_distribute_parallel_do_simd),
       unsigned(Directive::OMPD_teams_distribute_simd),
       unsigned(Directive::OMPD_teams_loop),
+      unsigned(Directive::OMPD_fuse),
       unsigned(Directive::OMPD_tile),
       unsigned(Directive::OMPD_unroll),
   };

flang/lib/Parser/unparse.cpp

@@ -2720,8 +2720,7 @@ class UnparseVisitor {
   }
   void Unparse(const OpenMPLoopConstruct &x) {
     Walk(std::get<OmpBeginLoopDirective>(x.t));
-    Walk(std::get<std::optional<std::variant<DoConstruct,
-            common::Indirection<parser::OpenMPLoopConstruct>>>>(x.t));
+    Walk(std::get<std::list<parser::NestedConstruct>>(x.t));
     Walk(std::get<std::optional<OmpEndLoopDirective>>(x.t));
   }
   void Unparse(const BasedPointer &x) {

flang/lib/Semantics/canonicalize-omp.cpp

@@ -10,6 +10,7 @@
 #include "flang/Parser/parse-tree-visitor.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/semantics.h"
+#include "flang/Semantics/openmp-directive-sets.h"
 
 // After Loop Canonicalization, rewrite OpenMP parse tree to make OpenMP
 // Constructs more structured which provide explicit scopes for later
@@ -137,30 +138,42 @@ class CanonicalizationOfOmp {
           "A DO loop must follow the %s directive"_err_en_US,
           parser::ToUpperCaseLetters(dirName.source.ToString()));
     };
-    auto tileUnrollError = [](const parser::OmpDirectiveName &dirName,
+    auto transformUnrollError = [](const parser::OmpDirectiveName &dirName,
                                parser::Messages &messages) {
       messages.Say(dirName.source,
-          "If a loop construct has been fully unrolled, it cannot then be tiled"_err_en_US,
+          "If a loop construct has been fully unrolled, it cannot then be further transformed"_err_en_US,
           parser::ToUpperCaseLetters(dirName.source.ToString()));
     };
+    auto missingEndFuse = [](auto &dir, auto &messages) {
+      messages.Say(dir.source,
+          "The %s construct requires the END FUSE directive"_err_en_US,
+          parser::ToUpperCaseLetters(dir.source.ToString()));
+    };
+
+    bool endFuseNeeded = beginName.v == llvm::omp::Directive::OMPD_fuse;
 
     nextIt = it;
-    while (++nextIt != block.end()) {
+    nextIt++;
+    while (nextIt != block.end()) {
       // Ignore compiler directives.
-      if (GetConstructIf<parser::CompilerDirective>(*nextIt))
+      if (GetConstructIf<parser::CompilerDirective>(*nextIt)) {
+        nextIt++;
         continue;
+      }
 
       if (auto *doCons{GetConstructIf<parser::DoConstruct>(*nextIt)}) {
         if (doCons->GetLoopControl()) {
           // move DoConstruct
-          std::get<std::optional<std::variant<parser::DoConstruct,
-              common::Indirection<parser::OpenMPLoopConstruct>>>>(x.t) =
-              std::move(*doCons);
+          std::get<std::list<parser::NestedConstruct>>(x.t).push_back(
+              std::move(*doCons));
           nextIt = block.erase(nextIt);
           // try to match OmpEndLoopDirective
           if (nextIt != block.end()) {
             if (auto *endDir{
                     GetConstructIf<parser::OmpEndLoopDirective>(*nextIt)}) {
+              auto &endDirName = endDir->DirName();
+              if (endDirName.v == llvm::omp::Directive::OMPD_fuse)
+                endFuseNeeded = false;
               std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
                   std::move(*endDir);
               nextIt = block.erase(nextIt);
@@ -170,17 +183,37 @@ class CanonicalizationOfOmp {
           messages_.Say(beginName.source,
               "DO loop after the %s directive must have loop control"_err_en_US,
               parser::ToUpperCaseLetters(beginName.source.ToString()));
+          endFuseNeeded = false;
         }
       } 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 = ompLoopCons->BeginDir();
         auto &nestedBeginName = nestedBeginDirective.DirName();
-        if ((nestedBeginName.v == llvm::omp::Directive::OMPD_unroll ||
-                nestedBeginName.v == llvm::omp::Directive::OMPD_tile) &&
-            !(nestedBeginName.v == llvm::omp::Directive::OMPD_unroll &&
-                beginName.v == llvm::omp::Directive::OMPD_tile)) {
+        if (llvm::omp::loopTransformationSet.test(nestedBeginName.v)) {
+          if (nestedBeginName.v == llvm::omp::Directive::OMPD_unroll &&
+              llvm::omp::loopTransformationSet.test(beginName.v)) {
+            // if a loop has been unrolled, the user can not then tile that loop
+            // as it has been unrolled
+            const parser::OmpClauseList &unrollClauseList{
+              nestedBeginDirective.Clauses()};
+            if (unrollClauseList.v.empty()) {
+              // if the clause list is empty for an unroll construct, we assume
+              // the loop is being fully unrolled
+              transformUnrollError(beginName, messages_);
+              endFuseNeeded = false;
+            } else {
+              // parse the clauses for the unroll directive to find the full
+              // clause
+              for (auto &clause : unrollClauseList.v) {
+                if (clause.Id() == llvm::omp::OMPC_full) {
+                  transformUnrollError(beginName, messages_);
+                  endFuseNeeded = false;
+                }
+              }
+            }
+          }
           // iterate through the remaining block items to find the end directive
           // for the unroll/tile directive.
           parser::Block::iterator endIt;
@@ -190,6 +223,8 @@ class CanonicalizationOfOmp {
                     GetConstructIf<parser::OmpEndLoopDirective>(*endIt)}) {
               auto &endDirName = endDir->DirName();
               if (endDirName.v == beginName.v) {
+                if (endDirName.v == llvm::omp::Directive::OMPD_fuse)
+                  endFuseNeeded = false;
                 std::get<std::optional<parser::OmpEndLoopDirective>>(x.t) =
                     std::move(*endDir);
                 endIt = block.erase(endIt);
@@ -199,43 +234,30 @@ class CanonicalizationOfOmp {
             ++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)}}};
+          auto &loopConsList = std::get<std::list<parser::NestedConstruct>>(x.t);
+          loopConsList.push_back(parser::NestedConstruct{
+              common::Indirection{std::move(*ompLoopCons)}});
           nextIt = block.erase(nextIt);
-        } else if (nestedBeginName.v == llvm::omp::Directive::OMPD_unroll &&
-            beginName.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
-          const parser::OmpClauseList &unrollClauseList{
-              nestedBeginDirective.Clauses()};
-          if (unrollClauseList.v.empty()) {
-            // if the clause list is empty for an unroll construct, we assume
-            // the loop is being fully unrolled
-            tileUnrollError(beginName, messages_);
-          } else {
-            // parse the clauses for the unroll directive to find the full
-            // clause
-            for (auto &clause : unrollClauseList.v) {
-              if (clause.Id() == llvm::omp::OMPC_full) {
-                tileUnrollError(beginName, messages_);
-              }
-            }
-          }
         } else {
           messages_.Say(nestedBeginName.source,
               "Only Loop Transformation Constructs or Loop Nests can be nested within Loop Constructs"_err_en_US,
               parser::ToUpperCaseLetters(nestedBeginName.source.ToString()));
+          endFuseNeeded = false;
         }
       } else {
         missingDoConstruct(beginName, messages_);
+        endFuseNeeded = false;
       }
+      if (endFuseNeeded)
+        continue;
       // If we get here, we either found a loop, or issued an error message.
       return;
     }
     if (nextIt == block.end()) {
-      missingDoConstruct(beginName, messages_);
+      if (endFuseNeeded)
+        missingEndFuse(beginName, messages_);
+      else
+        missingDoConstruct(beginName, messages_);
     }
   }
 

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

@@ -286,10 +286,9 @@ void OmpStructureChecker::Enter(const parser::OpenMPLoopConstruct &x) {
   }
   SetLoopInfo(x);
 
-  auto &optLoopCons = std::get<std::optional<parser::NestedConstruct>>(x.t);
-  if (optLoopCons.has_value()) {
-    if (const auto &doConstruct{
-            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+  auto &loopConsList = std::get<std::list<parser::NestedConstruct>>(x.t);
+  for (auto &loopCons : loopConsList) {
+    if (const auto &doConstruct{std::get_if<parser::DoConstruct>(&loopCons)}) {
       const auto &doBlock{std::get<parser::Block>(doConstruct->t)};
       CheckNoBranching(doBlock, beginName.v, beginName.source);
     }
@@ -315,10 +314,10 @@ const parser::Name OmpStructureChecker::GetLoopIndex(
 }
 
 void OmpStructureChecker::SetLoopInfo(const parser::OpenMPLoopConstruct &x) {
-  auto &optLoopCons = std::get<std::optional<parser::NestedConstruct>>(x.t);
-  if (optLoopCons.has_value()) {
+  auto &loopConsList = std::get<std::list<parser::NestedConstruct>>(x.t);
+  if (loopConsList.size() == 1) {
     if (const auto &loopConstruct{
-            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+            std::get_if<parser::DoConstruct>(&loopConsList.front())}) {
       const parser::DoConstruct *loop{&*loopConstruct};
       if (loop && loop->IsDoNormal()) {
         const parser::Name &itrVal{GetLoopIndex(loop)};
@@ -330,10 +329,10 @@ void OmpStructureChecker::SetLoopInfo(const parser::OpenMPLoopConstruct &x) {
 
 void OmpStructureChecker::CheckLoopItrVariableIsInt(
     const parser::OpenMPLoopConstruct &x) {
-  auto &optLoopCons = std::get<std::optional<parser::NestedConstruct>>(x.t);
-  if (optLoopCons.has_value()) {
+  auto &loopConsList = std::get<std::list<parser::NestedConstruct>>(x.t);
+  for (auto &loopCons : loopConsList) {
     if (const auto &loopConstruct{
-            std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+            std::get_if<parser::DoConstruct>(&loopCons)}) {
 
       for (const parser::DoConstruct *loop{&*loopConstruct}; loop;) {
         if (loop->IsDoNormal()) {
@@ -418,19 +417,20 @@ void OmpStructureChecker::CheckDistLinear(
 
     // Match the loop index variables with the collected symbols from linear
     // clauses.
-    auto &optLoopCons = std::get<std::optional<parser::NestedConstruct>>(x.t);
-    if (optLoopCons.has_value()) {
+    auto &loopConsList = std::get<std::list<parser::NestedConstruct>>(x.t);
+    for (auto &loopCons : loopConsList) {
+      std::int64_t collapseVal_ = collapseVal;
       if (const auto &loopConstruct{
-              std::get_if<parser::DoConstruct>(&*optLoopCons)}) {
+              std::get_if<parser::DoConstruct>(&loopCons)}) {
         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) {
+            collapseVal_--;
+            if (collapseVal_ == 0) {
               break;
             }
           }