[OpenMP][NFC] Minimize memory usage and copying of `OMPTraitInfo`s

[jdoerfert]

Summary

See rational here: https://reviews.llvm.org/D71830#1922656

[rnk]

lgtm, thanks!

I also noticed that Parser.h includes OpenMPClause.h just for this class. OpenMPClause.h is pretty big. It's for this family of related methods:

/// Parse a property kind into \p TIProperty for the selector set \p Set and
/// selector \p Selector.
void parseOMPTraitPropertyKind(OMPTraitInfo::OMPTraitProperty &TIProperty,
                               llvm::omp::TraitSet Set,
                               llvm::omp::TraitSelector Selector,
                               llvm::StringMap<SourceLocation> &Seen);

The use of an inner class here makes it impossible to forward declare OMPTraitProperty. Would you be opposed to moving it out of line?

The other popular includer of OpenMPClause.h is Attr.h, which is where I started my investigation. I think after your change, we can forward declare it. I'll look into it.

[jdoerfert]

In D76173#1922916, @rnk wrote:

lgtm, thanks!

I also noticed that Parser.h includes OpenMPClause.h just for this class. OpenMPClause.h is pretty big. It's for this family of related methods:

/// Parse a property kind into \p TIProperty for the selector set \p Set and
/// selector \p Selector.
void parseOMPTraitPropertyKind(OMPTraitInfo::OMPTraitProperty &TIProperty,
                               llvm::omp::TraitSet Set,
                               llvm::omp::TraitSelector Selector,
                               llvm::StringMap<SourceLocation> &Seen);

The use of an inner class here makes it impossible to forward declare OMPTraitProperty. Would you be opposed to moving it out of line?

Like D76184 ?

The other popular includer of OpenMPClause.h is Attr.h, which is where I started my investigation. I think after your change, we can forward declare it. I'll look into it.

Took care of that in D76184 as well. I will "soonish" redesign OpenMPClause.h completely but that will take more time.

clang/include/clang/AST/ASTContext.h

 class ObjCImplementationDecl;
 class ObjCInterfaceDecl;
 class ObjCIvarDecl;
 class ObjCMethodDecl;
 class ObjCPropertyDecl;
 class ObjCPropertyImplDecl;
 class ObjCProtocolDecl;
 class ObjCTypeParamDecl;
+struct OMPTraitInfo;
 struct ParsedTargetAttr;
 class Preprocessor;
 class Stmt;
 class StoredDeclsMap;
 class TargetAttr;
 class TargetInfo;
 class TemplateDecl;
 class TemplateParameterList;
struct SectionInfo {
     SectionInfo(DeclaratorDecl *Decl,
                 SourceLocation PragmaSectionLocation,
                 int SectionFlags)
         : Decl(Decl), PragmaSectionLocation(PragmaSectionLocation),
           SectionFlags(SectionFlags) {}
   };
 
   llvm::StringMap<SectionInfo> SectionInfos;
+
+  /// Return a new OMPTraitInfo object owned by the Sema.
+  OMPTraitInfo &getNewOMPTraitInfo();
+
+private:
+  /// All OMPTraitInfo objects live in this collection, one per `omp [begin]
+  /// declare variant` directive.
+  SmallVector<OMPTraitInfo *, 4> OMPTraitInfoVector;
 };
 
 /// Utility function for constructing a nullary selector.
 inline Selector GetNullarySelector(StringRef name, ASTContext &Ctx) {
   IdentifierInfo* II = &Ctx.Idents.get(name);
   return Ctx.Selectors.getSelector(0, &II);
 }
 

clang/include/clang/AST/OpenMPClause.h

 #include "clang/Basic/OpenMPKinds.def"
 };
 
 /// Helper data structure representing the traits in a match clause of an
 /// `declare variant` or `metadirective`. The outer level is an ordered
 /// collection of selector sets, each with an associated kind and an ordered
 /// collection of selectors. A selector has a kind, an optional score/condition,
 /// and an ordered collection of properties.
