[flang-commits] [flang] [flang][OpenMP] Identify DO loops affected by loop-associated construct (PR #191719)

[Krzysztof Parzyszek via flang-commits]

================
@@ -1203,6 +1203,97 @@ std::optional<int64_t> GetMinimumSequenceCount(
   return GetMinimumSequenceCount(std::nullopt, std::nullopt);
 }
 
+/// Collect the DO loops that are affected directly by the given loop
+/// transformation. Not all DO loops nested in the associated nest are
+/// affected by the top-level loop transformation, e.g.
+///
+/// !$omp do collapse(5)                           | [2]
+/// !$omp tile sizes(2, 2)  | [1]                  | <- nest of 4 loops
+/// do i = 1, 10            | <- affected by TILE  |    generated by TILE
+///   do j = 1, 10          | <-                   |
+///     do k = 1, 10                               | <- affected by DO
+///     end do
+///   end do
+/// end do
+///
+/// The two DO loops (i and j) in [1] are affected by the TILE construct.
+/// The k DO loop is affected by the DO construct [2].
+/// For the top-level DO COLLAPSE(5) construct, the k loop is the only
+/// directly affected loop.
+std::optional<std::vector<const parser::DoConstruct *>> CollectAffectedDoLoops(
+    const parser::OpenMPLoopConstruct &x, unsigned version,
+    SemanticsContext *semaCtx) {
+  std::vector<const parser::DoConstruct *> result;
+  const parser::OmpDirectiveSpecification &spec{x.BeginDir()};
+
+  auto [depth, _]{GetAffectedNestDepthWithReason(spec, version, semaCtx)};
+
+  // If the depth is absent, then there is some issue. Leave it alone here,
+  // and let the semantic checks diagnose the problem.
+  if (!depth) {
+    return std::nullopt;
+  }
+  if (*depth.value <= 0) {
+    return result;
+  }
+
+  // The algorithm is to descend down the nest and keep track of intervening
+  // constructs and how many loops they consume and produce. This is similar
+  // to traversing an expression tree to identify the operands to the top-
+  // level operation:
+  //
+  //   ... + + x y z w ...
+  //       ^ ^     ^
+  //       | |     |
+  //       | |     +-- produces 1 value consumed by the first +
+  //       | +-- produces 1 value, but first consumes 2
+  //       +-- consumes 2 operands
+  //
+  // The analogous result here would be "z" as the operand to the first +.
+
+  int64_t produced{0};
+  int64_t consuming{0};
+  int64_t level{*depth.value};
+
+  auto visit{[&](const LoopSequence &nest, auto &&self) -> void {
+    const parser::ExecutionPartConstruct *owner{nest.owner()};
+
+    if (auto *doLoop{parser::Unwrap<parser::DoConstruct>(owner)}) {
+      if (consuming == 0) {
+        result.push_back(doLoop);
+        ++produced;
+      } else {
+        --consuming;
+      }
+    } else if (auto *omp{parser::Unwrap<parser::OpenMPLoopConstruct>(owner)}) {
+      const parser::OmpDirectiveSpecification &ods{omp->BeginDir()};
+      auto [cons, _1]{GetAffectedNestDepthWithReason(ods, version, semaCtx)};
+      auto [prod, _2]{GetGeneratedNestDepthWithReason(ods, version, semaCtx)};
+      if (!cons || !prod) {
+        return;
+      }
+      if (*prod.value <= consuming) {
+        consuming -= *prod.value;
+      } else {
+        produced += (*prod.value - consuming);
+        consuming = 0;
+      }
+      consuming += *cons.value;
+    }
+
+    if (produced < level) {
+      for (const LoopSequence &child : nest.children()) {
----------------
kparzysz wrote:

`depth` is only take from OpenMP constructs.  If this is the only nest in the program, `CollectAffectedDoLoops` won't be called, since it requires OpenMPLoopConstruct as a parameter.

If this is a sub-nest of a loop construct, `depth` will be obtained from the construct and loops in this nest will be visited, provided that their "level" (i.e. how deep they are in the nest) does not exceed the depth.

https://github.com/llvm/llvm-project/pull/191719