[flang][OpenMP] Fix metadirective loop semantic checks (#208844)

This is a follow-up to https://github.com/llvm/llvm-project/pull/207088.

Post-merge review feedback:

* Rename mayVariantBeSelected to MayVariantBeSelected and move it to
  the shared OpenMP semantic utilities so declare-variant checking can
  reuse it.
* Explain why a WHEN selector is recorded only when its modifier list
  contains exactly one element.
* Diagnose loop-associated variants that end a declaration-only
  program unit, such as a module with no execution part.

Additional correctness fixes:

* Handle consecutive metadirectives and modules or submodules with
  contained procedures without silently dropping an earlier
  loop-associated variant.
* Prevent a loop-associated variant in an interface, BLOCK construct,
  conditional branch, or nested body from being checked against an
  unrelated loop outside that scope.
* Avoid false or missing loop diagnostics when match_any or match_none
  combines target selectors with user or construct selectors.

Assisted with codex.

flang/include/flang/Semantics/openmp-utils.h

@@ -223,6 +223,16 @@ private:
   std::string features_;
 };
 
+/// True if a variant guarded by \p selector may be selected in the current
+/// compilation context.
+///
+/// False only if the complete selector is provably inapplicable. A null
+/// \p selector, or one using an unsupported feature (see
+/// \c FindUnsupportedSelectorFeature), is treated as possibly selectable.
+bool MayVariantBeSelected(
+    const parser::traits::OmpContextSelectorSpecification *selector,
+    SemanticsContext &context, OmpVariantMatchContext &matchContext);
+
 std::vector<SomeExpr> GetTopLevelDesignators(const SomeExpr &expr);
 const SomeExpr *HasStorageOverlap(
     const SomeExpr &base, llvm::ArrayRef<SomeExpr> exprs);

flang/lib/Semantics/check-omp-structure.cpp

@@ -75,6 +75,16 @@ void OmpStructureChecker::Enter(const parser::ProgramUnit &) { //
   declareVariantPairs_.clear();
 }
 
