[CSSPGO] Add attribute metadata for context profile

[wenlei]

Summary

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.

Note that we only populate inline attribute for pseudo-probe profile now.

Test Plan:

[wenlei]

update

[hoy]

Thanks for working on this. This gets us closer to a global inlining decision making for thinLTO.

llvm/lib/ProfileData/SampleProfReader.cpp

91–92

Please update comment here and also the format comment at the beginning of SampleProfReader.h .

llvm/tools/llvm-profgen/ProfileGenerator.cpp

556

Should caller profile be marked inlined? The callee was inlined into the caller, but the caller may not be inlined into caller's caller.

630

I'm wondering this should also be done for non-probe CS profile generation.

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

91–92

Updated.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

556

Good catch, you're right.

630

Yeah, this is a TODO as mentioned in the change description. The change to support that for non-probe profile will actually be bigger since we put everything into flat context string out of unwinder, so wanted to do them separately. Now I left a TODO comment as well, but let me take another look..

[wenlei]

Address Hongtao's comments.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

630

Turns out it can be done easily still, with the help of context key. Updated to support line CS profile too in this change.

[wenlei]

Populate attribute for line-based profile too.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

630

Thanks for adding the support. Might also need to change existing regression tests if they are failed.

Nit: change the description to remove that TODO. Otherwise LGTM.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

630

Tests are all passing, the commit message is updated, but not sure why it wasn't sync to Phabricator.

llvm/include/llvm/ProfileData/SampleProf.h

enum class SecProfSummaryFlags : uint32_t {
   /// SecFlagContext means this is context-sensitive profile for
   /// CSSPGO
   SecFlagFullContext = (1 << 1)
 };
 
 enum class SecFuncMetadataFlags : uint32_t {
   SecFlagInvalid = 0,
   SecFlagIsProbeBased = (1 << 0),
+  SecFlagHasAttribute = (1 << 1)
 };
 
 // Verify section specific flag is used for the correct section.
 template <class SecFlagType>
 static inline void verifySecFlag(SecType Type, SecFlagType Flag) {
   // No verification is needed for common flags.
   if (std::is_same<SecCommonFlags, SecFlagType>())
     return;

 enum ContextStateMask {
   UnknownContext = 0x0,   // Profile without context
   RawContext = 0x1,       // Full context profile from input profile
   SyntheticContext = 0x2, // Synthetic context created for context promotion
   InlinedContext = 0x4,   // Profile for context that is inlined into caller
   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`.
 // It can accept and represent both full context string as well as context-less
 // function name.
 // Example of full context string (note the wrapping `[]`):
 //    `[main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]`
 // Example of context-less function name (same as AutoFDO):
 //    `_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);
   }
 
   // Promote context by removing top frames (represented by `ContextStrToRemove`).
   // Note that with string representation of context, the promotion is effectively
   // a substr operation with `ContextStrToRemove` removed from left.
   void promoteOnPath(StringRef ContextStrToRemove) {
     assert(FullContext.startswith(ContextStrToRemove));
if (!EntrySplit.second.empty()) {
 
       // Get discriminator
       if (!LocSplit.second.empty())
         LocSplit.second.getAsInteger(10, LineLoc.Discriminator);
     }
   }
 
   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; }
   bool hasContext() const { return State != UnknownContext; }
   bool isBaseContext() const { return CallingContext.empty(); }
   StringRef getNameWithoutContext() const { return Name; }
   StringRef getCallingContext() const { return CallingContext; }
   StringRef getNameWithContext(bool WithBracket = false) const {
private:
   // Full context string including calling context and leaf function name
   StringRef FullContext;
   // Function name for the associated sample profile
   StringRef Name;
   // Calling context (leaf function excluded) for the associated sample profile
   StringRef CallingContext;
   // State of the associated sample profile
   uint32_t State;
+  // Attribute of the associated sample profile
+  uint32_t Attributes;
 };
 
 class FunctionSamples;
 class SampleProfileReaderItaniumRemapper;
 
 using BodySampleMap = std::map<LineLocation, SampleRecord>;
 // NOTE: Using a StringMap here makes parsed profiles consume around 17% more
 // memory, which is *very* significant for large profiles.

