[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) {
----------------
kparzysz wrote:

I've added changes to diagnose `cons`/`prod` being std::nullopt.

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