[CSSPGO] Add attribute metadata for context profile

This changes adds attribute field for metadata of context profile. Currently we have an inline attribute that indicates whether the leaf frame corresponding to a context profile was inlined in previous build.

This will be used to help estimating inlining and be taken into account when trimming context. Changes for that in llvm-profgen will follow. It will also help tuning.

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

llvm/include/llvm/ProfileData/SampleProf.h

@@ -196,6 +196,7 @@ enum class SecProfSummaryFlags : uint32_t {
 enum class SecFuncMetadataFlags : uint32_t {
   SecFlagInvalid = 0,
   SecFlagIsProbeBased = (1 << 0),
+  SecFlagHasAttribute = (1 << 1)
 };
 
 // Verify section specific flag is used for the correct section.
@@ -385,6 +386,13 @@ enum ContextStateMask {
   MergedContext = 0x8     // Profile for context merged into base profile
 };
 
+// Attribute of context associated with FunctionSamples
+enum ContextAttributeMask {
+  ContextNone = 0x0,
+  ContextWasInlined = 0x1,      // Leaf of context was inlined in previous build
+  ContextShouldBeInlined = 0x2, // Leaf of context should be inlined
+};
+
 // Sample context for FunctionSamples. It consists of the calling context,
 // the function name and context state. Internally sample context is represented
 // using StringRef, which is also the input for constructing a `SampleContext`.
@@ -396,9 +404,9 @@ enum ContextStateMask {
 //    `_Z8funcLeafi`
 class SampleContext {
 public:
-  SampleContext() : State(UnknownContext) {}
-  SampleContext(StringRef ContextStr,
-                ContextStateMask CState = UnknownContext) {
+  SampleContext() : State(UnknownContext), Attributes(ContextNone) {}
+  SampleContext(StringRef ContextStr, ContextStateMask CState = UnknownContext)
+      : Attributes(ContextNone) {
     setContext(ContextStr, CState);
   }
 
@@ -443,6 +451,10 @@ public:
   }
 
   operator StringRef() const { return FullContext; }
+  bool hasAttribute(ContextAttributeMask A) { return Attributes & (uint32_t)A; }
+  void setAttribute(ContextAttributeMask A) { Attributes |= (uint32_t)A; }
+  uint32_t getAllAttributes() { return Attributes; }
+  void setAllAttributes(uint32_t A) { Attributes = A; }
   bool hasState(ContextStateMask S) { return State & (uint32_t)S; }
   void setState(ContextStateMask S) { State |= (uint32_t)S; }
   void clearState(ContextStateMask S) { State &= (uint32_t)~S; }
@@ -503,6 +515,8 @@ private:
   StringRef CallingContext;
   // State of the associated sample profile
   uint32_t State;
+  // Attribute of the associated sample profile
+  uint32_t Attributes;
 };
 
 class FunctionSamples;

llvm/include/llvm/ProfileData/SampleProfReader.h

@@ -28,6 +28,7 @@
 //       offsetA1[.discriminator]: number_of_samples [fn7:num fn8:num ... ]
 //       ...
 //      !CFGChecksum: num
+//      !Attribute: flags
 //
 // This is a nested tree in which the indentation represents the nesting level
 // of the inline stack. There are no blank lines in the file. And the spacing
@@ -127,6 +128,8 @@
 //
 // a. CFG Checksum (a.k.a. function hash):
 //   !CFGChecksum: 12345
+// b. CFG Checksum (see ContextAttributeMask):
+//   !Atribute: 1
 //
 //
 // Binary format
@@ -647,7 +650,7 @@ protected:
   std::error_code readSecHdrTableEntry(uint32_t Idx);
   std::error_code readSecHdrTable();
 
-  std::error_code readFuncMetadata();
+  std::error_code readFuncMetadata(bool ProfileHasAttribute);
   std::error_code readFuncOffsetTable();
   std::error_code readFuncProfiles();
   std::error_code readMD5NameTable();

llvm/lib/ProfileData/SampleProfReader.cpp

@@ -88,13 +88,22 @@ static bool isOffsetLegal(unsigned L) { return (L & 0xffff) == L; }
 /// Possible metadata:
 /// - CFG Checksum information:
 ///     !CFGChecksum: 12345
+/// - CFG Checksum information:
+///     !Attributes: 1
 /// Stores the FunctionHash (a.k.a. CFG Checksum) into \p FunctionHash.
-static bool parseMetadata(const StringRef &Input, uint64_t &FunctionHash) {
-  if (!Input.startswith("!CFGChecksum:"))
-    return false;
+static bool parseMetadata(const StringRef &Input, uint64_t &FunctionHash,
+                          uint32_t &Attributes) {
+  if (Input.startswith("!CFGChecksum:")) {
+    StringRef CFGInfo = Input.substr(strlen("!CFGChecksum:")).trim();
+    return !CFGInfo.getAsInteger(10, FunctionHash);
+  }
 
-  StringRef CFGInfo = Input.substr(strlen("!CFGChecksum:")).trim();
-  return !CFGInfo.getAsInteger(10, FunctionHash);
+  if (Input.startswith("!Attributes:")) {
+    StringRef Attrib = Input.substr(strlen("!Attributes:")).trim();
+    return !Attrib.getAsInteger(10, Attributes);
+  }
+
+  return false;
 }
 
 enum class LineType {
@@ -119,7 +128,7 @@ static bool ParseLine(const StringRef &Input, LineType &LineTy, uint32_t &Depth,
                       uint64_t &NumSamples, uint32_t &LineOffset,
                       uint32_t &Discriminator, StringRef &CalleeName,
                       DenseMap<StringRef, uint64_t> &TargetCountMap,
-                      uint64_t &FunctionHash) {
+                      uint64_t &FunctionHash, uint32_t &Attributes) {
   for (Depth = 0; Input[Depth] == ' '; Depth++)
     ;
   if (Depth == 0)
@@ -127,7 +136,7 @@ static bool ParseLine(const StringRef &Input, LineType &LineTy, uint32_t &Depth,
 
   if (Depth == 1 && Input[Depth] == '!') {
     LineTy = LineType::Metadata;
-    return parseMetadata(Input.substr(Depth), FunctionHash);
+    return parseMetadata(Input.substr(Depth), FunctionHash, Attributes);
   }
 
   size_t n1 = Input.find(':');
@@ -270,9 +279,11 @@ std::error_code SampleProfileReaderText::readImpl() {
       DenseMap<StringRef, uint64_t> TargetCountMap;
       uint32_t Depth, LineOffset, Discriminator;
       LineType LineTy;
-      uint64_t FunctionHash;
+      uint64_t FunctionHash = 0;
+      uint32_t Attributes = 0;
       if (!ParseLine(*LineIt, LineTy, Depth, NumSamples, LineOffset,
-                     Discriminator, FName, TargetCountMap, FunctionHash)) {
+                     Discriminator, FName, TargetCountMap, FunctionHash,
+                     Attributes)) {
         reportError(LineIt.line_number(),
                     "Expected 'NUM[.NUM]: NUM[ mangled_name:NUM]*', found " +
                         *LineIt);
@@ -312,8 +323,12 @@ std::error_code SampleProfileReaderText::readImpl() {
       }
       case LineType::Metadata: {
         FunctionSamples &FProfile = *InlineStack.back();
-        FProfile.setFunctionHash(FunctionHash);
-        ++ProbeProfileCount;
+        if (FunctionHash) {
+          FProfile.setFunctionHash(FunctionHash);
+          ++ProbeProfileCount;
+        }
+        if (Attributes)
+          FProfile.getContext().setAllAttributes(Attributes);
         SeenMetadata = true;
         break;
       }
@@ -601,13 +616,16 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     if (std::error_code EC = readFuncOffsetTable())
       return EC;
     break;
-  case SecFuncMetadata:
+  case SecFuncMetadata: {
     ProfileIsProbeBased =
         hasSecFlag(Entry, SecFuncMetadataFlags::SecFlagIsProbeBased);
     FunctionSamples::ProfileIsProbeBased = ProfileIsProbeBased;
-    if (std::error_code EC = readFuncMetadata())
+    bool HasAttribute =
+        hasSecFlag(Entry, SecFuncMetadataFlags::SecFlagHasAttribute);
+    if (std::error_code EC = readFuncMetadata(HasAttribute))
       return EC;
     break;
+  }
   case SecProfileSymbolList:
     if (std::error_code EC = readProfileSymbolList())
       return EC;
@@ -941,23 +959,31 @@ std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
   return SampleProfileReaderBinary::readNameTable();
 }
 
-std::error_code SampleProfileReaderExtBinaryBase::readFuncMetadata() {
-  if (!ProfileIsProbeBased)
-    return sampleprof_error::success;
+std::error_code
+SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
   while (Data < End) {
     auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
-    auto Checksum = readNumber<uint64_t>();
-    if (std::error_code EC = Checksum.getError())
-      return EC;
-
     SampleContext FContext(*FName);
-    // No need to load metadata for profiles that are not loaded in the current
-    // module.
-    if (Profiles.count(FContext))
-      Profiles[FContext].setFunctionHash(*Checksum);
+    bool ProfileInMap = Profiles.count(FContext);
+
+    if (ProfileIsProbeBased) {
+      auto Checksum = readNumber<uint64_t>();
+      if (std::error_code EC = Checksum.getError())
+        return EC;
+      if (ProfileInMap)
+        Profiles[FContext].setFunctionHash(*Checksum);
+    }
+
+    if (ProfileHasAttribute) {
+      auto Attributes = readNumber<uint32_t>();
+      if (std::error_code EC = Attributes.getError())
+        return EC;
+      if (ProfileInMap)
+        Profiles[FContext].getContext().setAllAttributes(*Attributes);
+    }
   }
 
   assert(Data == End && "More data is read than expected");

llvm/lib/ProfileData/SampleProfWriter.cpp

@@ -170,12 +170,15 @@ std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
     const StringMap<FunctionSamples> &Profiles) {
-  if (!FunctionSamples::ProfileIsProbeBased)
+  if (!FunctionSamples::ProfileIsProbeBased && !FunctionSamples::ProfileIsCS)
     return sampleprof_error::success;
   auto &OS = *OutputStream;
   for (const auto &Entry : Profiles) {
     writeNameIdx(Entry.first());
-    encodeULEB128(Entry.second.getFunctionHash(), OS);
+    if (FunctionSamples::ProfileIsProbeBased)
+      encodeULEB128(Entry.second.getFunctionHash(), OS);
+    if (FunctionSamples::ProfileIsCS)
+      encodeULEB128(Entry.second.getContext().getAllAttributes(), OS);
   }
   return sampleprof_error::success;
 }
@@ -239,6 +242,8 @@ std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsProbeBased);
   if (Type == SecProfSummary && FunctionSamples::ProfileIsCS)
     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFullContext);
+  if (Type == SecFuncMetadata && FunctionSamples::ProfileIsCS)
+    addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagHasAttribute);
 
   uint64_t SectionStart = markSectionStart(Type, LayoutIdx);
   switch (Type) {
@@ -417,6 +422,10 @@ std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
       OS.indent(Indent + 1);
       OS << "!CFGChecksum: " << S.getFunctionHash() << "\n";
     }
+    if (FunctionSamples::ProfileIsCS) {
+      OS.indent(Indent + 1);
+      OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
+    }
   }
 
   return sampleprof_error::success;

llvm/test/tools/llvm-profdata/Inputs/cs-sample.proftext

@@ -4,6 +4,7 @@
  3: 287884
  4: 287864 _Z3fibi:315608
  15: 23
+ !Attributes: 0
 [main:3.1 @ _Z5funcBi:1 @ _Z8funcLeafi]:500853:20
  0: 15
  1: 15
@@ -12,25 +13,32 @@
  10: 23324
  11: 23327 _Z3fibi:25228
  15: 11
+ !Attributes: 1
 [main]:154:0
  2: 12
  3: 18 _Z5funcAi:11
  3.1: 18 _Z5funcBi:19
+ !Attributes: 0
 [external:12 @ main]:154:12
  2: 12
  3: 10 _Z5funcAi:7
  3.1: 10 _Z5funcBi:11
+ !Attributes: 0
 [main:3.1 @ _Z5funcBi]:120:19
  0: 19
  1: 19 _Z8funcLeafi:20
  3: 12
+ !Attributes: 1
 [externalA:17 @ _Z5funcBi]:120:3
  0: 3
  1: 3
+ !Attributes: 0
 [external:10 @ _Z5funcBi]:120:10
  0: 10
  1: 10
+ !Attributes: 0
 [main:3 @ _Z5funcAi]:99:11
  0: 10
  1: 10 _Z8funcLeafi:11
  3: 24
+ !Attributes: 0

llvm/test/tools/llvm-profdata/Inputs/pseudo-probe-profile.proftext

@@ -6,3 +6,4 @@ foo:3200:13
  5: 7 _Z3foov:5 _Z3barv:2
  6: 6 _Z3barv:4 _Z3foov:2
  !CFGChecksum: 563022570642068
+ !Attributes: 0

llvm/test/tools/llvm-profgen/inline-cs-pseudoprobe.test

@@ -9,7 +9,7 @@
 ; CHECK-NEXT: 6: 15
 ; CHECK-NEXT: 8: 14 bar:14
 ; CHECK-NEXT: !CFGChecksum: 138950591924
-; CHECK-NEXT:[main:2 @ foo:8 @ bar]:28:14
+; CHECK:[main:2 @ foo:8 @ bar]:28:14
 ; CHECK-NEXT: 1: 14
 ; CHECK-NEXT: 2: 18446744073709551615
 ; CHECK-NEXT: 3: 18446744073709551615

llvm/test/tools/llvm-profgen/merge-cold-profile.test

@@ -14,6 +14,7 @@
 ; CHECK-NEXT: 7: 2 fb:2
 ; CHECK-NEXT: 8: 1 fa:1
 ; CHECK-NEXT: !CFGChecksum: 120515930909
+; CHECK-NEXT: !Attributes: 0
 ; CHECK-NEXT:[main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:13:4
 ; CHECK-NEXT: 1: 4
 ; CHECK-NEXT: 2: 3
@@ -29,6 +30,7 @@
 ; CHECK-KEEP-COLD-NEXT: 5: 4 fb:4
 ; CHECK-KEEP-COLD-NEXT: 6: 3 fa:3
 ; CHECK-KEEP-COLD-NEXT: !CFGChecksum: 72617220756
+; CHECK-KEEP-COLD-NEXT: !Attributes: 0
 ; CHECK-KEEP-COLD-NEXT:[fa]:14:4
 ; CHECK-KEEP-COLD-NEXT: 1: 4
 ; CHECK-KEEP-COLD-NEXT: 3: 4

llvm/test/tools/llvm-profgen/noinline-cs-pseudoprobe.test

@@ -8,7 +8,7 @@
 ; CHECK-NEXT: 6: 15
 ; CHECK-NEXT: 8: 15 bar:15
 ; CHECK-NEXT: !CFGChecksum: 138950591924
-; CHECK-NEXT:[main:2 @ foo:8 @ bar]:30:15
+; CHECK:[main:2 @ foo:8 @ bar]:30:15
 ; CHECK-NEXT: 1: 15
 ; CHECK-NEXT: 2: 18446744073709551615
 ; CHECK-NEXT: 3: 18446744073709551615

llvm/tools/llvm-profgen/PerfReader.cpp

@@ -92,7 +92,8 @@ void VirtualUnwinder::unwindBranchWithinFrame(UnwindState &State) {
 std::shared_ptr<StringBasedCtxKey> FrameStack::getContextKey() {
   std::shared_ptr<StringBasedCtxKey> KeyStr =
       std::make_shared<StringBasedCtxKey>();
-  KeyStr->Context = Binary->getExpandedContextStr(Stack);
+  KeyStr->Context =
+      Binary->getExpandedContextStr(Stack, KeyStr->WasLeafInlined);
   if (KeyStr->Context.empty())
     return nullptr;
   KeyStr->genHashCode();

llvm/tools/llvm-profgen/PerfReader.h

@@ -311,7 +311,8 @@ struct ContextKey {
 // String based context id
 struct StringBasedCtxKey : public ContextKey {
   std::string Context;
-  StringBasedCtxKey() : ContextKey(CK_StringBased){};
+  bool WasLeafInlined;
+  StringBasedCtxKey() : ContextKey(CK_StringBased), WasLeafInlined(false){};
   static bool classof(const ContextKey *K) {
     return K->getKind() == CK_StringBased;
   }

llvm/tools/llvm-profgen/ProfileGenerator.cpp

@@ -188,10 +188,13 @@ void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
 }
 
 FunctionSamples &
-CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr) {
+CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr,
+                                                 bool WasLeafInlined) {
   auto Ret = ProfileMap.try_emplace(ContextStr, FunctionSamples());
   if (Ret.second) {
     SampleContext FContext(Ret.first->first(), RawContext);
+    if (WasLeafInlined)
+      FContext.setAttribute(ContextWasInlined);
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(FContext);
   }
@@ -208,7 +211,7 @@ void CSProfileGenerator::generateProfile() {
       StringRef ContextId(CtxKey->Context);
       // Get or create function profile for the range
       FunctionSamples &FunctionProfile =
-          getFunctionProfileForContext(ContextId);
+          getFunctionProfileForContext(ContextId, CtxKey->WasLeafInlined);
 
       // Fill in function body samples
       populateFunctionBodySamples(FunctionProfile, CI.second.RangeCounter,
@@ -428,6 +431,7 @@ void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
     assert(Ret.second && "Must be a unique context");
     SampleContext FContext(Ret.first->first(), RawContext);
     FunctionSamples &FProfile = Ret.first->second;
+    FContext.setAllAttributes(FProfile.getContext().getAllAttributes());
     FProfile.setName(FContext.getNameWithContext(true));
     FProfile.setContext(FContext);
   }
@@ -587,7 +591,7 @@ void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
     SmallVectorImpl<std::string> &ContextStrStack,
-    const PseudoProbeFuncDesc *LeafFuncDesc) {
+    const PseudoProbeFuncDesc *LeafFuncDesc, bool WasLeafInlined) {
   assert(ContextStrStack.size() && "Profile context must have the leaf frame");
   // Compress the context string except for the leaf frame
   std::string LeafFrame = ContextStrStack.back();
@@ -608,7 +612,7 @@ FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
   OContextStr << StringRef(LeafFrame).split(":").first.str();
 
   FunctionSamples &FunctionProile =
-      getFunctionProfileForContext(OContextStr.str());
+      getFunctionProfileForContext(OContextStr.str(), WasLeafInlined);
   FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
   return FunctionProile;
 }
@@ -619,13 +623,11 @@ FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
   // Explicitly copy the context for appending the leaf context
   SmallVector<std::string, 16> ContextStrStackCopy(ContextStrStack.begin(),
                                                    ContextStrStack.end());
-  Binary->getInlineContextForProbe(LeafProbe, ContextStrStackCopy);
-  // Note that the context from probe doesn't include leaf frame,
-  // hence we need to retrieve and append the leaf frame.
+  Binary->getInlineContextForProbe(LeafProbe, ContextStrStackCopy, true);
   const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->GUID);
-  ContextStrStackCopy.emplace_back(FuncDesc->FuncName + ":" +
-                                   Twine(LeafProbe->Index).str());
-  return getFunctionProfileForLeafProbe(ContextStrStackCopy, FuncDesc);
+  bool WasLeafInlined = LeafProbe->InlineTree->hasInlineSite();
+  return getFunctionProfileForLeafProbe(ContextStrStackCopy, FuncDesc,
+                                        WasLeafInlined);
 }
 
 } // end namespace sampleprof

llvm/tools/llvm-profgen/ProfileGenerator.h

@@ -174,7 +174,8 @@ public:
 
 protected:
   // Lookup or create FunctionSamples for the context
-  FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
+  FunctionSamples &getFunctionProfileForContext(StringRef ContextId,
+                                                bool WasLeafInlined = false);
   // Merge cold context profile whose total sample is below threshold
   // into base profile.
   void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap);
@@ -229,7 +230,8 @@ private:
   // Helper function to get FunctionSamples for the leaf inlined context
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
-                                 const PseudoProbeFuncDesc *LeafFuncDesc);
+                                 const PseudoProbeFuncDesc *LeafFuncDesc,
+                                 bool WasLeafInlined);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,