-struct OMPTraitInfo {
+class OMPTraitInfo {
+  /// Private constructor accesible only by ASTContext.
+  OMPTraitInfo() {}
+  friend class ASTContext;
+
+public:
   struct OMPTraitProperty {
     llvm::omp::TraitProperty Kind = llvm::omp::TraitProperty::invalid;
   };
   struct OMPTraitSelector {
     Expr *ScoreOrCondition = nullptr;
     llvm::omp::TraitSelector Kind = llvm::omp::TraitSelector::invalid;
-    llvm::SmallVector<OMPTraitProperty, 4> Properties;
+    llvm::SmallVector<OMPTraitProperty, 1> Properties;
   };
   struct OMPTraitSet {
     llvm::omp::TraitSet Kind = llvm::omp::TraitSet::invalid;
-    llvm::SmallVector<OMPTraitSelector, 4> Selectors;
+    llvm::SmallVector<OMPTraitSelector, 2> Selectors;
   };
 
   /// The outermost level of selector sets.
-  llvm::SmallVector<OMPTraitSet, 4> Sets;
+  llvm::SmallVector<OMPTraitSet, 2> Sets;
 
   bool anyScoreOrCondition(
       const llvm::function_ref<bool(Expr *&, bool /* IsScore */)> &Cond) {
     return llvm::any_of(Sets, [&Cond](OMPTraitInfo::OMPTraitSet &Set) {
       return llvm::any_of(
           Set.Selectors, [&Cond](OMPTraitInfo::OMPTraitSelector &Selector) {
             return Cond(Selector.ScoreOrCondition,
                         /* IsScore */ Selector.Kind !=
public:
   /// this method.
   void getAsVariantMatchInfo(ASTContext &ASTCtx,
                              llvm::omp::VariantMatchInfo &VMI) const;
 
   /// Print a human readable representation into \p OS.
   void print(llvm::raw_ostream &OS, const PrintingPolicy &Policy) const;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const OMPTraitInfo &TI);
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const OMPTraitInfo *TI);
 
 } // namespace clang
 
 #endif // LLVM_CLANG_AST_OPENMPCLAUSE_H

clang/include/clang/Basic/Attr.td

def OMPDeclareVariant : InheritableAttr {
   let HasCustomParsing = 1;
   let InheritEvenIfAlreadyPresent = 1;
   let Documentation = [OMPDeclareVariantDocs];
   let Args = [
     ExprArgument<"VariantFuncRef">,
     OMPTraitInfoArgument<"TraitInfos">,
   ];
   let AdditionalMembers = [{
-    OMPTraitInfo &getTraitInfo() { return traitInfos; }
+    OMPTraitInfo &getTraitInfo() { return *traitInfos; }
     void printPrettyPragma(raw_ostream & OS, const PrintingPolicy &Policy)
         const;
   }];
 }
 
 def InternalLinkage : InheritableAttr {
   let Spellings = [Clang<"internal_linkage">];
   let Subjects = SubjectList<[Var, Function, CXXRecord]>;

clang/include/clang/Serialization/ASTRecordReader.h

public:
   /// Read a CXXCtorInitializer array, advancing Idx.
   CXXCtorInitializer **readCXXCtorInitializers();
 
   CXXTemporary *readCXXTemporary() {
     return Reader->ReadCXXTemporary(*F, Record, Idx);
   }
 
   /// Read an OMPTraitInfo object, advancing Idx.
-  OMPTraitInfo readOMPTraitInfo();
+  OMPTraitInfo *readOMPTraitInfo();
 
   /// Read an OpenMP clause, advancing Idx.
   OMPClause *readOMPClause();
 
   /// Read a source location, advancing Idx.
   SourceLocation readSourceLocation() {
     return Reader->ReadSourceLocation(*F, Record, Idx);
   }

clang/include/clang/Serialization/ASTRecordWriter.h

public:
   void AddUnresolvedSet(const ASTUnresolvedSet &Set);
 
   /// Emit a CXXCtorInitializer array.
   void AddCXXCtorInitializers(ArrayRef<CXXCtorInitializer *> CtorInits);
 
   void AddCXXDefinitionData(const CXXRecordDecl *D);
 
   /// Write an OMPTraitInfo object.
-  void writeOMPTraitInfo(const OMPTraitInfo &TI);
+  void writeOMPTraitInfo(const OMPTraitInfo *TI);
 
   void writeOMPClause(OMPClause *C);
 
   /// Emit a string.
   void AddString(StringRef Str) {
     return Writer->AddString(Str, *Record);
   }
 

clang/lib/AST/ASTContext.cpp

for (llvm::DenseMap<const Decl*, AttrVec*>::iterator A = DeclAttrs.begin(),
        A != AEnd; ++A)
     A->second->~AttrVec();
 
   for (const auto &Value : ModuleInitializers)
     Value.second->~PerModuleInitializers();
 
   for (APValue *Value : APValueCleanups)
     Value->~APValue();
+
+  // Destroy the OMPTraitInfo objects that life here.
+  llvm::DeleteContainerPointers(OMPTraitInfoVector);
 }
 
 void ASTContext::setTraversalScope(const std::vector<Decl *> &TopLevelDecls) {
   TraversalScope = TopLevelDecls;
   getParentMapContext().clear();
 }
 
 void ASTContext::AddDeallocation(void (*Callback)(void *), void *Data) const {
Target->getCPUSpecificCPUDispatchFeatures(
         SD->getCPUName(GD.getMultiVersionIndex())->getName(), FeaturesTmp);
     std::vector<std::string> Features(FeaturesTmp.begin(), FeaturesTmp.end());
     Target->initFeatureMap(FeatureMap, getDiagnostics(), TargetCPU, Features);
   } else {
     Target->initFeatureMap(FeatureMap, getDiagnostics(), TargetCPU,
                            Target->getTargetOpts().Features);
   }
 }
+
+OMPTraitInfo &ASTContext::getNewOMPTraitInfo() {
+  OMPTraitInfoVector.push_back(new OMPTraitInfo());
+  return *OMPTraitInfoVector.back();
+}

clang/lib/AST/AttrImpl.cpp

 void OMPDeclareVariantAttr::printPrettyPragma(
     raw_ostream &OS, const PrintingPolicy &Policy) const {
   if (const Expr *E = getVariantFuncRef()) {
     OS << "(";
     E->printPretty(OS, nullptr, Policy);
     OS << ")";
   }
   OS << " match(";
-  traitInfos.print(OS, Policy);
+  traitInfos->print(OS, Policy);
   OS << ")";
 }
 
 #include "clang/AST/AttrImpl.inc"

clang/lib/AST/OpenMPClause.cpp

 
 llvm::raw_ostream &clang::operator<<(llvm::raw_ostream &OS,
                                      const OMPTraitInfo &TI) {
   LangOptions LO;
   PrintingPolicy Policy(LO);
   TI.print(OS, Policy);
   return OS;
 }
+llvm::raw_ostream &clang::operator<<(llvm::raw_ostream &OS,
+                                     const OMPTraitInfo *TI) {
+  return TI ? OS << *TI : OS;
+}

clang/lib/CodeGen/CGOpenMPRuntime.cpp

 static const FunctionDecl *getDeclareVariantFunction(CodeGenModule &CGM,
                                                      const FunctionDecl *FD) {
   if (!FD->hasAttrs() || !FD->hasAttr<OMPDeclareVariantAttr>())
     return FD;
 
   SmallVector<Expr *, 8> VariantExprs;
   SmallVector<VariantMatchInfo, 8> VMIs;
   for (const auto *A : FD->specific_attrs<OMPDeclareVariantAttr>()) {
-    const OMPTraitInfo &TI = A->getTraitInfos();
+    const OMPTraitInfo &TI = *A->getTraitInfos();
     VMIs.push_back(VariantMatchInfo());
     TI.getAsVariantMatchInfo(CGM.getContext(), VMIs.back());
     VariantExprs.push_back(A->getVariantFuncRef());
   }
 
   OMPContext Ctx(CGM.getLangOpts().OpenMPIsDevice, CGM.getTriple());
   // FIXME: Keep the context in the OMPIRBuilder so we can add constructs as we
   //        build them.

clang/lib/Parse/ParseOpenMP.cpp

if (T.expectAndConsume(diag::err_expected_lparen_after,
     while (!SkipUntil(tok::annot_pragma_openmp_end, StopBeforeMatch))
       ;
     // Skip the last annot_pragma_openmp_end.
     (void)ConsumeAnnotationToken();
     return;
   }
 
   // Parse inner context selectors.
-  OMPTraitInfo TI;
+  OMPTraitInfo &TI = Actions.getASTContext().getNewOMPTraitInfo();
   parseOMPContextSelectors(Loc, TI);
 
   // Parse ')'
   (void)T.consumeClose();
 
   Optional<std::pair<FunctionDecl *, Expr *>> DeclVarData =
       Actions.checkOpenMPDeclareVariantFunction(
           Ptr, AssociatedFunction.get(), TI,

clang/lib/Sema/SemaOpenMP.cpp

Sema::checkOpenMPDeclareVariantFunction(Sema::DeclGroupPtrTy DG,
   return std::make_pair(FD, cast<Expr>(DRE));
 }
 
 void Sema::ActOnOpenMPDeclareVariantDirective(FunctionDecl *FD,
                                               Expr *VariantRef,
                                               OMPTraitInfo &TI,
                                               SourceRange SR) {
   auto *NewAttr =
-      OMPDeclareVariantAttr::CreateImplicit(Context, VariantRef, TI, SR);
+      OMPDeclareVariantAttr::CreateImplicit(Context, VariantRef, &TI, SR);
   FD->addAttr(NewAttr);
 }
 
 void Sema::markOpenMPDeclareVariantFuncsReferenced(SourceLocation Loc,
                                                    FunctionDecl *Func,
                                                    bool MightBeOdrUse) {
   assert(LangOpts.OpenMP && "Expected OpenMP mode.");
 

clang/lib/Sema/SemaTemplateInstantiateDecl.cpp

if (Expr *E = Attr.getVariantFuncRef()) {
     // only place where it is used.
     EnterExpressionEvaluationContext Unevaluated(
         S, Sema::ExpressionEvaluationContext::ConstantEvaluated);
     VariantFuncRef = Subst(E);
   }
 
   // Copy the template version of the OMPTraitInfo and run substitute on all
   // score and condition expressiosn.
-  OMPTraitInfo TI = Attr.getTraitInfos();
+  OMPTraitInfo &TI = S.getASTContext().getNewOMPTraitInfo();
+  TI = *Attr.getTraitInfos();
 
   // Try to substitute template parameters in score and condition expressions.
   auto SubstScoreOrConditionExpr = [&S, Subst](Expr *&E, bool) {
     if (E) {
       EnterExpressionEvaluationContext Unevaluated(
           S, Sema::ExpressionEvaluationContext::ConstantEvaluated);
       ExprResult ER = Subst(E);
       if (ER.isUsable())

clang/lib/Serialization/ASTReader.cpp

 }
 
 void OMPClauseReader::VisitOMPOrderClause(OMPOrderClause *C) {
   C->setKind(Record.readEnum<OpenMPOrderClauseKind>());
   C->setLParenLoc(Record.readSourceLocation());
   C->setKindKwLoc(Record.readSourceLocation());
 }
 
-OMPTraitInfo ASTRecordReader::readOMPTraitInfo() {
-  OMPTraitInfo TI;
+OMPTraitInfo *ASTRecordReader::readOMPTraitInfo() {
+  OMPTraitInfo &TI = getContext().getNewOMPTraitInfo();
   TI.Sets.resize(readUInt32());
   for (auto &Set : TI.Sets) {
     Set.Kind = readEnum<llvm::omp::TraitSet>();
     Set.Selectors.resize(readUInt32());
     for (auto &Selector : Set.Selectors) {
       Selector.Kind = readEnum<llvm::omp::TraitSelector>();
       Selector.ScoreOrCondition = nullptr;
       if (readBool())
         Selector.ScoreOrCondition = readExprRef();
       Selector.Properties.resize(readUInt32());
       for (auto &Property : Selector.Properties)
         Property.Kind = readEnum<llvm::omp::TraitProperty>();
     }
   }
-  return TI;
+  return &TI;
 }

clang/lib/Serialization/ASTReaderDecl.cpp

public:
   IdentifierInfo *readIdentifier() {
     return Reader.readIdentifier();
   }
 
   VersionTuple readVersionTuple() {
     return Reader.readVersionTuple();
   }
 
-  OMPTraitInfo readOMPTraitInfo() { return Reader.readOMPTraitInfo(); }
+  OMPTraitInfo *readOMPTraitInfo() { return Reader.readOMPTraitInfo(); }
 
   template <typename T> T *GetLocalDeclAs(uint32_t LocalID) {
     return Reader.GetLocalDeclAs<T>(LocalID);
   }
 };
 }
 
 Attr *ASTRecordReader::readAttr() {

clang/lib/Serialization/ASTWriter.cpp

 }
 
 void OMPClauseWriter::VisitOMPOrderClause(OMPOrderClause *C) {
   Record.writeEnum(C->getKind());
   Record.AddSourceLocation(C->getLParenLoc());
   Record.AddSourceLocation(C->getKindKwLoc());
 }
 
-void ASTRecordWriter::writeOMPTraitInfo(const OMPTraitInfo &TI) {
-  writeUInt32(TI.Sets.size());
-  for (const auto &Set : TI.Sets) {
+void ASTRecordWriter::writeOMPTraitInfo(const OMPTraitInfo *TI) {
+  writeUInt32(TI->Sets.size());
+  for (const auto &Set : TI->Sets) {
     writeEnum(Set.Kind);
     writeUInt32(Set.Selectors.size());
     for (const auto &Selector : Set.Selectors) {
       writeEnum(Selector.Kind);
       writeBool(Selector.ScoreOrCondition);
       if (Selector.ScoreOrCondition)
         writeExprRef(Selector.ScoreOrCondition);
       writeUInt32(Selector.Properties.size());
       for (const auto &Property : Selector.Properties)
         writeEnum(Property.Kind);
     }
   }
 }

clang/utils/TableGen/ClangAttrEmitter.cpp

return StringSwitch<std::string>(type)
       .EndsWith("Decl *", "Record.GetLocalDeclAs<" +
                               std::string(type.data(), 0, type.size() - 1) +
                               ">(Record.readInt())")
       .Case("TypeSourceInfo *", "Record.readTypeSourceInfo()")
       .Case("Expr *", "Record.readExpr()")
       .Case("IdentifierInfo *", "Record.readIdentifier()")
       .Case("StringRef", "Record.readString()")
       .Case("ParamIdx", "ParamIdx::deserialize(Record.readInt())")
-      .Case("OMPTraitInfo", "Record.readOMPTraitInfo()")
+      .Case("OMPTraitInfo *", "Record.readOMPTraitInfo()")
       .Default("Record.readInt()");
 }
 
 // Get a type that is suitable for storing an object of the specified type.
 static StringRef getStorageType(StringRef type) {
   return StringSwitch<StringRef>(type)
     .Case("StringRef", "std::string")
     .Default(type);
 }
 
 // Assumes that the way to get the value is SA->getname()
 static std::string WritePCHRecord(StringRef type, StringRef name) {
   return "Record." +
          StringSwitch<std::string>(type)
              .EndsWith("Decl *", "AddDeclRef(" + std::string(name) + ");\n")
              .Case("TypeSourceInfo *",
                    "AddTypeSourceInfo(" + std::string(name) + ");\n")
              .Case("Expr *", "AddStmt(" + std::string(name) + ");\n")
              .Case("IdentifierInfo *",
                    "AddIdentifierRef(" + std::string(name) + ");\n")
              .Case("StringRef", "AddString(" + std::string(name) + ");\n")
              .Case("ParamIdx",
                    "push_back(" + std::string(name) + ".serialize());\n")
-             .Case("OMPTraitInfo",
+             .Case("OMPTraitInfo *",
                    "writeOMPTraitInfo(" + std::string(name) + ");\n")
              .Default("push_back(" + std::string(name) + ");\n");
 }
 
 // Normalize attribute name by removing leading and trailing
 // underscores. For example, __foo, foo__, __foo__ would
 // become foo.
 static StringRef NormalizeAttrName(StringRef AttrName) {
void writeDump(raw_ostream &OS) const override {
            << getUpperName() << "\";\n";
       } else if (type == "int" || type == "unsigned") {
         OS << "    OS << \" \" << SA->get" << getUpperName() << "();\n";
       } else if (type == "ParamIdx") {
         if (isOptional())
           OS << "    if (SA->get" << getUpperName() << "().isValid())\n  ";
         OS << "    OS << \" \" << SA->get" << getUpperName()
            << "().getSourceIndex();\n";
-      } else if (type == "OMPTraitInfo") {
+      } else if (type == "OMPTraitInfo *") {
         OS << "    OS << \" \" << SA->get" << getUpperName() << "();\n";
       } else {
         llvm_unreachable("Unknown SimpleArgument type!");
       }
     }
   };
 
   class DefaultSimpleArgument : public SimpleArgument {
else if (ArgName == "VariadicParamOrParamIdxArgument")
     Ptr = std::make_unique<VariadicParamOrParamIdxArgument>(Arg, Attr);
   else if (ArgName == "ParamIdxArgument")
     Ptr = std::make_unique<SimpleArgument>(Arg, Attr, "ParamIdx");
   else if (ArgName == "VariadicIdentifierArgument")
     Ptr = std::make_unique<VariadicIdentifierArgument>(Arg, Attr);
   else if (ArgName == "VersionArgument")
     Ptr = std::make_unique<VersionArgument>(Arg, Attr);
   else if (ArgName == "OMPTraitInfoArgument")
-    Ptr = std::make_unique<SimpleArgument>(Arg, Attr, "OMPTraitInfo");
+    Ptr = std::make_unique<SimpleArgument>(Arg, Attr, "OMPTraitInfo *");
 
   if (!Ptr) {
     // Search in reverse order so that the most-derived type is handled first.
     ArrayRef<std::pair<Record*, SMRange>> Bases = Search->getSuperClasses();
     for (const auto &Base : llvm::reverse(Bases)) {
       if ((Ptr = createArgument(Arg, Attr, Base.first)))
         break;
     }