[clang] [flang] [llvm] [clang][flang][OpenMP] Fix context selector matching and scoring (PR #224431)

[via llvm-commits]

================
@@ -1473,27 +1490,78 @@ semantics::omp::OmpVariantMatchContext makeVariantMatchContext(
 }
 
 void collectEnclosingConstructTraits(
-    mlir::Operation *op,
+    AbstractConverter &converter, const pft::Evaluation *evaluation,
     llvm::SmallVectorImpl<llvm::omp::TraitProperty> &constructTraits) {
-  // Collect enclosing OpenMP operations so variants chosen by an outer
-  // metadirective are part of this metadirective's context. For example, an
-  // inner metadirective inside `target` and an outer-selected `parallel` must
-  // be able to match construct={target, parallel}. The final reverse yields
-  // outermost-to-innermost order as required by OMPContext.
-  for (; op; op = op->getParentOp()) {
-    if (mlir::isa<mlir::omp::WsloopOp>(op))
-      constructTraits.push_back(llvm::omp::TraitProperty::construct_for_for);
-    if (mlir::isa<mlir::omp::ParallelOp>(op))
-      constructTraits.push_back(
-          llvm::omp::TraitProperty::construct_parallel_parallel);
-    if (mlir::isa<mlir::omp::TeamsOp>(op))
-      constructTraits.push_back(
-          llvm::omp::TraitProperty::construct_teams_teams);
-    if (mlir::isa<mlir::omp::TargetOp>(op))
-      constructTraits.push_back(
-          llvm::omp::TraitProperty::construct_target_target);
+  const auto *loopControl =
+      converter.getStateStack().getStackTop<LoopControlContext>();
+  // Lastprivate can re-evaluate bounds after lowering the loop body, leaving
+  // a body evaluation current. Use the owning directive's ancestors so the
+  // loop itself is not added before filtering its entered constituents below.
+  if (loopControl)
+    evaluation = &loopControl->evaluation;
+
+  llvm::SmallVector<const OpenMPContextFrame *, 4> frames;
+  converter.getStateStack().stackWalk<OpenMPContextFrame>(
+      [&](OpenMPContextFrame &frame) {
+        frames.push_back(&frame);
+        return mlir::WalkResult::advance();
+      });
+  std::reverse(frames.begin(), frames.end());
+  llvm::SmallVector<bool, 4> usedFrames(frames.size(), false);
+
+  llvm::SmallVector<const pft::Evaluation *, 8> ancestors;
+  for (const pft::Evaluation *parent = evaluation ? evaluation->parentConstruct
+                                                  : nullptr;
+       parent; parent = parent->parentConstruct) {
+    ancestors.push_back(parent);
+  }
+  std::reverse(ancestors.begin(), ancestors.end());
+
+  auto append = [&](llvm::omp::Directive directive) {
+    semantics::omp::AppendDirectiveContextTraits(directive, constructTraits);
+  };
+  for (const pft::Evaluation *ancestor : ancestors) {
+    const auto *omp = ancestor->getIf<parser::OpenMPConstruct>();
+    if (!omp)
+      continue;
+    llvm::omp::Directive directive{parser::omp::GetOmpDirectiveName(*omp).v};
+    // An ancestor supplies the full source context, including constituents
+    // whose bodies also have active frames. Count each construct only once.
+    for (auto [index, frame] : llvm::enumerate(frames))
+      if (&frame->evaluation == ancestor)
+        usedFrames[index] = true;
+    if (directive != llvm::omp::Directive::OMPD_metadirective) {
+      append(directive);
+      continue;
+    }
+    for (const OpenMPContextFrame *frame : frames) {
+      if (&frame->evaluation == ancestor && frame->isReplacement) {
+        append(frame->directive);
+        break;
+      }
+    }
+  }
+
+  // Include entered constituents while their own evaluation is current, e.g.
+  // PARALLEL when lowering the bounds of PARALLEL DO. A selected directive
+  // contributes here only through its entered constituents, so its own clause
+  // expressions have the same context as a directly written directive's.
+  bool insideLoop = false;
+  for (auto [index, frame] : llvm::enumerate(frames)) {
+    if (usedFrames[index] || frame->isReplacement)
+      continue;
+    if (loopControl && &frame->evaluation == &loopControl->evaluation) {
+      // Keep the prefix before the first loop-associated constituent. For
----------------
MattPD wrote:

This prefix comes only from entered frames, but host evaluation runs before the loop directive's frames are entered. The host-evaluated bound of a combined loop directive therefore misses the directive's own leading constituents, while the lastprivate copy-back inside the TARGET region sees them. The two evaluations select different variants, and the copy-back never stores the last value:

```fortran
module m
contains
  pure integer function variant(n)
    integer, intent(in) :: n
    variant = n + 1
  end function
  pure integer function bound(n)
    integer, intent(in) :: n
    !$omp declare variant(variant) match(construct={target, parallel})
    bound = n
  end function
  subroutine s(n, x)
    integer :: n, x, i
    !$omp target parallel do lastprivate(x) map(tofrom: x)
    do i = 1, bound(n)
      x = i
    end do
  end subroutine
end module
program p
  use m
  integer :: x
  x = -7
  call s(5, x)
  print *, x
end program
```

Built with `flang -fopenmp -fopenmp-version=52` and no offload target, it prints `-7` at 89a879d and at the merge base. `bound` sets the trip count, while the copy-back compares against the result of `variant`. In device compilation, both evaluations call `variant`.

With `!$omp target` around a separate `!$omp parallel do`, the host-evaluated bound calls a `construct={target}` variant. The copy-back and device compilation call a separate `construct={parallel}` variant.

Could host evaluation include the loop directive's constituents before its first loop-associated one, or could the copy-back reuse the evaluated loop bound?

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