llvm/tools/llvm-profgen/ProfiledBinary.cpp

@@ -131,8 +131,9 @@ bool ProfiledBinary::inlineContextEqual(uint64_t Address1,
                     Context2.begin(), Context2.begin() + Context2.size() - 1);
 }
 
-std::string ProfiledBinary::getExpandedContextStr(
-    const SmallVectorImpl<uint64_t> &Stack) const {
+std::string
+ProfiledBinary::getExpandedContextStr(const SmallVectorImpl<uint64_t> &Stack,
+                                      bool &WasLeafInlined) const {
   std::string ContextStr;
   SmallVector<std::string, 16> ContextVec;
   // Process from frame root to leaf
@@ -143,6 +144,9 @@ std::string ProfiledBinary::getExpandedContextStr(
     // processing
     if (ExpandedContext.empty())
       return std::string();
+    // Set WasLeafInlined to the size of inlined frame count for the last
+    // address which is leaf
+    WasLeafInlined = (ExpandedContext.size() > 1);
     for (const auto &Loc : ExpandedContext) {
       ContextVec.push_back(getCallSite(Loc));
     }

llvm/tools/llvm-profgen/ProfiledBinary.h

@@ -239,8 +239,8 @@ public:
   // Get the context string of the current stack with inline context filled in.
   // It will search the disassembling info stored in Offset2LocStackMap. This is
   // used as the key of function sample map
-  std::string
-  getExpandedContextStr(const SmallVectorImpl<uint64_t> &Stack) const;
+  std::string getExpandedContextStr(const SmallVectorImpl<uint64_t> &Stack,
+                                    bool &WasLeafInlined) const;
 
   const PseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);