[flang-commits] [flang] [flang][OpenMP] Lower DO and SIMD variants in metadirectives (PR #210810)

[via flang-commits]

================
@@ -753,11 +754,253 @@ void OmpStructureChecker::Leave(const parser::OmpDirectiveSpecification &x) {
   }
 }
 
+void OmpStructureChecker::BeginMetadirectiveSelection() {
+  metadirectiveSelectionStarts_.push_back(metadirectiveLoopVariants_.size());
+}
+
+void OmpStructureChecker::EndMetadirectiveSelection(
+    const parser::OmpClauseList &clauses) {
+  CHECK(!metadirectiveSelectionStarts_.empty());
+  std::size_t firstVariant{metadirectiveSelectionStarts_.back()};
+  metadirectiveSelectionStarts_.pop_back();
+  CHECK(!metadirectiveConstructContexts_.empty());
+
+  // Nested dynamic selections form a cross product. Bound it and leave all
+  // variants conservatively reachable if the exact context set grows larger.
+  constexpr std::size_t maxConstructTraitAlternatives{64};
+  auto isInvariantCondition =
+      [&](const MetadirectiveConditionConstraint &condition) {
+        const SomeExpr *expr{GetExpr(context_, *condition.expr)};
+        if (!expr)
+          return false;
+        auto symbols{evaluate::CollectSymbols(*expr)};
+        return !symbols.empty() &&
+            llvm::all_of(symbols, [](const Symbol &symbol) {
+              if (IsNamedConstant(symbol))
+                return true;
+              return IsIntentIn(symbol) && !IsPointer(symbol) &&
+                  !evaluate::IsCoarray(symbol) &&
+                  !symbol.attrs().HasAny({Attr::ASYNCHRONOUS, Attr::VOLATILE});
+            });
+      };
+  auto sameCondition = [&](const MetadirectiveConditionConstraint &left,
+                           const MetadirectiveConditionConstraint &right) {
+    // Repeated conditions within one metadirective are evaluated at the same
+    // selection point. Across nested metadirectives, only correlate expressions
+    // whose values cannot have changed between the two selections.
+    if (left.owner != right.owner &&
+        (!isInvariantCondition(left) || !isInvariantCondition(right)))
+      return false;
+    if (left.expr == right.expr)
+      return true;
+    const SomeExpr *leftExpr{GetExpr(context_, *left.expr)};
+    const SomeExpr *rightExpr{GetExpr(context_, *right.expr)};
+    return leftExpr && rightExpr && *leftExpr == *rightExpr;
+  };
+  auto addCondition =
+      [&](llvm::SmallVectorImpl<MetadirectiveConditionConstraint> &conditions,
+          MetadirectiveConditionConstraint condition) {
+        for (const MetadirectiveConditionConstraint &existing : conditions) {
+          if (sameCondition(existing, condition))
+            return existing.value == condition.value;
+        }
+        conditions.push_back(condition);
+        return true;
+      };
+  auto sameConditions =
+      [&](llvm::ArrayRef<MetadirectiveConditionConstraint> left,
+          llvm::ArrayRef<MetadirectiveConditionConstraint> right) {
+        if (left.size() != right.size())
+          return false;
+        return llvm::all_of(left, [&](const auto &leftCondition) {
+          return llvm::any_of(right, [&](const auto &rightCondition) {
+            return leftCondition.value == rightCondition.value &&
+                sameCondition(leftCondition, rightCondition);
+          });
+        });
+      };
+  auto containsAlternative =
+      [&](llvm::ArrayRef<MetadirectiveConstructAlternative> alternatives,
+          const MetadirectiveConstructAlternative &alternative) {
+        return llvm::any_of(alternatives, [&](const auto &existing) {
+          return existing.traits == alternative.traits &&
+              sameConditions(existing.conditions, alternative.conditions);
+        });
+      };
+
+  llvm::SmallVector<MetadirectiveConstructAlternative, 4>
+      constructTraitAlternatives(1);
+  auto appendAlternatives =
+      [&](llvm::ArrayRef<MetadirectiveConstructAlternative> suffixes) {
+        if (suffixes.empty())
+          return false;
+
+        llvm::SmallVector<MetadirectiveConstructAlternative, 4>
+            combinedAlternatives;
+        for (const MetadirectiveConstructAlternative &prefix :
+            constructTraitAlternatives) {
+          for (const MetadirectiveConstructAlternative &suffix : suffixes) {
+            MetadirectiveConstructAlternative combined{prefix};
+            combined.traits.append(suffix.traits);
+            bool isCompatible{llvm::all_of(suffix.conditions,
+                [&](MetadirectiveConditionConstraint condition) {
+                  return addCondition(combined.conditions, condition);
+                })};
+            if (!isCompatible ||
+                containsAlternative(combinedAlternatives, combined))
+              continue;
+            if (combinedAlternatives.size() == maxConstructTraitAlternatives)
+              return false;
+            combinedAlternatives.push_back(std::move(combined));
+          }
+        }
+        if (combinedAlternatives.empty())
+          return false;
+        constructTraitAlternatives = std::move(combinedAlternatives);
+        return true;
+      };
+
+  // Interleave each enclosing metadirective's selected traits with source
+  // constructs at the metadirective's position in the construct stack.
+  std::size_t metadirectiveContextIndex{0};
+  for (const LoopOrConstruct &item : constructStack_) {
+    if (const auto *construct{
+            std::get_if<const parser::OpenMPConstruct *>(&item)}) {
+      llvm::omp::Directive directive{
+          parser::omp::GetOmpDirectiveName(**construct).v};
+      if (directive == llvm::omp::Directive::OMPD_metadirective) {
+        CHECK(
+            metadirectiveContextIndex < metadirectiveConstructContexts_.size());
+        const MetadirectiveConstructContext &context{
+            metadirectiveConstructContexts_[metadirectiveContextIndex++]};
+        bool isCurrentMetadirective{metadirectiveContextIndex ==
+            metadirectiveConstructContexts_.size()};
+        if (!isCurrentMetadirective &&
+            !appendAlternatives(context.alternatives))
+          return;
+        continue;
+      }
+
+      ConstructTraitSequence sourceTraits;
+      AppendConstructTraitsForDirective(directive, sourceTraits);
+      for (MetadirectiveConstructAlternative &alternative :
+          constructTraitAlternatives)
+        alternative.traits.append(sourceTraits);
+    }
+  }
+
+  struct ReachableVariantPath {
+    const parser::OmpDirectiveSpecification *specification;
+    llvm::SmallVector<MetadirectiveConditionConstraint, 2> conditions;
+  };
+  auto getReachableVariantPaths =
+      [&](const MetadirectiveCandidateSet &candidateSet,
+          const OmpVariantMatchContext &matchContext) {
+        llvm::SmallVector<unsigned, 4> candidates;
+        candidates.reserve(candidateSet.candidates.size());
+        for (unsigned i{0}; i < candidateSet.candidates.size(); ++i)
+          candidates.push_back(i);
+
+        llvm::SmallVector<ReachableVariantPath, 4> paths;
+        llvm::SmallVector<MetadirectiveConditionConstraint, 2> pathConditions;
+        while (true) {
+          std::optional<unsigned> selected{SelectBestMetadirectiveCandidate(
+              candidates, candidateSet.candidates, matchContext)};
+          if (!selected) {
+            paths.push_back({candidateSet.fallback, pathConditions});
+            break;
+          }
+
+          const MetadirectiveCandidate &candidate{
+              candidateSet.candidates[*selected]};
+          if (!candidate.dynamicCondition) {
+            paths.push_back({candidate.spec, pathConditions});
+            break;
+          }
+
+          MetadirectiveConditionConstraint selectedCondition{
+              candidate.dynamicCondition->expr, candidate.conditionShouldBeTrue,
+              &clauses};
+          auto selectedPathConditions{pathConditions};
+          if (addCondition(selectedPathConditions, selectedCondition))
+            paths.push_back(
+                {candidate.spec, std::move(selectedPathConditions)});
+
+          MetadirectiveConditionConstraint failedCondition{
+              candidate.dynamicCondition->expr,
+              !candidate.conditionShouldBeTrue, &clauses};
+          if (!addCondition(pathConditions, failedCondition))
+            break;
+          candidates = GetMetadirectiveElsePathCandidates(
+              *selected, candidates, candidateSet.candidates, context_);
+        }
+        return paths;
+      };
+
+  llvm::SmallVector<const parser::OmpDirectiveSpecification *, 4>
+      reachableVariants;
+  llvm::SmallVector<MetadirectiveConstructAlternative, 2> currentAlternatives;
+  for (const MetadirectiveConstructAlternative &constructAlternative :
+      constructTraitAlternatives) {
+    OmpVariantMatchContext matchContext{context_, constructAlternative.traits};
+    std::optional<MetadirectiveCandidateSet> candidateSet{
+        BuildMetadirectiveCandidateSet(clauses, context_, matchContext)};
+    if (!candidateSet) {
+      // Keep every variant when selection cannot yet model a selector.
+      return;
+    }
+
+    for (const ReachableVariantPath &path :
+        getReachableVariantPaths(*candidateSet, matchContext)) {
+      MetadirectiveConstructAlternative currentAlternative;
+      if (path.specification)
+        AppendConstructTraitsForDirective(
+            path.specification->DirId(), currentAlternative.traits);
+      currentAlternative.conditions = constructAlternative.conditions;
+      bool isCompatible{llvm::all_of(
+          path.conditions, [&](MetadirectiveConditionConstraint condition) {
+            return addCondition(currentAlternative.conditions, condition);
+          })};
+      if (!isCompatible)
+        continue;
+
+      if (!llvm::is_contained(reachableVariants, path.specification))
+        reachableVariants.push_back(path.specification);
+      if (!containsAlternative(currentAlternatives, currentAlternative)) {
+        if (currentAlternatives.size() == maxConstructTraitAlternatives)
+          return;
+        currentAlternatives.push_back(std::move(currentAlternative));
+      }
+    }
+  }
+  if (currentAlternatives.empty())
+    return;
+
+  // Make every reachable replacement and its accumulated path constraints
+  // available to metadirectives in a delimited body.
+  MetadirectiveConstructContext &currentContext{
+      metadirectiveConstructContexts_.back()};
----------------
MattPD wrote:

This context reaches a metadirective in a delimited body, but not one whose loop is a parse-tree sibling. Such a metadirective never reaches `constructStack_`, and `Leave(parser::OmpMetadirectiveDirective)` pops the context before the loop body is walked:

```fortran
subroutine sp(n, a)
  integer :: n, a(n), i
  real :: r
  !$omp metadirective when(implementation={vendor(llvm)}: do) otherwise(nothing)
  do i = 1, n
    !$omp metadirective when(construct={do}: nothing) otherwise(do)
    do r = 1, 2
      a(i) = int(r)
    end do
  end do
end subroutine
```

The inner metadirective reports `The DO loop iteration variable must be of integer type`. Declaring `r` as an integer shows what lowering actually chose: a single `omp.wsloop`, privatizing `i`, with the inner loop emitted as a plain `fir.do_loop`. Lowering selected `nothing` for the inner metadirective, because the enclosing construct is a `do`, so semantics rejected a variant that lowering had already discarded.

`Enter(parser::ExecutionPartConstruct)` in this file already walks from a metadirective to its sibling loop to validate `metadirectiveLoopVariants_`. Would holding the selected context across that same window cover the non-delimited spelling too?

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