[MemProf] Extend CallSite information to include potential callees.

llvm/include/llvm/ProfileData/MemProf.h

@@ -342,6 +342,28 @@ using CallStackId = uint64_t;
 // A type representing the index into the call stack array.
 using LinearCallStackId = uint32_t;
 
+// Holds call site information with indexed frame contents.
+struct IndexedCallSiteInfo {
+  // The call stack ID for this call site
+  CallStackId CSId = 0;
+  // The GUIDs of the callees at this call site
+  SmallVector<GlobalValue::GUID, 1> CalleeGuids;
+
+  IndexedCallSiteInfo() = default;
+  IndexedCallSiteInfo(CallStackId CSId) : CSId(CSId) {}
+  IndexedCallSiteInfo(CallStackId CSId,
+                      SmallVector<GlobalValue::GUID, 1> CalleeGuids)
+      : CSId(CSId), CalleeGuids(std::move(CalleeGuids)) {}
+
+  bool operator==(const IndexedCallSiteInfo &Other) const {
+    return CSId == Other.CSId && CalleeGuids == Other.CalleeGuids;
+  }
+
+  bool operator!=(const IndexedCallSiteInfo &Other) const {
+    return !operator==(Other);
+  }
+};
+
 // Holds allocation information in a space efficient format where frames are
 // represented using unique identifiers.
 struct IndexedAllocationInfo {
@@ -410,7 +432,7 @@ struct IndexedMemProfRecord {
   // list of inline locations in bottom-up order i.e. from leaf to root. The
   // inline location list may include additional entries, users should pick
   // the last entry in the list with the same function GUID.
-  llvm::SmallVector<CallStackId> CallSiteIds;
+  llvm::SmallVector<IndexedCallSiteInfo> CallSites;
 
   void clear() { *this = IndexedMemProfRecord(); }
 
@@ -427,7 +449,7 @@ struct IndexedMemProfRecord {
     if (Other.AllocSites != AllocSites)
       return false;
 
-    if (Other.CallSiteIds != CallSiteIds)
+    if (Other.CallSites != CallSites)
       return false;
     return true;
   }
@@ -455,14 +477,37 @@ struct IndexedMemProfRecord {
   static GlobalValue::GUID getGUID(const StringRef FunctionName);
 };
 
+// Holds call site information with frame contents inline.
+struct CallSiteInfo {
+  // The frames in the call stack
+  std::vector<Frame> Frames;
+
+  // The GUIDs of the callees at this call site
+  SmallVector<GlobalValue::GUID, 1> CalleeGuids;
+
+  CallSiteInfo() = default;
+  CallSiteInfo(std::vector<Frame> Frames) : Frames(std::move(Frames)) {}
+  CallSiteInfo(std::vector<Frame> Frames,
+               SmallVector<GlobalValue::GUID, 1> CalleeGuids)
+      : Frames(std::move(Frames)), CalleeGuids(std::move(CalleeGuids)) {}
+
+  bool operator==(const CallSiteInfo &Other) const {
+    return Frames == Other.Frames && CalleeGuids == Other.CalleeGuids;
+  }
+
+  bool operator!=(const CallSiteInfo &Other) const {
+    return !operator==(Other);
+  }
+};
+
 // Holds the memprof profile information for a function. The internal
 // representation stores frame contents inline. This representation should
 // be used for small amount of temporary, in memory instances.
 struct MemProfRecord {
   // Same as IndexedMemProfRecord::AllocSites with frame contents inline.
   llvm::SmallVector<AllocationInfo> AllocSites;
   // Same as IndexedMemProfRecord::CallSites with frame contents inline.
-  llvm::SmallVector<std::vector<Frame>> CallSites;
+  llvm::SmallVector<CallSiteInfo> CallSites;
 
   MemProfRecord() = default;
 
@@ -476,8 +521,8 @@ struct MemProfRecord {
 
     if (!CallSites.empty()) {
       OS << "    CallSites:\n";
-      for (const std::vector<Frame> &Frames : CallSites) {
-        for (const Frame &F : Frames) {
+      for (const CallSiteInfo &CS : CallSites) {
+        for (const Frame &F : CS.Frames) {
           OS << "    -\n";
           F.printYAML(OS);
         }

llvm/include/llvm/ProfileData/MemProfYAML.h

@@ -155,6 +155,14 @@ template <> struct MappingTraits<memprof::AllocationInfo> {
 // In YAML, we use GUIDMemProfRecordPair instead of MemProfRecord so that we can
 // treat the GUID and the fields within MemProfRecord at the same level as if
 // the GUID were part of MemProfRecord.
+template <> struct MappingTraits<memprof::CallSiteInfo> {
+  static void mapping(IO &Io, memprof::CallSiteInfo &CS) {
+    Io.mapRequired("Frames", CS.Frames);
+    // Keep this optional to make it easier to write tests.
+    Io.mapOptional("CalleeGuids", CS.CalleeGuids);
+  }
+};
+
 template <> struct MappingTraits<memprof::GUIDMemProfRecordPair> {
   static void mapping(IO &Io, memprof::GUIDMemProfRecordPair &Pair) {
     Io.mapRequired("GUID", Pair.GUID);
@@ -174,6 +182,7 @@ template <> struct MappingTraits<memprof::AllMemProfData> {
 LLVM_YAML_IS_SEQUENCE_VECTOR(memprof::Frame)
 LLVM_YAML_IS_SEQUENCE_VECTOR(std::vector<memprof::Frame>)
 LLVM_YAML_IS_SEQUENCE_VECTOR(memprof::AllocationInfo)
+LLVM_YAML_IS_SEQUENCE_VECTOR(memprof::CallSiteInfo)
 LLVM_YAML_IS_SEQUENCE_VECTOR(memprof::GUIDMemProfRecordPair)
 
 #endif // LLVM_PROFILEDATA_MEMPROFYAML_H_

llvm/lib/ProfileData/MemProf.cpp

@@ -64,7 +64,7 @@ static size_t serializedSizeV2(const IndexedMemProfRecord &Record,
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
   // The CallStackId
-  Result += Record.CallSiteIds.size() * sizeof(CallStackId);
+  Result += Record.CallSites.size() * sizeof(CallStackId);
   return Result;
 }
 
@@ -78,7 +78,7 @@ static size_t serializedSizeV3(const IndexedMemProfRecord &Record,
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
   // The linear call stack ID.
-  Result += Record.CallSiteIds.size() * sizeof(LinearCallStackId);
+  Result += Record.CallSites.size() * sizeof(LinearCallStackId);
   return Result;
 }
 
@@ -106,9 +106,9 @@ static void serializeV2(const IndexedMemProfRecord &Record,
   }
 
   // Related contexts.
-  LE.write<uint64_t>(Record.CallSiteIds.size());
-  for (const auto &CSId : Record.CallSiteIds)
-    LE.write<CallStackId>(CSId);
+  LE.write<uint64_t>(Record.CallSites.size());
+  for (const auto &CS : Record.CallSites)
+    LE.write<CallStackId>(CS.CSId);
 }
 
 static void serializeV3(
@@ -127,10 +127,10 @@ static void serializeV3(
   }
 
   // Related contexts.
-  LE.write<uint64_t>(Record.CallSiteIds.size());
-  for (const auto &CSId : Record.CallSiteIds) {
-    assert(MemProfCallStackIndexes.contains(CSId));
-    LE.write<LinearCallStackId>(MemProfCallStackIndexes[CSId]);
+  LE.write<uint64_t>(Record.CallSites.size());
+  for (const auto &CS : Record.CallSites) {
+    assert(MemProfCallStackIndexes.contains(CS.CSId));
+    LE.write<LinearCallStackId>(MemProfCallStackIndexes[CS.CSId]);
   }
 }
 
@@ -170,11 +170,11 @@ static IndexedMemProfRecord deserializeV2(const MemProfSchema &Schema,
   // Read the callsite information.
   const uint64_t NumCtxs =
       endian::readNext<uint64_t, llvm::endianness::little>(Ptr);
-  Record.CallSiteIds.reserve(NumCtxs);
+  Record.CallSites.reserve(NumCtxs);
   for (uint64_t J = 0; J < NumCtxs; J++) {
     CallStackId CSId =
         endian::readNext<CallStackId, llvm::endianness::little>(Ptr);
-    Record.CallSiteIds.push_back(CSId);
+    Record.CallSites.emplace_back(CSId);
   }
 
   return Record;
@@ -202,15 +202,15 @@ static IndexedMemProfRecord deserializeV3(const MemProfSchema &Schema,
   // Read the callsite information.
   const uint64_t NumCtxs =
       endian::readNext<uint64_t, llvm::endianness::little>(Ptr);
-  Record.CallSiteIds.reserve(NumCtxs);
+  Record.CallSites.reserve(NumCtxs);
   for (uint64_t J = 0; J < NumCtxs; J++) {
     // We are storing LinearCallStackId in CallSiteIds, which is a vector of
     // CallStackId.  Assert that CallStackId is no smaller than
     // LinearCallStackId.
     static_assert(sizeof(LinearCallStackId) <= sizeof(CallStackId));
     LinearCallStackId CSId =
         endian::readNext<LinearCallStackId, llvm::endianness::little>(Ptr);
-    Record.CallSiteIds.push_back(CSId);
+    Record.CallSites.emplace_back(CSId);
   }
 
   return Record;
@@ -241,9 +241,11 @@ MemProfRecord IndexedMemProfRecord::toMemProfRecord(
     Record.AllocSites.push_back(std::move(AI));
   }
 
-  Record.CallSites.reserve(CallSiteIds.size());
-  for (CallStackId CSId : CallSiteIds)
-    Record.CallSites.push_back(Callback(CSId));
+  Record.CallSites.reserve(CallSites.size());
+  for (const IndexedCallSiteInfo &CS : CallSites) {
+    std::vector<Frame> Frames = Callback(CS.CSId);
+    Record.CallSites.emplace_back(std::move(Frames), CS.CalleeGuids);
+  }
 
   return Record;
 }

llvm/lib/ProfileData/MemProfReader.cpp

@@ -521,7 +521,7 @@ Error RawMemProfReader::mapRawProfileToRecords() {
     // we insert a new entry for callsite data if we need to.
     IndexedMemProfRecord &Record = MemProfData.Records[Id];
     for (LocationPtr Loc : Locs)
-      Record.CallSiteIds.push_back(MemProfData.addCallStack(*Loc));
+      Record.CallSites.emplace_back(MemProfData.addCallStack(*Loc));
   }
 
   return Error::success();
@@ -808,10 +808,10 @@ void YAMLMemProfReader::parse(StringRef YAMLData) {
       IndexedRecord.AllocSites.emplace_back(CSId, AI.Info);
     }
 
-    // Populate CallSiteIds.
+    // Populate CallSites with CalleeGuids.
     for (const auto &CallSite : Record.CallSites) {
-      CallStackId CSId = AddCallStack(CallSite);
-      IndexedRecord.CallSiteIds.push_back(CSId);
+      CallStackId CSId = AddCallStack(CallSite.Frames);
+      IndexedRecord.CallSites.emplace_back(CSId, CallSite.CalleeGuids);
     }
 
     MemProfData.Records.try_emplace(GUID, std::move(IndexedRecord));

llvm/lib/Transforms/Instrumentation/MemProfiler.cpp

@@ -954,7 +954,7 @@ undriftMemProfRecord(const DenseMap<uint64_t, LocToLocMap> &UndriftMaps,
     UndriftCallStack(AS.CallStack);
 
   for (auto &CS : MemProfRec.CallSites)
-    UndriftCallStack(CS);
+    UndriftCallStack(CS.Frames);
 }
 
 static void
@@ -1048,15 +1048,16 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
     // Need to record all frames from leaf up to and including this function,
     // as any of these may or may not have been inlined at this point.
     unsigned Idx = 0;
-    for (auto &StackFrame : CS) {
+    for (auto &StackFrame : CS.Frames) {
       uint64_t StackId = computeStackId(StackFrame);
-      LocHashToCallSites[StackId].insert(ArrayRef<Frame>(CS).drop_front(Idx++));
+      LocHashToCallSites[StackId].insert(
+          ArrayRef<Frame>(CS.Frames).drop_front(Idx++));
       ProfileHasColumns |= StackFrame.Column;
       // Once we find this function, we can stop recording.
       if (StackFrame.Function == FuncGUID)
         break;
     }
-    assert(Idx <= CS.size() && CS[Idx - 1].Function == FuncGUID);
+    assert(Idx <= CS.Frames.size() && CS.Frames[Idx - 1].Function == FuncGUID);
   }
 
   auto GetOffset = [](const DILocation *DIL) {

llvm/test/Transforms/PGOProfile/memprof-call-site-at-alloc-site.ll

@@ -20,7 +20,8 @@ HeapProfileRecords:
           TotalLifetime:   1000000
           TotalLifetimeAccessDensity: 1
     CallSites:
-      - - { Function: _Z3foov, LineOffset: 6, Column: 12, IsInlineFrame: false }
+      - Frames:
+        - { Function: _Z3foov, LineOffset: 6, Column: 12, IsInlineFrame: false }
 ...
 
 ;--- memprof-call-site-at-alloc-site.ll

llvm/test/Transforms/PGOProfile/memprof-dump-matched-call-sites.ll

@@ -42,16 +42,19 @@ HeapProfileRecords:
   - GUID:            main
     AllocSites:      []
     CallSites:
-      - - { Function: main, LineOffset: 1, Column: 3, IsInlineFrame: false }
+      - Frames:
+        - { Function: main, LineOffset: 1, Column: 3, IsInlineFrame: false }
   - GUID:            _ZL2f1v
     AllocSites:      []
     CallSites:
-      - - { Function: _ZL2f2v, LineOffset: 0, Column: 28, IsInlineFrame: true }
+      - Frames:
+        - { Function: _ZL2f2v, LineOffset: 0, Column: 28, IsInlineFrame: true }
         - { Function: _ZL2f1v, LineOffset: 0, Column: 54, IsInlineFrame: false }
   - GUID:            _ZL2f2v
     AllocSites:      []
     CallSites:
-      - - { Function: _ZL2f2v, LineOffset: 0, Column: 28, IsInlineFrame: true }
+      - Frames:
+        - { Function: _ZL2f2v, LineOffset: 0, Column: 28, IsInlineFrame: true }
         - { Function: _ZL2f1v, LineOffset: 0, Column: 54, IsInlineFrame: false }
   - GUID:            _Z2f3v
     AllocSites:

llvm/test/Transforms/PGOProfile/memprof-undrift.test

@@ -35,11 +35,13 @@ HeapProfileRecords:
   - GUID:            _Z3aaav
     AllocSites:      []
     CallSites:
-      - - { Function: _Z3aaav, LineOffset: 5, Column: 33, IsInlineFrame: false }
+      - Frames:
+        - { Function: _Z3aaav, LineOffset: 5, Column: 33, IsInlineFrame: false }
   - GUID:            _Z6middlev
     AllocSites:      []
     CallSites:
-      - - { Function: _Z6middlev, LineOffset: 5, Column: 33, IsInlineFrame: false }
+      - Frames:
+        - { Function: _Z6middlev, LineOffset: 5, Column: 33, IsInlineFrame: false }
   - GUID:            _Z3foov
     AllocSites:
       - Callstack:
@@ -77,7 +79,8 @@ HeapProfileRecords:
   - GUID:            _Z3bbbv
     AllocSites:      []
     CallSites:
-      - - { Function: _Z3bbbv, LineOffset: 5, Column: 33, IsInlineFrame: false }
+      - Frames:
+        - { Function: _Z3bbbv, LineOffset: 5, Column: 33, IsInlineFrame: false }
 ...
 ;--- memprof_undrift.ll
 define dso_local ptr @_Z3foov() !dbg !5 {

llvm/test/tools/llvm-profdata/memprof-yaml.test

@@ -1,9 +1,11 @@
-; REQUIRES: x86_64-linux
 ; RUN: split-file %s %t
 ; RUN: llvm-profdata merge %t/memprof-in.yaml -o %t/memprof-out.indexed
 ; RUN: llvm-profdata show --memory %t/memprof-out.indexed > %t/memprof-out.yaml
 ; RUN: cmp %t/memprof-in.yaml %t/memprof-out.yaml
 
+; This test is expected to fail until the profile format is updated to handle CalleeGuids.
+; XFAIL: *
+
 ; Verify that the YAML output is identical to the YAML input.
 ;--- memprof-in.yaml
 ---
@@ -27,8 +29,12 @@ HeapProfileRecords:
           TotalLifetime:   777
           TotalLifetimeAccessDensity: 888
     CallSites:
-      - - { Function: 0x5555555555555555, LineOffset: 55, Column: 50, IsInlineFrame: true }
+      - Frames:
+        - { Function: 0x5555555555555555, LineOffset: 55, Column: 50, IsInlineFrame: true }
         - { Function: 0x6666666666666666, LineOffset: 66, Column: 60, IsInlineFrame: false }
-      - - { Function: 0x7777777777777777, LineOffset: 77, Column: 70, IsInlineFrame: true }
+        CalleeGuids: [0x100, 0x200]
+      - Frames:
+        - { Function: 0x7777777777777777, LineOffset: 77, Column: 70, IsInlineFrame: true }
         - { Function: 0x8888888888888888, LineOffset: 88, Column: 80, IsInlineFrame: false }
+        CalleeGuids: [0x300]
 ...

llvm/unittests/ProfileData/InstrProfTest.cpp

@@ -397,7 +397,7 @@ makeRecordV2(std::initializer_list<::llvm::memprof::CallStackId> AllocFrames,
   for (const auto &CSId : AllocFrames)
     MR.AllocSites.emplace_back(CSId, Block, Schema);
   for (const auto &CSId : CallSiteFrames)
-    MR.CallSiteIds.push_back(CSId);
+    MR.CallSites.push_back(llvm::memprof::IndexedCallSiteInfo(CSId));
   return MR;
 }