[flang][OpenMP] Canonicalize loops with intervening OpenMP constructs (#169191)

Example based on the gfortran test a.6.1.f90
```
  do 100 i = 1,10
  !$omp do
    do 100 j = 1,10
      call work(i,j)
    100 continue
```

During canonicalization of label-DO loops, if the body of an OpenMP
construct ends with a label, treat the label as ending the construct
itself.

This will also allow handling of cases like
```
  do 100 i = 1, 10
  !$omp atomic write
  100 x = i
```
which we were unable to before.

5 files changed, 141 insertions, 3 deletions

diff --git a/flang/include/flang/Parser/openmp-utils.h b/flang/include/flang/Parser/openmp-utils.h
index d4b739f4c952..b72164e6cef4 100644
--- a/flang/include/flang/Parser/openmp-utils.h
+++ b/flang/include/flang/Parser/openmp-utils.h
@@ -135,6 +135,7 @@ template <typename T> struct IsStatement<Statement<T>> {
 };
 
 std::optional<Label> GetStatementLabel(const ExecutionPartConstruct &x);
+std::optional<Label> GetFinalLabel(const OpenMPConstruct &x);
 
 const OmpObjectList *GetOmpObjectList(const OmpClause &clause);
 
diff --git a/flang/lib/Parser/openmp-utils.cpp b/flang/lib/Parser/openmp-utils.cpp
index 3201e5149e27..ab2ed0641f4c 100644
--- a/flang/lib/Parser/openmp-utils.cpp
+++ b/flang/lib/Parser/openmp-utils.cpp
@@ -15,6 +15,7 @@
 #include "flang/Common/indirection.h"
 #include "flang/Common/template.h"
 #include "flang/Common/visit.h"
+#include "flang/Parser/tools.h"
 
 #include <tuple>
 #include <type_traits>
@@ -77,6 +78,44 @@ std::optional<Label> GetStatementLabel(const ExecutionPartConstruct &x) {
   return GetStatementLabelHelper(x);
 }
 
+static std::optional<Label> GetFinalLabel(const Block &x) {
+  if (!x.empty()) {
+    const ExecutionPartConstruct &last{x.back()};
+    if (auto *omp{Unwrap<OpenMPConstruct>(last)}) {
+      return GetFinalLabel(*omp);
+    } else if (auto *doLoop{Unwrap<DoConstruct>(last)}) {
+      return GetFinalLabel(std::get<Block>(doLoop->t));
+    } else {
+      return GetStatementLabel(x.back());
+    }
+  } else {
+    return std::nullopt;
+  }
+}
+
+std::optional<Label> GetFinalLabel(const OpenMPConstruct &x) {
+  return common::visit(
+      [](auto &&s) -> std::optional<Label> {
+        using TypeS = llvm::remove_cvref_t<decltype(s)>;
+        if constexpr (std::is_same_v<TypeS, OpenMPSectionsConstruct>) {
+          auto &list{std::get<std::list<OpenMPConstruct>>(s.t)};
+          if (!list.empty()) {
+            return GetFinalLabel(list.back());
+          } else {
+            return std::nullopt;
+          }
+        } else if constexpr ( //
+            std::is_same_v<TypeS, OpenMPLoopConstruct> ||
+            std::is_same_v<TypeS, OpenMPSectionConstruct> ||
+            std::is_base_of_v<OmpBlockConstruct, TypeS>) {
+          return GetFinalLabel(std::get<Block>(s.t));
+        } else {
+          return std::nullopt;
+        }
+      },
+      x.u);
+}
+
 const OmpObjectList *GetOmpObjectList(const OmpClause &clause) {
   // Clauses with OmpObjectList as its data member
   using MemberObjectListClauses = std::tuple<OmpClause::Copyin,
diff --git a/flang/lib/Semantics/canonicalize-do.cpp b/flang/lib/Semantics/canonicalize-do.cpp
index ef20cffc3e0c..d98ea48f67f2 100644
--- a/flang/lib/Semantics/canonicalize-do.cpp
+++ b/flang/lib/Semantics/canonicalize-do.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "canonicalize-do.h"
+#include "flang/Parser/openmp-utils.h"
 #include "flang/Parser/parse-tree-visitor.h"
 
 namespace Fortran::parser {
@@ -87,6 +88,12 @@ public:
                 [&](Statement<ActionStmt> &actionStmt) {
                   CanonicalizeIfMatch(block, stack, i, actionStmt);
                 },
+                [&](common::Indirection<OpenMPConstruct> &construct) {
+                  // If the body of the OpenMP construct ends with a label,
+                  // treat the label as ending the construct itself.
+                  CanonicalizeIfMatch(
+                      block, stack, i, omp::GetFinalLabel(construct.value()));
+                },
             },
             executableConstruct->u);
       }
