[OpenMP][NFC] Reduce instantiation time with different data structure

[jdoerfert]

Summary

See rational here: https://reviews.llvm.org/D71847#1922648

[jdoerfert]

Tests pass fine, @rnk, let me know if this solves the issue you saw.

[rnk]

Thanks, yes, this is what I had in mind.

llvm/include/llvm/Frontend/OpenMP/OMPContext.h

49

LLVM doesn't use the all uppercase naming style for enums, I'd use something like Count, Num, or Max.

One technique you could use to avoid the sentinel enum is to make an alias for the last enum, as is done for builtin types here:
https://github.com/llvm/llvm-project/blob/master/clang/include/clang/AST/Type.h#L2514
Then you won't have to handle NUM_PROPERTIES below, and you would use TraitProperty::Max + 1 for the BitVector.

152

This is the same instantiation, so we don't get any savings unless the same technique is applied here. I started by trying to change this type, but it was challenging, so I raised the issue to see if you'd be OK with it.

llvm/lib/Frontend/OpenMP/OMPContext.cpp

132

Worth fixing this formatting nit

[jdoerfert]

Addressed comments

[rnk]

lgtm, as long as you are OK with the casts as the code author.

llvm/lib/Frontend/OpenMP/OMPContext.cpp

130–132

I expected BitVector to expose an API to do this without iterating the bits individually, but I can't find one. test seems close, but not quite right. Oh well.

llvm/include/llvm/Frontend/OpenMP/OMPContext.h

 /// contexts as used by `[begin/end] declare variant` and `metadirective`.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OPENMP_CONTEXT_H
 #define LLVM_OPENMP_CONTEXT_H
 
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 namespace llvm {
 namespace omp {
 

 enum class TraitSelector {
 #define OMP_TRAIT_SELECTOR(Enum, ...) Enum,
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 };
 
 /// IDs for all OpenMP context trait properties (host/gpu/bsc/llvm/...)
 enum class TraitProperty {
 #define OMP_TRAIT_PROPERTY(Enum, ...) Enum,
+#define OMP_LAST_TRAIT_PROPERTY(Enum) Last = Enum
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 };
 
 /// Parse \p Str and return the trait set it matches or TraitSet::invalid.
 TraitSet getOpenMPContextTraitSetKind(StringRef Str);
 
 /// Return the trait set for which \p Selector is a selector.
 TraitSet getOpenMPContextTraitSetForSelector(TraitSelector Selector);
 
 /// Return the trait set for which \p Property is a property.
void addTrait(TraitProperty Property, APInt *Score = nullptr) {
     addTrait(getOpenMPContextTraitSetForProperty(Property), Property, Score);
   }
   /// Add the trait \p Property which is in set \p Set to the required trait
   /// set. If \p Score is not null, it recorded as well. If \p Set is the
   /// `construct` set it is recorded in-order in the ConstructTraits as well.
   void addTrait(TraitSet Set, TraitProperty Property, APInt *Score = nullptr) {
     if (Score)
       ScoreMap[Property] = *Score;
-    RequiredTraits.insert(Property);
+    RequiredTraits.set(unsigned(Property));
     if (Set == TraitSet::construct)
       ConstructTraits.push_back(Property);
   }
 
-  SmallSet<TraitProperty, 8> RequiredTraits;
+  BitVector RequiredTraits = BitVector(unsigned(TraitProperty::Last) + 1);
   SmallVector<TraitProperty, 8> ConstructTraits;
   SmallDenseMap<TraitProperty, APInt> ScoreMap;
 };
 
 /// The context for a source location is made up of active property traits,
 /// e.g., device={kind(host)}, and constructs traits which describe the nesting
 /// in OpenMP constructs at the location.
 struct OMPContext {
   OMPContext(bool IsDeviceCompilation, Triple TargetTriple);
 
   void addTrait(TraitProperty Property) {
     addTrait(getOpenMPContextTraitSetForProperty(Property), Property);
   }
   void addTrait(TraitSet Set, TraitProperty Property) {
-    ActiveTraits.insert(Property);
+    ActiveTraits.set(unsigned(Property));
     if (Set == TraitSet::construct)
       ConstructTraits.push_back(Property);
   }
 
-  SmallSet<TraitProperty, 8> ActiveTraits;
+  BitVector ActiveTraits = BitVector(unsigned(TraitProperty::Last) + 1);
   SmallVector<TraitProperty, 8> ConstructTraits;
 };
 
 /// 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.
 bool isVariantApplicableInContext(const VariantMatchInfo &VMI,
                                   const OMPContext &Ctx);
 

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 #define OMP_TRAIT_SET(Enum, Str)
 #endif
 #ifndef OMP_TRAIT_SELECTOR
 #define OMP_TRAIT_SELECTOR(Enum, TraitSetEnum, Str, RequiresProperty)
 #endif
 #ifndef OMP_TRAIT_PROPERTY
 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)
 #endif
