[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.

1 files changed, 14 insertions, 3 deletions

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{}};
       }