@@ -97,10 +104,14 @@ private:
   template <typename T>
   void CanonicalizeIfMatch(Block &originalBlock, std::vector<LabelInfo> &stack,
       Block::iterator &i, Statement<T> &statement) {
-    if (!stack.empty() && statement.label &&
-        stack.back().label == *statement.label) {
+    CanonicalizeIfMatch(originalBlock, stack, i, statement.label);
+  }
+
+  void CanonicalizeIfMatch(Block &originalBlock, std::vector<LabelInfo> &stack,
+      Block::iterator &i, std::optional<Label> label) {
+    if (!stack.empty() && label && stack.back().label == *label) {
       auto currentLabel{stack.back().label};
-      if constexpr (std::is_same_v<T, common::Indirection<EndDoStmt>>) {
+      if (Unwrap<EndDoStmt>(*i)) {
         std::get<ExecutableConstruct>(i->u).u = Statement<ActionStmt>{
             std::optional<Label>{currentLabel}, ContinueStmt{}};
       }
diff --git a/flang/test/Parser/OpenMP/atomic-label-do.f90 b/flang/test/Parser/OpenMP/atomic-label-do.f90
new file mode 100644
index 000000000000..06197587b2d1
--- /dev/null
+++ b/flang/test/Parser/OpenMP/atomic-label-do.f90
@@ -0,0 +1,39 @@
+!RUN: %flang_fc1 -fdebug-unparse -fopenmp %s | FileCheck --ignore-case --check-prefix="UNPARSE" %s
+!RUN: %flang_fc1 -fdebug-dump-parse-tree -fopenmp %s | FileCheck --check-prefix="PARSE-TREE" %s
+
+subroutine f
+  integer :: i, x
+  do 100 i = 1, 10
+  !$omp atomic write
+  100 x = i
+end
+
+!UNPARSE: SUBROUTINE f
+!UNPARSE:  INTEGER i, x
+!UNPARSE:  DO i=1_4,10_4
+!UNPARSE: !$OMP ATOMIC WRITE
+!UNPARSE:   100  x=i
+!UNPARSE:  END DO
+!UNPARSE: END SUBROUTINE
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> DoConstruct
+!PARSE-TREE: | NonLabelDoStmt
+!PARSE-TREE: | | LoopControl -> LoopBounds
+!PARSE-TREE: | | | Scalar -> Name = 'i'
+!PARSE-TREE: | | | Scalar -> Expr = '1_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '1'
+!PARSE-TREE: | | | Scalar -> Expr = '10_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '10'
+!PARSE-TREE: | Block
+!PARSE-TREE: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPAtomicConstruct
+!PARSE-TREE: | | | OmpBeginDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = atomic
+!PARSE-TREE: | | | | OmpClauseList -> OmpClause -> Write
+!PARSE-TREE: | | | | Flags = None
+!PARSE-TREE: | | | Block
+!PARSE-TREE: | | | | ExecutionPartConstruct -> ExecutableConstruct -> ActionStmt -> AssignmentStmt = 'x=i'
+!PARSE-TREE: | | | | | Variable = 'x'
+!PARSE-TREE: | | | | | | Designator -> DataRef -> Name = 'x'
+!PARSE-TREE: | | | | | Expr = 'i'
+!PARSE-TREE: | | | | | | Designator -> DataRef -> Name = 'i'
+!PARSE-TREE: | EndDoStmt ->
diff --git a/flang/test/Parser/OpenMP/cross-label-do.f90 b/flang/test/Parser/OpenMP/cross-label-do.f90
new file mode 100644
index 000000000000..52ac264756df
--- /dev/null
+++ b/flang/test/Parser/OpenMP/cross-label-do.f90
@@ -0,0 +1,48 @@
+!RUN: %flang_fc1 -fdebug-unparse -fopenmp %s | FileCheck --ignore-case --check-prefix="UNPARSE" %s
+!RUN: %flang_fc1 -fdebug-dump-parse-tree -fopenmp %s | FileCheck --check-prefix="PARSE-TREE" %s
+
+subroutine f00
+  integer :: i, j
+  do 100 i = 1,10
+!$omp do
+    do 100 j = 1,10
+    100 continue
+end
+
+!UNPARSE:  SUBROUTINE f00
+!UNPARSE:   INTEGER i, j
+!UNPARSE:   DO i=1_4,10_4
+!UNPARSE: !$OMP DO
+!UNPARSE:    DO j=1_4,10_4
+!UNPARSE:     100 CONTINUE
+!UNPARSE:    END DO
+!UNPARSE:   END DO
+!UNPARSE:  END SUBROUTINE
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> DoConstruct
+!PARSE-TREE: | NonLabelDoStmt
+!PARSE-TREE: | | LoopControl -> LoopBounds
+!PARSE-TREE: | | | Scalar -> Name = 'i'
+!PARSE-TREE: | | | Scalar -> Expr = '1_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '1'
+!PARSE-TREE: | | | Scalar -> Expr = '10_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '10'
+!PARSE-TREE: | Block
+!PARSE-TREE: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPLoopConstruct
+!PARSE-TREE: | | | OmpBeginLoopDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = do
+!PARSE-TREE: | | | | OmpClauseList ->
+!PARSE-TREE: | | | | Flags = None
+!PARSE-TREE: | | | Block
+!PARSE-TREE: | | | | ExecutionPartConstruct -> ExecutableConstruct -> DoConstruct
+!PARSE-TREE: | | | | | NonLabelDoStmt
+!PARSE-TREE: | | | | | | LoopControl -> LoopBounds
+!PARSE-TREE: | | | | | | | Scalar -> Name = 'j'
+!PARSE-TREE: | | | | | | | Scalar -> Expr = '1_4'
+!PARSE-TREE: | | | | | | | | LiteralConstant -> IntLiteralConstant = '1'
+!PARSE-TREE: | | | | | | | Scalar -> Expr = '10_4'
+!PARSE-TREE: | | | | | | | | LiteralConstant -> IntLiteralConstant = '10'
+!PARSE-TREE: | | | | | Block
+!PARSE-TREE: | | | | | | ExecutionPartConstruct -> ExecutableConstruct -> ActionStmt -> ContinueStmt
+!PARSE-TREE: | | | | | EndDoStmt ->
+!PARSE-TREE: | EndDoStmt ->