[OpenMP] Add match_{all,any,none} declare variant selector extensions.

By default, all traits in the OpenMP context selector have to match for
it to be acceptable. Though, we sometimes want a single property out of
multiple to match (=any) or no match at all (=none). We offer these
choices as extensions via
  `implementation={extension(match_{all,any,none})}`
to the user. The choice will affect the entire context selector not only
the traits following the match property.

The first user will be D75788. There we can replace
```
  #pragma omp begin declare variant match(device={arch(nvptx64)})
  #define __CUDA__

  #include <__clang_cuda_cmath.h>

  // TODO: Hack until we support an extension to the match clause that allows "or".
  #undef __CLANG_CUDA_CMATH_H__

  #undef __CUDA__
  #pragma omp end declare variant

  #pragma omp begin declare variant match(device={arch(nvptx)})
  #define __CUDA__

  #include <__clang_cuda_cmath.h>

  #undef __CUDA__
  #pragma omp end declare variant
```
with the much simpler
```
  #pragma omp begin declare variant match(device={arch(nvptx, nvptx64)}, implementation={extension(match_any)})
  #define __CUDA__

  #include <__clang_cuda_cmath.h>

  #undef __CUDA__
  #pragma omp end declare variant
```

Reviewed By: mikerice

Differential Revision: https://reviews.llvm.org/D77414

include/llvm/Frontend/OpenMP/OMPContext.h

@@ -154,9 +154,12 @@ struct OMPContext {
 };
 
 /// Return true if \p VMI is applicable in \p Ctx, that is, all traits required
-/// by \p VMI are available in the OpenMP context \p Ctx.
+/// by \p VMI are available in the OpenMP context \p Ctx. If \p DeviceSetOnly is
+/// true, only the device selector set, if present, are checked. Note that we
+/// still honor extension traits provided by the user.
 bool isVariantApplicableInContext(const VariantMatchInfo &VMI,
-                                  const OMPContext &Ctx);
+                                  const OMPContext &Ctx,
+                                  bool DeviceSetOnly = false);
 
 /// Return the index (into \p VMIs) of the variant with the highest score
 /// from the ones applicble in \p Ctx. See llvm::isVariantApplicableInContext.

include/llvm/Frontend/OpenMP/OMPKinds.def

@@ -664,6 +664,9 @@ __OMP_TRAIT_PROPERTY(implementation, vendor, ti)
 __OMP_TRAIT_PROPERTY(implementation, vendor, unknown)
 
 __OMP_TRAIT_SELECTOR(implementation, extension, true)
+__OMP_TRAIT_PROPERTY(implementation, extension, match_all)
+__OMP_TRAIT_PROPERTY(implementation, extension, match_any)
+__OMP_TRAIT_PROPERTY(implementation, extension, match_none)
 
 __OMP_TRAIT_SET(user)
 

lib/Frontend/OpenMP/OMPContext.cpp