+#ifndef OMP_LAST_TRAIT_PROPERTY
+#define OMP_LAST_TRAIT_PROPERTY(Enum)
+#endif
 
 #define __OMP_TRAIT_SET(Name) OMP_TRAIT_SET(Name, #Name)
 #define __OMP_TRAIT_SELECTOR(TraitSet, Name, RequiresProperty)                 \
   OMP_TRAIT_SELECTOR(TraitSet##_##Name, TraitSet, #Name, RequiresProperty)
 #define __OMP_TRAIT_SELECTOR_AND_PROPERTY(TraitSet, Name)                      \
   OMP_TRAIT_SELECTOR(TraitSet##_##Name, TraitSet, #Name, false)                \
   OMP_TRAIT_PROPERTY(TraitSet##_##Name##_##Name, TraitSet, TraitSet##_##Name,  \
                      #Name)

 #endif
 
 __OMP_REQUIRES_TRAIT(unified_address)
 __OMP_REQUIRES_TRAIT(unified_shared_memory)
 __OMP_REQUIRES_TRAIT(reverse_offload)
 __OMP_REQUIRES_TRAIT(dynamic_allocators)
 __OMP_REQUIRES_TRAIT(atomic_default_mem_order)
 
+OMP_LAST_TRAIT_PROPERTY(
+    implementation_atomic_default_mem_order_atomic_default_mem_order)
+
 #undef __OMP_TRAIT_SELECTOR_AND_PROPERTY
 #undef OMP_TRAIT_SELECTOR
 #undef __OMP_TRAIT_SELECTOR
 #undef OMP_TRAIT_PROPERTY
+#undef OMP_LAST_TRAIT_PROPERTY
 #undef __OMP_TRAIT_PROPERTY
 #undef __OMP_REQUIRES_TRAIT
 #undef OMP_REQUIRES_TRAIT
 ///}

llvm/lib/Frontend/OpenMP/OMPContext.cpp

 #define DEBUG_TYPE "openmp-ir-builder"
 
 using namespace llvm;
 using namespace omp;
 
 OMPContext::OMPContext(bool IsDeviceCompilation, Triple TargetTriple) {
   // Add the appropriate device kind trait based on the triple and the
   // IsDeviceCompilation flag.
-  ActiveTraits.insert(IsDeviceCompilation ? TraitProperty::device_kind_nohost
-                                          : TraitProperty::device_kind_host);
+  ActiveTraits.set(unsigned(IsDeviceCompilation
+                                ? TraitProperty::device_kind_nohost
+                                : TraitProperty::device_kind_host));
   switch (TargetTriple.getArch()) {
   case Triple::arm:
   case Triple::armeb:
   case Triple::aarch64:
   case Triple::aarch64_be:
   case Triple::aarch64_32:
   case Triple::mips:
   case Triple::mipsel:
   case Triple::mips64:
   case Triple::mips64el:
   case Triple::ppc:
   case Triple::ppc64:
   case Triple::ppc64le:
   case Triple::x86:
   case Triple::x86_64:
-    ActiveTraits.insert(TraitProperty::device_kind_cpu);
+    ActiveTraits.set(unsigned(TraitProperty::device_kind_cpu));
     break;
   case Triple::amdgcn:
   case Triple::nvptx:
   case Triple::nvptx64:
-    ActiveTraits.insert(TraitProperty::device_kind_gpu);
+    ActiveTraits.set(unsigned(TraitProperty::device_kind_gpu));
     break;
   default:
     break;
   }
 
   // Add the appropriate device architecture trait based on the triple.
 #define OMP_TRAIT_PROPERTY(Enum, TraitSetEnum, TraitSelectorEnum, Str)         \
   if (TraitSelector::TraitSelectorEnum == TraitSelector::device_arch)          \
     if (TargetTriple.getArch() == TargetTriple.getArchTypeForLLVMName(Str))    \
-      ActiveTraits.insert(TraitProperty::Enum);
+      ActiveTraits.set(unsigned(TraitProperty::Enum));
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 
   // TODO: What exactly do we want to see as device ISA trait?
   //       The discussion on the list did not seem to have come to an agreed
   //       upon solution.
 
   // LLVM is the "OpenMP vendor" but we could also interpret vendor as the
   // target vendor.
-  ActiveTraits.insert(TraitProperty::implementation_vendor_llvm);
+  ActiveTraits.set(unsigned(TraitProperty::implementation_vendor_llvm));
 
   // The user condition true is accepted but not false.
-  ActiveTraits.insert(TraitProperty::user_condition_true);
+  ActiveTraits.set(unsigned(TraitProperty::user_condition_true));
 
   // This is for sure some device.
-  ActiveTraits.insert(TraitProperty::device_kind_any);
+  ActiveTraits.set(unsigned(TraitProperty::device_kind_any));
 
   LLVM_DEBUG({
     dbgs() << "[" << DEBUG_TYPE
            << "] New OpenMP context with the following properties:\n";
-    for (auto &Property : ActiveTraits)
+    for (const auto &SetBitsIt : ActiveTraits.set_bits()) {
+      TraitProperty Property = TraitProperty(SetBitsIt);
       dbgs() << "\t " << getOpenMPContextTraitPropertyFullName(Property)
              << "\n";
+    }
   });
 }
 
 /// Return true if \p C0 is a subset of \p C1. Note that both arrays are
 /// expected to be sorted.
 template <typename T> static bool isSubset(ArrayRef<T> C0, ArrayRef<T> C1) {
 #ifdef EXPENSIVE_CHECKS
   assert(std::is_sorted(C0.begin(), C0.end()) &&
static bool isStrictSubset(ArrayRef<T> C0, ArrayRef<T> C1) {
   return isSubset<T>(C0, C1);
 }
 
 static bool isStrictSubset(const VariantMatchInfo &VMI0,
                            const VariantMatchInfo &VMI1) {
   // If all required traits are a strict subset and the ordered vectors storing
   // the construct traits, we say it is a strict subset. Note that the latter
   // relation is not required to be strict.
-  return set_is_strict_subset(VMI0.RequiredTraits, VMI1.RequiredTraits) &&
-         isSubset<TraitProperty>(VMI0.ConstructTraits, VMI1.ConstructTraits);
+  if (VMI0.RequiredTraits.count() >= VMI1.RequiredTraits.count())
+    return false;
+  for (const auto &SetBitsIt : VMI0.RequiredTraits.set_bits())
+    if (!VMI1.RequiredTraits.test(SetBitsIt))
+      return false;
+  if (!isSubset<TraitProperty>(VMI0.ConstructTraits, VMI1.ConstructTraits))
+    return false;
+  return true;
 }
 
 static int isVariantApplicableInContextHelper(
     const VariantMatchInfo &VMI, const OMPContext &Ctx,
     SmallVectorImpl<unsigned> *ConstructMatches) {
 
-  for (TraitProperty Property : VMI.RequiredTraits) {
+  for (const auto &SetBitsIt : VMI.RequiredTraits.set_bits()) {
+    TraitProperty Property = TraitProperty(SetBitsIt);
 
-    bool IsActiveTrait = Ctx.ActiveTraits.count(Property);
+    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;
     }
   }
 

 }
 
 static APInt getVariantMatchScore(const VariantMatchInfo &VMI,
                                   const OMPContext &Ctx,
                                   SmallVectorImpl<unsigned> &ConstructMatches) {
   APInt Score(64, 1);
 
   unsigned NoConstructTraits = VMI.ConstructTraits.size();
-  for (TraitProperty Property : VMI.RequiredTraits) {
+  for (const auto &SetBitsIt : VMI.RequiredTraits.set_bits()) {
+    TraitProperty Property = TraitProperty(SetBitsIt);
     // If there is a user score attached, use it.
     if (VMI.ScoreMap.count(Property)) {
       const APInt &UserScore = VMI.ScoreMap.lookup(Property);
       assert(UserScore.uge(0) && "Expect non-negative user scores!");
       Score += UserScore.getZExtValue();
       continue;
     }