llvm/include/llvm/ProfileData/SampleProfReader.h

private:
 
   BumpPtrAllocator Allocator;
 
 protected:
   std::vector<SecHdrTableEntry> SecHdrTable;
   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();
   std::error_code readNameTableSec(bool IsMD5);
   std::error_code readProfileSymbolList();
 
   virtual std::error_code readHeader() override;
   virtual std::error_code verifySPMagic(uint64_t Magic) override = 0;

llvm/lib/ProfileData/SampleProfReader.cpp

 
 /// Returns true if line offset \p L is legal (only has 16 bits).
 static bool isOffsetLegal(unsigned L) { return (L & 0xffff) == L; }
 
 /// Parse \p Input that contains metadata.
 /// Possible metadata:
 /// - CFG Checksum information:
 ///     !CFGChecksum: 12345
 /// 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);
   }
 
+  if (Input.startswith("!Attributes:")) {
+    StringRef Attrib = Input.substr(strlen("!Attributes:")).trim();
+    return !Attrib.getAsInteger(10, Attributes);
+    ;
+  }
+
+  return false;
+}
+
 enum class LineType {
   CallSiteProfile,
   BodyProfile,
   Metadata,
 };
 
 /// Parse \p Input as line sample.
 ///
 /// \param Input input line.
 /// \param LineTy Type of this line.
 /// \param Depth the depth of the inline stack.
 /// \param NumSamples total samples of the line/inlined callsite.
 /// \param LineOffset line offset to the start of the function.
 /// \param Discriminator discriminator of the line.
 /// \param TargetCountMap map from indirect call target to count.
 /// \param FunctionHash the function's CFG hash, used by pseudo probe.
 ///
 /// returns true if parsing is successful.
 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)
     return false;
 
   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(':');
   StringRef Loc = Input.substr(Depth, n1 - Depth);
   size_t n2 = Loc.find('.');
   if (n2 == StringRef::npos) {
     if (Loc.getAsInteger(10, LineOffset) || !isOffsetLegal(LineOffset))
       return false;
if ((*LineIt)[0] != ' ') {
       InlineStack.clear();
       InlineStack.push_back(&FProfile);
     } else {
       uint64_t NumSamples;
       StringRef FName;
       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);
         return sampleprof_error::malformed;
       }
       if (SeenMetadata && LineTy != LineType::Metadata) {
         // Metadata must be put at the end of a function profile.
         reportError(LineIt.line_number(),
if ((*LineIt)[0] != ' ') {
                                   name_count.second));
         }
         MergeResult(Result, FProfile.addBodySamples(LineOffset, Discriminator,
                                                     NumSamples));
         break;
       }
       case LineType::Metadata: {
         FunctionSamples &FProfile = *InlineStack.back();
+        if (FunctionHash) {
           FProfile.setFunctionHash(FunctionHash);
           ++ProbeProfileCount;
+        }
+        if (Attributes)
+          FProfile.getContext().setAllAttributes(Attributes);
         SeenMetadata = true;
         break;
       }
       }
     }
   }
 
   assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
   case SecLBRProfile:
     if (std::error_code EC = readFuncProfiles())
       return EC;
     break;
   case SecFuncOffsetTable:
     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;
     break;
   default:
     if (std::error_code EC = readCustomSection(Entry))
       return EC;
     break;

 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
   if (IsMD5)
     return readMD5NameTable();
   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;
 
+    SampleContext FContext(*FName);
+    bool ProfileInMap = Profiles.count(FContext);
+
+    if (ProfileIsProbeBased) {
       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))
+      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");
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderCompactBinary::readNameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;