+void OmpStructureChecker::Leave(const parser::ProgramUnit &) {
+  if (!metadirectiveLoopVariants_.empty()) {
+    // A declaration-only unit (module, submodule, or block data) has no
+    // execution part to follow the metadirective, so its loop-associated
+    // variants were never validated. A subprogram validates them while
+    // scanning the execution part, leaving none pending here.
+    CheckMetadirectiveVariantsWithoutLoop();
+  }
+}
+
 void OmpStructureChecker::Enter(const parser::MainProgram &x) {
   using StatementProgramStmt = parser::Statement<parser::ProgramStmt>;
   if (auto &stmt{std::get<std::optional<StatementProgramStmt>>(x.t)}) {
@@ -187,12 +197,38 @@ void OmpStructureChecker::Enter(const parser::EndMpSubprogramStmt &x) {
   scopeStack_.pop_back();
 }
 
+void OmpStructureChecker::BeginMetadirectiveVariantScope() {
+  metadirectiveVariantScopeStarts_.push_back(metadirectiveLoopVariants_.size());
+}
+
+void OmpStructureChecker::EndMetadirectiveVariantScope() {
+  CHECK(!metadirectiveVariantScopeStarts_.empty());
+  std::size_t firstVariant{metadirectiveVariantScopeStarts_.back()};
+  metadirectiveVariantScopeStarts_.pop_back();
+  if (firstVariant < metadirectiveLoopVariants_.size()) {
+    // Diagnose variants that were recorded in this scope but not consumed by
+    // one of its executable constructs, preserving variants from an enclosing
+    // scope.
+    CheckMetadirectiveVariantsWithoutLoop(firstVariant);
+  }
+}
+
+void OmpStructureChecker::Enter(const parser::Block &) {
+  BeginMetadirectiveVariantScope();
+}
+
+void OmpStructureChecker::Leave(const parser::Block &) {
+  EndMetadirectiveVariantScope();
+}
+
 void OmpStructureChecker::Enter(const parser::BlockConstruct &x) {
+  BeginMetadirectiveVariantScope();
   auto &endBlockStmt{std::get<parser::Statement<parser::EndBlockStmt>>(x.t)};
   scopeStack_.push_back(&context_.FindScope(endBlockStmt.source));
 }
 
 void OmpStructureChecker::Leave(const parser::BlockConstruct &x) {
+  EndMetadirectiveVariantScope();
   scopeStack_.pop_back();
 }
 
@@ -204,6 +240,23 @@ void OmpStructureChecker::Enter(const parser::ModuleSubprogram &) {
   ClearLabels();
 }
 
+void OmpStructureChecker::Enter(const parser::ModuleSubprogramPart &) {
+  if (!metadirectiveLoopVariants_.empty()) {
+    // The enclosing module or submodule has no execution part. Diagnose its
+    // pending loop-associated variants before a contained procedure starts a
+    // new specification part and resets the worklist.
+    CheckMetadirectiveVariantsWithoutLoop();
+  }
+}
+
+void OmpStructureChecker::Enter(const parser::InterfaceBody &) {
+  BeginMetadirectiveVariantScope();
+}
+
+void OmpStructureChecker::Leave(const parser::InterfaceBody &) {
+  EndMetadirectiveVariantScope();
+}
+
 void OmpStructureChecker::Enter(const parser::SpecificationPart &) {
   // Clear pending metadirective loop variants at the start of a program unit.
   // An empty partStack_ marks the unit's top-level specification part, so a

flang/lib/Semantics/check-omp-structure.h

@@ -60,6 +60,7 @@ public:
   OmpStructureChecker(SemanticsContext &context);
 
   void Enter(const parser::ProgramUnit &);
+  void Leave(const parser::ProgramUnit &);
   void Enter(const parser::MainProgram &);
   void Leave(const parser::MainProgram &);
   void Enter(const parser::BlockData &);
@@ -74,10 +75,15 @@ public:
   void Enter(const parser::EndFunctionStmt &);
   void Enter(const parser::MpSubprogramStmt &);
   void Enter(const parser::EndMpSubprogramStmt &);
+  void Enter(const parser::Block &);
+  void Leave(const parser::Block &);
   void Enter(const parser::BlockConstruct &);
   void Leave(const parser::BlockConstruct &);
   void Enter(const parser::InternalSubprogram &);
   void Enter(const parser::ModuleSubprogram &);
+  void Enter(const parser::ModuleSubprogramPart &);
+  void Enter(const parser::InterfaceBody &);
+  void Leave(const parser::InterfaceBody &);
 
   void Enter(const parser::SpecificationPart &);
   void Leave(const parser::SpecificationPart &);
@@ -282,8 +288,11 @@ private:
   void CheckScanModifier(const parser::OmpClause::Reduction &x);
   void CheckDistLinear(const parser::OpenMPLoopConstruct &x);
 
+  void BeginMetadirectiveVariantScope();
+  void EndMetadirectiveVariantScope();
+
   // check-omp-variant.cpp
-  void CheckMetadirectiveVariantsWithoutLoop();
+  void CheckMetadirectiveVariantsWithoutLoop(std::size_t firstVariant = 0);
   void CheckOmpDeclareVariantDirective(
       const parser::OmpDeclareVariantDirective &);
   void CheckDeclareVariantUserConditions(const parser::OmpContextSelector &);
@@ -497,6 +506,7 @@ private:
     const parser::OmpDirectiveSpecification *spec;
   };
   std::vector<MetadirectiveLoopVariant> metadirectiveLoopVariants_;
+  std::vector<std::size_t> metadirectiveVariantScopeStarts_;
   const parser::traits::OmpContextSelectorSpecification *currentWhenSelector_{
       nullptr};
 

flang/lib/Semantics/check-omp-variant.cpp

@@ -45,7 +45,10 @@ void OmpStructureChecker::Enter(const parser::OmpClause::When &x) {
   OmpVerifyModifiers(
       x.v, llvm::omp::OMPC_when, GetContext().clauseSource, context_);
   // Record this WHEN clause's context selector so the variant directive it
-  // controls can be paired with it for static-applicability matching.
+  // controls can be paired with it for static-applicability matching. A
+  // well-formed WHEN clause has exactly one modifier, its context selector;
+  // pair it only in that case, which also makes front() safe. Any other count
+  // is malformed and already diagnosed by OmpVerifyModifiers above.
   if (const auto &modifiers{std::get<0>(x.v.t)};
       modifiers && modifiers->size() == 1) {
     currentWhenSelector_ =
@@ -603,48 +606,12 @@ void OmpStructureChecker::Leave(const parser::OmpDirectiveSpecification &x) {
 
 void OmpStructureChecker::Enter(const parser::OmpMetadirectiveDirective &x) {
   EnterDirectiveNest(MetadirectiveNest);
-  metadirectiveLoopVariants_.clear();
 }
 
 void OmpStructureChecker::Leave(const parser::OmpMetadirectiveDirective &) {
   ExitDirectiveNest(MetadirectiveNest);
 }
 
-// Return true if the variant guarded by `selector` might be selected, so its
-// associated loop must still be checked. Skip it only when it is provably
-// unselectable.
-static bool mayVariantBeSelected(
-    const parser::traits::OmpContextSelectorSpecification *selector,
-    SemanticsContext &context, OmpVariantMatchContext &matchContext) {
-  if (!selector ||
-      FindUnsupportedSelectorFeature(*selector, context) !=
-          UnsupportedSelectorFeature::None) {
-    return true;
-  }
-  llvm::omp::VariantMatchInfo vmi;
-  (void)MakeVariantMatchInfo(vmi, *selector, context);
-  const auto &required{vmi.RequiredTraits};
-  using TP = llvm::omp::TraitProperty;
-  // In the default "match all" mode, a required trait that can never match
-  // (e.g. a compile-time-false condition or an invalid device/implementation
-  // property) makes the variant unselectable. match_any and match_none tolerate
-  // such traits, so only reject variants in the default mode.
-  bool matchAll{
-      !required.test(unsigned(TP::implementation_extension_match_any)) &&
-      !required.test(unsigned(TP::implementation_extension_match_none))};
-  if (matchAll &&
-      (required.test(unsigned(TP::user_condition_false)) ||
-          required.test(unsigned(TP::invalid)))) {
-    return false;
-  }
-  // Matching a device or implementation selector requires a known target.
-  if (context.targetTriple().empty()) {
-    return true;
-  }
-  return llvm::omp::isVariantApplicableInContext(
-      vmi, matchContext, /*DeviceOrImplementationSetOnly=*/true);
-}
-
 // Check a loop-associated metadirective's variants against the loop nest they
 // apply to. The nest is not attached to the directive in the parse tree. It is
 // the next executable construct, either a following sibling or the first
@@ -710,7 +677,7 @@ void OmpStructureChecker::Enter(const parser::ExecutionPartConstruct &x) {
     const parser::OmpDirectiveSpecification *spec{variant.spec};
     // Skip variants that can never be selected on this compilation target so
     // that their associated loop is not diagnosed.
-    if (!mayVariantBeSelected(variant.selector, context_, matchContext)) {
+    if (!MayVariantBeSelected(variant.selector, context_, matchContext)) {
       continue;
     }
     auto assoc{llvm::omp::getDirectiveAssociation(spec->DirId())};
@@ -748,16 +715,24 @@ void OmpStructureChecker::Enter(const parser::ExecutionPartConstruct &x) {
 // loop nest, either because the metadirective is the last construct in the
 // execution part or because a non-loop construct follows it. Variants that
 // cannot be selected on this target are skipped.
-void OmpStructureChecker::CheckMetadirectiveVariantsWithoutLoop() {
+void OmpStructureChecker::CheckMetadirectiveVariantsWithoutLoop(
+    std::size_t firstVariant) {
+  CHECK(firstVariant <= metadirectiveLoopVariants_.size());
   std::vector<MetadirectiveLoopVariant> variants;
-  variants.swap(metadirectiveLoopVariants_);
+  if (firstVariant == 0) {
+    variants.swap(metadirectiveLoopVariants_);
+  } else {
+    auto first{metadirectiveLoopVariants_.begin() + firstVariant};
+    variants.assign(first, metadirectiveLoopVariants_.end());
+    metadirectiveLoopVariants_.erase(first, metadirectiveLoopVariants_.end());
+  }
 
   OmpVariantMatchContext matchContext{context_};
   const auto MsgShouldContainDoOr{
       "This construct should contain a DO-loop or a loop-%s-generating construct"_err_en_US};
 
   for (const MetadirectiveLoopVariant &variant : variants) {
-    if (!mayVariantBeSelected(variant.selector, context_, matchContext)) {
+    if (!MayVariantBeSelected(variant.selector, context_, matchContext)) {
       continue;
     }
     auto assoc{llvm::omp::getDirectiveAssociation(variant.spec->DirId())};

flang/lib/Semantics/openmp-utils.cpp

@@ -2420,6 +2420,82 @@ std::optional<DynamicUserCondition> MakeVariantMatchInfo(
   return dynamicCond;
 }
 
+bool MayVariantBeSelected(
+    const parser::traits::OmpContextSelectorSpecification *selector,
+    SemanticsContext &context, OmpVariantMatchContext &matchContext) {
+  if (!selector ||
+      FindUnsupportedSelectorFeature(*selector, context) !=
+          UnsupportedSelectorFeature::None) {
+    return true;
+  }
+  llvm::omp::VariantMatchInfo vmi;
+  (void)MakeVariantMatchInfo(vmi, *selector, context);
+  const auto &required{vmi.RequiredTraits};
+  using TP = llvm::omp::TraitProperty;
+
+  enum class MatchKind { All, Any, None };
+  MatchKind matchKind{MatchKind::All};
+  if (required.test(unsigned(TP::implementation_extension_match_any))) {
+    matchKind = MatchKind::Any;
+  }
+  // Match-none takes precedence over match-any when both are present, matching
+  // isVariantApplicableInContextHelper.
+  if (required.test(unsigned(TP::implementation_extension_match_none))) {
+    matchKind = MatchKind::None;
+  }
+
+  bool userTrue{required.test(unsigned(TP::user_condition_true))};
+  bool userUnknown{required.test(unsigned(TP::user_condition_unknown))};
+  bool userFalse{required.test(unsigned(TP::user_condition_false))};
+  bool invalid{required.test(unsigned(TP::invalid))};
+
+  // The target-only LLVM matcher below skips user and construct traits while
+  // retaining the global match kind. Account for those skipped traits first;
+  // otherwise an empty filtered set incorrectly fails match-any and satisfies
+  // match-none.
+  switch (matchKind) {
+  case MatchKind::All:
+    if (userFalse || invalid) {
+      return false;
+    }
+    break;
+  case MatchKind::Any:
+    if (userTrue || userUnknown || !vmi.ConstructTraits.empty()) {
+      return true;
+    }
+    break;
+  case MatchKind::None:
+    if (userTrue) {
+      return false;
+    }
+    break;
+  }
+
+  // Without a target triple, do not reject device or implementation traits.
+  if (context.targetTriple().empty()) {
+    if (matchKind != MatchKind::Any) {
+      return true;
+    }
+    // No skipped trait can satisfy match-any here. Keep the selector
+    // conservatively selectable if it contains a device or implementation
+    // trait; otherwise no trait can satisfy match-any.
+    for (unsigned bit : required.set_bits()) {
+      TP property{static_cast<TP>(bit)};
+      llvm::omp::TraitSet set{
+          llvm::omp::getOpenMPContextTraitSetForProperty(property)};
+      if ((set == llvm::omp::TraitSet::device ||
+              set == llvm::omp::TraitSet::implementation) &&
+          llvm::omp::getOpenMPContextTraitSelectorForProperty(property) !=
+              llvm::omp::TraitSelector::implementation_extension) {
+        return true;
+      }
+    }
+    return false;
+  }
+  return llvm::omp::isVariantApplicableInContext(
+      vmi, matchContext, /*DeviceOrImplementationSetOnly=*/true);
+}
+
 OmpVariantMatchContext::OmpVariantMatchContext(bool isDeviceCompilation,
     llvm::Triple targetTriple, llvm::Triple targetOffloadTriple,
     std::string targetFeatures,

flang/test/Semantics/OpenMP/metadirective-loop-applicability.f90

@@ -116,3 +116,18 @@ subroutine f09(n, a)
     end do
   end do
 end subroutine
+
+! MATCH_ANY is satisfied by a compile-time-true user condition, so its
+! loop-associated variant still requires a loop.
+subroutine f10()
+  logical, parameter :: use_variant = .true.
+  !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+  !$omp metadirective when(user={condition(use_variant)}, implementation={extension(match_any)}: do) default(nothing)
+end subroutine
+
+! MATCH_NONE is not satisfied when the user condition is true, so its
+! loop-associated variant cannot be selected and needs no loop.
+subroutine f11()
+  logical, parameter :: use_variant = .true.
+  !$omp metadirective when(user={condition(use_variant)}, implementation={extension(match_none)}: do) default(nothing)
+end subroutine

flang/test/Semantics/OpenMP/metadirective-loop-nest.f90

@@ -152,3 +152,126 @@ end subroutine
 subroutine no_loop_dead_variant()
   !$omp metadirective when(device={kind(nohost)}: do) default(nothing)
 end subroutine
+
+! A loop-associated variant in a declarative context (e.g. a module
+! specification part) also has no loop nest to associate with.
+module no_loop_in_module
+  implicit none
+  !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+  !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+end module
+
+! Starting a contained procedure must not discard a pending variant from the
+! enclosing module specification part.
+module no_loop_in_module_with_contains
+  implicit none
+  !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+  !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+contains
+  subroutine contained()
+  end subroutine
+end module
+
+! A later specification-part metadirective must not discard an earlier one.
+module no_loop_before_another_metadirective
+  implicit none
+  !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+  !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+  !$omp metadirective when(implementation={vendor(llvm)}: parallel) default(nothing)
+end module
+
+! A loop in the enclosing subprogram cannot satisfy a variant from an
+! interface body, which has no execution part of its own.
+subroutine no_loop_in_interface_body(n, a)
+  integer :: n, a(n), i
+  interface
+    subroutine iface()
+      !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+      !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+    end subroutine
+  end interface
+  do i = 1, n
+    a(i) = i
+  end do
+end subroutine
+
+! Checking an interface body must preserve variants pending in the enclosing
+! subprogram for its execution part.
+subroutine no_loop_in_interface_body_preserves_outer(n, a)
+  integer :: n, a(n), i
+  !ERROR: This construct requires a perfect nest of depth 2, but the associated nest is a perfect nest of depth 1
+  !BECAUSE: COLLAPSE clause was specified with argument 2
+  !$omp metadirective when(implementation={vendor(llvm)}: do collapse(2)) default(nothing)
+  interface
+    subroutine iface()
+      !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+      !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+    end subroutine
+  end interface
+  do i = 1, n
+    a(i) = i
+  end do
+end subroutine
+
+! A loop after an empty BLOCK cannot satisfy a variant from the BLOCK's
+! specification part.
+subroutine no_loop_in_empty_block(n, a)
+  integer :: n, a(n), i
+  block
+    !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+    !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+  end block
+  do i = 1, n
+    a(i) = i
+  end do
+end subroutine
+
+! A loop in the BLOCK itself remains its associated loop.
+subroutine loop_in_block(n, a)
+  integer :: n, a(n), i
+  block
+    !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+    do i = 1, n
+      a(i) = i
+    end do
+  end block
+end subroutine
+
+! A loop in a mutually exclusive IF branch cannot satisfy the variant.
+subroutine no_loop_across_if_branch(n, a, flag)
+  integer :: n, a(n), i
+  logical :: flag
+  if (flag) then
+    !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+    !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+  else
+    do i = 1, n
+      a(i) = i
+    end do
+  end if
+end subroutine
+
+! A loop after a nested executable block cannot satisfy its variant.
+subroutine no_loop_after_if(n, a, flag)
+  integer :: n, a(n), i
+  logical :: flag
+  if (flag) then
+    !ERROR: This construct should contain a DO-loop or a loop-nest-generating construct
+    !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+  end if
+  do i = 1, n
+    a(i) = i
+  end do
+end subroutine
+
+! A loop in the same IF branch remains the associated loop.
+subroutine loop_in_if_branch(n, a, flag)
+  integer :: n, a(n), i
+  logical :: flag
+  if (flag) then
+    !$omp metadirective when(implementation={vendor(llvm)}: do) default(nothing)
+    do i = 1, n
+      a(i) = i
+    end do
+  end if
+end subroutine

flang/unittests/Semantics/OpenMPUtils.cpp

@@ -100,6 +100,53 @@ protected:
   }
 };
 
+class ContextSelectorFinder {
+public:
+  template <typename T> bool Pre(const T &) { return true; }
+  template <typename T> void Post(const T &) {}
+
+  bool Pre(
+      const parser::traits::OmpContextSelectorSpecification &contextSelector) {
+    selector = &contextSelector;
+    return false;
+  }
+
+  const parser::traits::OmpContextSelectorSpecification *selector{nullptr};
+};
+
+TEST_F(OpenMPUtilsTest, MayVariantBeSelectedMatchAnyWithoutTargetTriple) {
+  *inputFileOs << R"(
+      integer :: i
+      !$omp metadirective when(implementation={vendor(gnu), extension(match_any)}: do) otherwise(nothing)
+      do i = 1, 10
+      end do
+      end
+  )";
+  inputFileOs.reset();
+
+  compInst.getInvocation().getFrontendOpts().programAction = ParseSyntaxOnly;
+
+  bool success{executeCompilerInvocation(&compInst)};
+  ASSERT_TRUE(success);
+
+  std::optional<parser::Program> &parseTree{compInst.getParsing().parseTree()};
+  ASSERT_TRUE(parseTree.has_value());
+
+  ContextSelectorFinder finder;
+  parser::Walk(*parseTree, finder);
+  ASSERT_NE(finder.selector, nullptr);
+
+  semantics::SemanticsContext &semanticsContext{compInst.getSemanticsContext()};
+  semantics::omp::OmpVariantMatchContext matchContext{semanticsContext};
+  EXPECT_FALSE(semantics::omp::MayVariantBeSelected(
+      finder.selector, semanticsContext, matchContext));
+
+  semanticsContext.set_targetTriple("");
+  semantics::omp::OmpVariantMatchContext noTargetMatchContext{semanticsContext};
+  EXPECT_TRUE(semantics::omp::MayVariantBeSelected(
+      finder.selector, semanticsContext, noTargetMatchContext));
+}
+
 TEST_F(OpenMPUtilsTest, AffectedNestDepthNoClauses) {
   // Populate the input file with the pre-defined input and flush it.
   *inputFileOs << R"(