@@ -137,52 +137,121 @@ static bool isStrictSubset(const VariantMatchInfo &VMI0,
 
 static int isVariantApplicableInContextHelper(
     const VariantMatchInfo &VMI, const OMPContext &Ctx,
-    SmallVectorImpl<unsigned> *ConstructMatches) {
+    SmallVectorImpl<unsigned> *ConstructMatches, bool DeviceSetOnly) {
+
+  // The match kind determines if we need to match all traits, any of the
+  // traits, or none of the traits for it to be an applicable context.
+  enum MatchKind { MK_ALL, MK_ANY, MK_NONE };
+
+  MatchKind MK = MK_ALL;
+  // Determine the match kind the user wants, "all" is the default and provided
+  // to the user only for completeness.
+  if (VMI.RequiredTraits.test(
+          unsigned(TraitProperty::implementation_extension_match_any)))
+    MK = MK_ANY;
+  if (VMI.RequiredTraits.test(
+          unsigned(TraitProperty::implementation_extension_match_none)))
+    MK = MK_NONE;
+
+  // Helper to deal with a single property that was (not) found in the OpenMP
+  // context based on the match kind selected by the user via
+  // `implementation={extensions(match_[all,any,none])}'
+  auto HandleTrait = [MK](TraitProperty Property,
+                          bool WasFound) -> Optional<bool> /* Result */ {
+    // For kind "any" a single match is enough but we ignore non-matched
+    // properties.
+    if (MK == MK_ANY) {
+      if (WasFound)
+        return true;
+      return None;
+    }
+
+    // In "all" or "none" mode we accept a matching or non-matching property
+    // respectively and move on. We are not done yet!
+    if ((WasFound && MK == MK_ALL) || (!WasFound && MK == MK_NONE))
+      return None;
+
+    // We missed a property, provide some debug output and indicate failure.
+    LLVM_DEBUG({
+      if (MK == MK_ALL)
+        dbgs() << "[" << DEBUG_TYPE << "] Property "
+               << getOpenMPContextTraitPropertyName(Property)
+               << " was not in the OpenMP context but match kind is all.\n";
+      if (MK == MK_NONE)
+        dbgs() << "[" << DEBUG_TYPE << "] Property "
+               << getOpenMPContextTraitPropertyName(Property)
+               << " was in the OpenMP context but match kind is none.\n";
+    });
+    return false;
+  };
 
   for (unsigned Bit : VMI.RequiredTraits.set_bits()) {
     TraitProperty Property = TraitProperty(Bit);
+    if (DeviceSetOnly &&
+        getOpenMPContextTraitSetForProperty(Property) != TraitSet::device)
+      continue;
+
+    // So far all extensions are handled elsewhere, we skip them here as they
+    // are not part of the OpenMP context.
+    if (getOpenMPContextTraitSelectorForProperty(Property) ==
+        TraitSelector::implementation_extension)
+      continue;
 
     bool IsActiveTrait = Ctx.ActiveTraits.test(unsigned(Property));
-    if (!IsActiveTrait) {
-      LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE << "] Property "
-                        << getOpenMPContextTraitPropertyName(Property)
-                        << " was not in the OpenMP context.\n");
-      return false;
-    }
+    Optional<bool> Result = HandleTrait(Property, IsActiveTrait);
+    if (Result.hasValue())
+      return Result.getValue();
   }
 
-  // We could use isSubset here but we also want to record the match locations.
-  unsigned ConstructIdx = 0, NoConstructTraits = Ctx.ConstructTraits.size();
-  for (TraitProperty Property : VMI.ConstructTraits) {
-    assert(getOpenMPContextTraitSetForProperty(Property) ==
-               TraitSet::construct &&
-           "Variant context is ill-formed!");
+  if (!DeviceSetOnly) {
+    // We could use isSubset here but we also want to record the match
+    // locations.
+    unsigned ConstructIdx = 0, NoConstructTraits = Ctx.ConstructTraits.size();
+    for (TraitProperty Property : VMI.ConstructTraits) {
+      assert(getOpenMPContextTraitSetForProperty(Property) ==
+                 TraitSet::construct &&
+             "Variant context is ill-formed!");
 
-    // Verify the nesting.
-    bool FoundInOrder = false;
-    while (!FoundInOrder && ConstructIdx != NoConstructTraits)
-      FoundInOrder = (Ctx.ConstructTraits[ConstructIdx++] == Property);
-    if (ConstructMatches)
-      ConstructMatches->push_back(ConstructIdx - 1);
+      // Verify the nesting.
+      bool FoundInOrder = false;
+      while (!FoundInOrder && ConstructIdx != NoConstructTraits)
+        FoundInOrder = (Ctx.ConstructTraits[ConstructIdx++] == Property);
+      if (ConstructMatches)
+        ConstructMatches->push_back(ConstructIdx - 1);
 
-    if (!FoundInOrder) {
-      LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE << "] Construct property "
-                        << getOpenMPContextTraitPropertyName(Property)
-                        << " was not nested properly.\n");
-      return false;
+      Optional<bool> Result = HandleTrait(Property, FoundInOrder);
+      if (Result.hasValue())
+        return Result.getValue();
+
+      if (!FoundInOrder) {
+        LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE << "] Construct property "
+                          << getOpenMPContextTraitPropertyName(Property)
+                          << " was not nested properly.\n");
+        return false;
+      }
+
+      // TODO: Verify SIMD
     }
 
-    // TODO: Verify SIMD
+    assert(isSubset<TraitProperty>(VMI.ConstructTraits, Ctx.ConstructTraits) &&
+           "Broken invariant!");
+  }
+
+  if (MK == MK_ANY) {
+    LLVM_DEBUG(dbgs() << "[" << DEBUG_TYPE
+                      << "] None of the properties was in the OpenMP context "
+                         "but match kind is any.\n");
+    return false;
   }
 
-  assert(isSubset<TraitProperty>(VMI.ConstructTraits, Ctx.ConstructTraits) &&
-         "Broken invariant!");
   return true;
 }
 
 bool llvm::omp::isVariantApplicableInContext(const VariantMatchInfo &VMI,
-                                             const OMPContext &Ctx) {
-  return isVariantApplicableInContextHelper(VMI, Ctx, nullptr);
+                                             const OMPContext &Ctx,
+                                             bool DeviceSetOnly) {
+  return isVariantApplicableInContextHelper(
+      VMI, Ctx, /* ConstructMatches */ nullptr, DeviceSetOnly);
 }
 
 static APInt getVariantMatchScore(const VariantMatchInfo &VMI,
@@ -267,7 +336,8 @@ int llvm::omp::getBestVariantMatchForContext(
 
     SmallVector<unsigned, 8> ConstructMatches;
     // If the variant is not applicable its not the best.
-    if (!isVariantApplicableInContextHelper(VMI, Ctx, &ConstructMatches))
+    if (!isVariantApplicableInContextHelper(VMI, Ctx, &ConstructMatches,
+                                            /* DeviceSetOnly */ false))
       continue;
     // Check if its clearly not the best.
     APInt Score = getVariantMatchScore(VMI, Ctx, ConstructMatches);