llvm/lib/ProfileData/SampleProfWriter.cpp

for (auto entry : FuncOffsetTable) {
     encodeULEB128(entry.second, OS);
   }
   FuncOffsetTable.clear();
   return sampleprof_error::success;
 }
 
 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());
+    if (FunctionSamples::ProfileIsProbeBased)
       encodeULEB128(Entry.second.getFunctionHash(), OS);
+    if (FunctionSamples::ProfileIsCS)
+      encodeULEB128(Entry.second.getContext().getAllAttributes(), OS);
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
   if (!UseMD5)
     return SampleProfileWriterBinary::writeNameTable();
 
std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
     const StringMap<FunctionSamples> &ProfileMap) {
   // The setting of SecFlagCompress should happen before markSectionStart.
   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
     setToCompressSection(SecProfileSymbolList);
   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsProbeBased)
     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) {
   case SecProfSummary:
     computeSummary(ProfileMap);
     if (auto EC = writeSummary())
       return EC;
     break;
for (const auto &I : SortedCallsiteSamples.get())
     }
   Indent -= 1;
 
   if (Indent == 0) {
     if (FunctionSamples::ProfileIsProbeBased) {
       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;
 }
 
 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
   const auto &ret = NameTable.find(FName);
   if (ret == NameTable.end())

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

 [main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]:1467299:11
  0: 6
  1: 6
  3: 287884
  4: 287864 _Z3fibi:315608
  15: 23
+ !Attributes: 0
 [main:3.1 @ _Z5funcBi:1 @ _Z8funcLeafi]:500853:20
  0: 15
  1: 15
  3: 74946
  4: 74941 _Z3fibi:82359
  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

 foo:3200:13
  1: 13
  2: 7
  3: 18446744073709551615
  4: 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

 ; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-pseudoprobe.perfscript --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:     [main:2 @ foo]:74:0
 ; CHECK-NEXT: 2: 15
 ; CHECK-NEXT: 3: 15
 ; CHECK-NEXT: 4: 14
 ; CHECK-NEXT: 5: 1
 ; 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
 ; CHECK-NEXT: 4: 14
 ; CHECK-NEXT: !CFGChecksum: 72617220756
 
 
 ; CHECK-UNWINDER:      Binary(inline-cs-pseudoprobe.perfbin)'s Range Counter:

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

 ; Used the data from recursion-compression.test, refer it for the unmerged output
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=8
 ; RUN: FileCheck %s --input-file %t
 
 ; Test --csprof-keep-cold
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=100 --csprof-keep-cold
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-KEEP-COLD
 
 ; CHECK:     [fa]:14:4
 ; CHECK-NEXT: 1: 4
 ; CHECK-NEXT: 3: 4
 ; CHECK-NEXT: 4: 2
 ; CHECK-NEXT: 5: 1
 ; 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
 ; CHECK-NEXT: 3: 1
 ; CHECK-NEXT: 5: 4 fb:4
 ; CHECK-NEXT: 6: 1 fa:1
 ; CHECK-NEXT: !CFGChecksum: 72617220756
 
 ; CHECK-KEEP-COLD:     [fb]:19:6
 ; CHECK-KEEP-COLD-NEXT: 1: 6
 ; CHECK-KEEP-COLD-NEXT: 2: 3
 ; CHECK-KEEP-COLD-NEXT: 3: 3
 ; 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
 ; CHECK-KEEP-COLD-NEXT: 4: 2
 ; CHECK-KEEP-COLD-NEXT: 5: 1
 ; CHECK-KEEP-COLD-NEXT: 7: 2 fb:2
 ; CHECK-KEEP-COLD-NEXT: 8: 1 fa:1
 ; CHECK-KEEP-COLD-NEXT: !CFGChecksum: 120515930909

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

 ; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-pseudoprobe.perfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:     [main:2 @ foo]:75:0
 ; CHECK-NEXT: 2: 15
 ; CHECK-NEXT: 3: 15
 ; CHECK-NEXT: 4: 15
 ; 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
 ; CHECK-NEXT: 4: 15
 ; CHECK-NEXT: !CFGChecksum: 72617220756
 
 
 ; CHECK-UNWINDER:      Binary(noinline-cs-pseudoprobe.perfbin)'s Range Counter:

llvm/tools/llvm-profgen/ProfileGenerator.h

while (I <= MaxDedupSize) {
       I++;
       Context.resize(End);
       MaxDedupSize = std::min(static_cast<uint32_t>(End / 2), MaxDedupSize);
     }
   }
 
 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);
   void write(std::unique_ptr<SampleProfileWriter> Writer,
              StringMap<FunctionSamples> &ProfileMap) override;
 
 private:
   // Helper function for updating body sample for a leaf location in
