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

Right. I've done some work to remove emitting any diagnostic messages from the symbol resolution process, and delay that until the "proper" semantic checks.  The utility functions will provide some diagnostics (the "reasons"), but they remain unattached to anything.  They are appended to actual error messages (mostly) in check-omp-loop.cpp.

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