populateBodySamplesWithProbes(const RangeSample &RangeCounter,
                                 ProfiledBinary *Binary);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(
       const BranchSample &BranchCounter,
       SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary);
   // 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,
                                  const PseudoProbe *LeafProbe,
                                  ProfiledBinary *Binary);
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (Point.EndCount) {
       DisjointRanges[{BeginAddress, Address}] = Count;
       Count -= Point.EndCount;
       BeginAddress = Address + 1;
     }
   }
 }
 
 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);
   }
   return Ret.first->second;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
   // Add bracket for context key to support different profile binary format
   StringMap<FunctionSamples> CxtWithBracketPMap;
   for (const auto &Item : ProfileMap) {
     std::string ContextWithBracket = "[" + Item.first().str() + "]";
     auto Ret = CxtWithBracketPMap.try_emplace(ContextWithBracket, Item.second);
     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);
   }
   Writer->write(CxtWithBracketPMap);
 }
 
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
if (Probe->isEntry()) {
             FunctionProfile.getContext().getNameWithContext(true);
         StringRef CallerContextId;
         FrameLocation &&CallerLeafFrameLoc =
             getCallerContext(CalleeContextId, CallerContextId);
         uint64_t CallerIndex = CallerLeafFrameLoc.second.LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
         FunctionSamples &CallerProfile =
-            getFunctionProfileForContext(CallerContextId);
+            getFunctionProfileForContext(CallerContextId, true);
         CallerProfile.setFunctionHash(InlinerDesc->FuncHash);
         CallerProfile.addBodySamples(CallerIndex, 0, Count);
         CallerProfile.addTotalSamples(Count);
         CallerProfile.addCalledTargetSamples(
             CallerIndex, 0,
             FunctionProfile.getContext().getNameWithoutContext(), Count);
       }
     }
if (CalleeName.size() == 0)
       continue;
     FunctionProfile.addCalledTargetSamples(CallProbe->Index, 0, CalleeName,
                                            Count);
   }
 }
 
 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();
   ContextStrStack.pop_back();
   CSProfileGenerator::compressRecursionContext(ContextStrStack);
 
   std::ostringstream OContextStr;
   for (uint32_t I = 0; I < ContextStrStack.size(); I++) {
     if (OContextStr.str().size())
       OContextStr << " @ ";
     OContextStr << ContextStrStack[I];
   }
   // For leaf inlined context with the top frame, we should strip off the top
   // frame's probe id, like:
   // Inlined stack: [foo:1, bar:2], the ContextId will be "foo:1 @ bar"
   if (OContextStr.str().size())
     OContextStr << " @ ";
   OContextStr << StringRef(LeafFrame).split(":").first.str();
 
   FunctionSamples &FunctionProile =
-      getFunctionProfileForContext(OContextStr.str());
+      getFunctionProfileForContext(OContextStr.str(), WasLeafInlined);
   FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
   return FunctionProile;
 }
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
     SmallVectorImpl<std::string> &ContextStrStack, const PseudoProbe *LeafProbe,
     ProfiledBinary *Binary) {
   // 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
 } // end namespace llvm