[memprof] Store callsite metadata with memprof records.

[snehasish]

Summary

To ease profile annotation, each of the callsites in a function can be
annotated with profile data - "IR metadata format for MemProf" [1]. This
patch extends the on-disk serialized record format to store the debug
information for allocation callsites incl inline frames. This change is
incompatible with the existing format i.e. indexed profiles must be
regenerated, raw profiles are unaffected.

[1] https://groups.google.com/g/llvm-dev/c/aWHsdMxKAfE/m/WtEmRqyhAgAJ

[snehasish]

Add a comment and remove some commented code.

[tejohnson]

llvm/include/llvm/ProfileData/MemProf.h

237

It would be good to document the data structure choice - i.e. why this is a vector of vectors

286

Should this be comparing CallSites as well?

Also, it might be nice to refactor some of the code below into an operator== in the AllocationInfo struct.

llvm/include/llvm/ProfileData/RawMemProfReader.h

103

missing word(s) after "that"?

llvm/lib/ProfileData/RawMemProfReader.cpp

301

Document what the uint64_t values are.

324

Why can't the leaf function be an inline frame? I.e. what if we have:

foo() { malloc(); }
bar() { foo(); }

and we inline foo into bar.

Won't the leaf be foo which will be an inline frame?

332

If we want to avoid a hash table lookup to SymbolizedFrame for each of these addresses in the later loop, we could store a pointer to the the DenseMap iterator value_type for SymbolizedFrame[Address].

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

50

Can this test now also look for the CallSites?

llvm/unittests/ProfileData/InstrProfTest.cpp

228

Should this file also be testing CallSites?

283

Why is this test commented out now?

llvm/unittests/ProfileData/MemProfTest.cpp

189

Should this callstack now have 4 entries and can all be checked?

Additionally, should the record for foo now have CallSites since it shows up in the call stack for the allocation in CSM[0x2]?

[snehasish]

Address comments.

[snehasish]

Thanks for the detailed review! PTAL.

llvm/lib/ProfileData/RawMemProfReader.cpp

301

The comment here should really have been down below before the loop. Updated the comment since we now use a pointer to avoid another lookup in the map.

324

My usage of the term "leaf" is the cause of confusion, here I identify "last non-inlined" function in the callstack as leaf. Since we will annotate the profile post-inlining, this is the function where the records will be placed. In the example above, bar will have the memprof record entry. Let me know how you want me to reword this to make it clearer.

Also a note about ordering of entries in the callstack which represented by a list of MemProfRecord::Frames. For example

qux() { malloc(); }
inline foo() { qux(); }
bar() { foo(); }

is represented as:

bar (isInline: false); foo (isInline: true); qux (isInline:false)

as opposed to the expectation:

foo (isInline: true); bar (isInline: false); qux (isInline:false)

Based on your comment, this may be surprising since foo is the "leaf" function. I chose this representation to have deterministic access to the function id where the memprof record is to be applied, i.e. it is always at index 0. However, now that we have a map of function ids to memprof records in the reader (previously we just had a list of memprof records) we can change the order of to be more intuitive (i.e. on L339 we will append the frames in reverse). Wdyt?

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

50

The input for this test case only has main which calls malloc directly so there is no callsite data to add in this profile. I'll follow up with another llvm-profdata test which has some inlined allocation calls (also see discussion in leaf function comment).

llvm/unittests/ProfileData/InstrProfTest.cpp

283

Meant to remove it altogether since the failure path is no longer exercised by this test. The error is now guarded by an assert in mapRawProfileToRecords L331.

llvm/unittests/ProfileData/MemProfTest.cpp

189

Sure, added the check below.

[tejohnson]

llvm/lib/ProfileData/RawMemProfReader.cpp

324

My usage of the term "leaf" is the cause of confusion, here I identify "last non-inlined" function in the callstack as leaf. Since we will annotate the profile post-inlining, this is the function where the records will be placed. In the example above, bar will have the memprof record entry. Let me know how you want me to reword this to make it clearer.

An issue with this is that foo may have been inlined into bar before or after the point of matching. I.e. the inline frames in the profile are those inlined anywhere in the compiler: either in early inlining (which is done before typical PGO matching), or in the usual inline optimization pass (or in the rarer but supported case of an instrumented build with ThinLTO, in the LTO backend's invocation of the inliner).

So in my example, when we match we may or may not have inlined foo into bar at that point. Which raises the question of which function(s) should get the memprof data in the profile. I would argue that at the very least foo() should, even if it is marked as an inline frame (because when we match it may not yet have been inlined, and in fact based on profile info we may or may not make the same inlining decisions later in the compile either). Possibly it should also be on the non-inlined caller (and any other intervening inline frames), to make matching easier, in the case where the inline was done by the early inliner before matching.

Actually, just realized this same question affects where we attach the callsites profiles below given a set of inline frames.

Also a note about ordering of entries in the callstack which represented by a list of MemProfRecord::Frames. For example

qux() { malloc(); }
inline foo() { qux(); }
bar() { foo(); }

is represented as:

bar (isInline: false); foo (isInline: true); qux (isInline:false)

This makes sense to me since bar calls foo calls qux. Also, just a note that the isInline here assumes we only inlined foo into bar (which is suggested but not technically guaranteed by the inline keyword, and also we could have additional inlining by the optimizer).

as opposed to the expectation:

foo (isInline: true); bar (isInline: false); qux (isInline:false)

I'm not sure why this one would be the expectation? It is neither bottom up nor top down.

Based on your comment, this may be surprising since foo is the "leaf" function. I chose this representation to have deterministic access to the function id where the memprof record is to be applied, i.e. it is always at index 0. However, now that we have a map of function ids to memprof records in the reader (previously we just had a list of memprof records) we can change the order of to be more intuitive (i.e. on L339 we will append the frames in reverse). Wdyt?

[snehasish]

Discussed with @tejohnson offline how we can make the matching more robust.

[snehasish]

Update allocsite and callsite profile contents.

[snehasish]

@tejohnson I've updated the logic to track and annotate additional functions with the allocation site profiles. Please take a look at the example in MemProfTest.cpp L176 to see if it matches your expectations. Thanks!

[tejohnson]

llvm/include/llvm/ProfileData/MemProf.h

287

Specify ordering of inline frames (i.e. leaf to root).

llvm/lib/ProfileData/RawMemProfReader.cpp

299

maybe replace "to each callsite location" with "to the set of callsite locations"

333

I think we decided we need to do the same thing for the callsites that you are doing below for the memprof data on allocations, and put the callsites metadata on the inline functions as well, since the inlining at the point of matching may be less than the inlining in the final instrumented binary. Although in that case in the inline functions we only need the frames in Frames up to and including that function.

e.g.

void qux() { ... }
void foo() { qux(); }
void bar() { foo(); }

If instrumented binary inlined foo -> bar (but not qux->foo), Frames for the address in bar would include:
foo:0:13 (IsInlineFrame = true)
bar:0:13 (IsInlineFrame = false)

and we would want:

foo:

CallSites:
   [ foo:0:13 ]

bar:

CallSites:
   [ foo:0:13, bar:0:13 ]

345

It would be clearer I think to have this comment just above the start of this loop, to explain what the loop is doing.

llvm/unittests/ProfileData/MemProfTest.cpp

173

Is this printing meant to be here?

177

This is similar to the case I described in my comment in mapRawProfileToRecords(), where I think we should get CallSite metadata in foo() and in xyz() below. Of course, with the "inline" keyword it is likely to be early inlined before we do matching. But in the general case the symbolized stack frame may contain inline frames from a later inlining pass, which would not have been inlined at the point of matching yet. And we can't distinguish these cases when doing the profile symbolization / generation.

[snehasish]

Address comments.

[snehasish]

Fix lit test and some cleanup.

[snehasish]

Updated the patch to track callsites for all functions in the frame. PTAL, thanks!

[tejohnson]

lgtm

[snehasish]

With this patch on an internal binary, the time taken to merge indexed and raw profile increases to 85s (up from 48s in D120430). The zip profile size increases to ~520M (up from 247M in D118653). I'll look into de-duplicating the inline storage of callstacks and their frames in followup patches. Thanks for the review!

[hctim]

llvm/include/llvm/ProfileData/MemProf.h

264

Hi, looks like this patch causes the MemorySanitizer build bot to fail:

See instructions here on how to reproduce:

********************
Testing:  0.. 10.. 20.. 30.. 40.. 50.. 60.. 70.. 80.. 90
FAIL: LLVM-Unit :: ProfileData/./ProfileDataTests/MemProf.RecordSerializationRoundTrip (78166 of 82953)
******************** TEST 'LLVM-Unit :: ProfileData/./ProfileDataTests/MemProf.RecordSerializationRoundTrip' FAILED ********************
Script:
--
/b/sanitizer-x86_64-linux-fast/build/llvm_build_msan/unittests/ProfileData/./ProfileDataTests --gtest_filter=MemProf.RecordSerializationRoundTrip
--
Note: Google Test filter = MemProf.RecordSerializationRoundTrip
[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from MemProf
[ RUN      ] MemProf.RecordSerializationRoundTrip
==47398==WARNING: MemorySanitizer: use-of-uninitialized-value
    #0 0x4ae582 in llvm::memprof::MemProfRecord::AllocationInfo::operator==(llvm::memprof::MemProfRecord::AllocationInfo const&) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ProfileData/MemProf.h:236:11
    #1 0x4adb6a in operator!= /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ProfileData/MemProf.h:250:15
    #2 0x4adb6a in llvm::memprof::MemProfRecord::operator==(llvm::memprof::MemProfRecord const&) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ProfileData/MemProf.h:319:25
    #3 0x589fc6 in (anonymous namespace)::EqualsRecordMatcherP<llvm::memprof::MemProfRecord>::gmock_Impl<llvm::memprof::MemProfRecord const&>::MatchAndExplain(llvm::memprof::MemProfRecord const&, testing::MatchResultListener*) const /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/unittests/ProfileData/MemProfTest.cpp:111:11
    #4 0x585c50 in MatchAndExplain /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/include/gtest/gtest-matchers.h:264:19
    #5 0x585c50 in Matches /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/include/gtest/gtest-matchers.h:270:12
    #6 0x585c50 in operator()<llvm::memprof::MemProfRecord> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googlemock/include/gmock/gmock-matchers.h:1327:17
    #7 0x585c50 in (anonymous namespace)::MemProf_RecordSerializationRoundTrip_Test::TestBody() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/unittests/ProfileData/MemProfTest.cpp:300:3
    #8 0xa56c69 in HandleExceptionsInMethodIfSupported<testing::Test, void> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc
    #9 0xa56c69 in testing::Test::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2508:5
    #10 0xa5ba53 in testing::TestInfo::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2684:11
    #11 0xa5d14e in testing::TestSuite::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2816:28
    #12 0xa8be08 in testing::internal::UnitTestImpl::RunAllTests() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:5338:44
    #13 0xa8a443 in HandleExceptionsInMethodIfSupported<testing::internal::UnitTestImpl, bool> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc
    #14 0xa8a443 in testing::UnitTest::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:4925:10
    #15 0xa3b70e in RUN_ALL_TESTS /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/include/gtest/gtest.h:2472:46
    #16 0xa3b70e in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/UnitTestMain/TestMain.cpp:55:10
    #17 0x7fbb240a409a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a) (BuildId: 18b9a9a8c523e5cfe5b5d946d605d09242f09798)
    #18 0x379e39 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_msan/unittests/ProfileData/ProfileDataTests+0x379e39)

llvm/include/llvm/ProfileData/InstrProfWriter.h

 // based PGO and coverage.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_PROFILEDATA_INSTRPROFWRITER_H
 #define LLVM_PROFILEDATA_INSTRPROFWRITER_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/MemProf.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include <cstdint>
 #include <memory>
 
public:
   using ProfilingData = SmallDenseMap<uint64_t, InstrProfRecord>;
 
 private:
   bool Sparse;
   StringMap<ProfilingData> FunctionData;
 
   // A map to hold memprof data per function. The lower 64 bits obtained from
   // the md5 hash of the function name is used to index into the map.
-  memprof::FunctionMemProfMap MemProfData;
+  llvm::MapVector<Function::GUID, memprof::MemProfRecord> MemProfData;
 
   // An enum describing the attributes of the profile.
   InstrProfKind ProfileKind = InstrProfKind::Unknown;
   // Use raw pointer here for the incomplete type object.
   InstrProfRecordWriterTrait *InfoObj;
 
 public:
   InstrProfWriter(bool Sparse = false);
   ~InstrProfWriter();
 
   StringMap<ProfilingData> &getProfileData() { return FunctionData; }
 
   /// Add function counts for the given function. If there are already counts
   /// for this function and the hash and number of counts match, each counter is
   /// summed. Optionally scale counts by \p Weight.
   void addRecord(NamedInstrProfRecord &&I, uint64_t Weight,
                  function_ref<void(Error)> Warn);
   void addRecord(NamedInstrProfRecord &&I, function_ref<void(Error)> Warn) {
     addRecord(std::move(I), 1, Warn);
   }
 
-  void addRecord(const ::llvm::memprof::MemProfRecord &MR,
+  void addRecord(const Function::GUID Id, const memprof::MemProfRecord &Record,
                  function_ref<void(Error)> Warn);
 
   /// Merge existing function counts from the given writer.
   void mergeRecordsFromWriter(InstrProfWriter &&IPW,
                               function_ref<void(Error)> Warn);
 
   /// Write the profile to \c OS
   Error write(raw_fd_ostream &OS);

llvm/include/llvm/ProfileData/MemProf.h

#undef MIBEntryDef
       default:
         llvm_unreachable("Unknown meta type id, invalid input?");
       }
     }
   }
 
   // Print out the contents of the MemInfoBlock in YAML format.
   void printYAML(raw_ostream &OS) const {
-    OS << "    MemInfoBlock:\n";
+    OS << "      MemInfoBlock:\n";
 #define MIBEntryDef(NameTag, Name, Type)                                       \
-  OS << "      " << #Name << ": " << Name << "\n";
+  OS << "        " << #Name << ": " << Name << "\n";
 #include "llvm/ProfileData/MIBEntryDef.inc"
 #undef MIBEntryDef
   }
 
   // Define getters for each type which can be called by analyses.
 #define MIBEntryDef(NameTag, Name, Type)                                       \
   Type get##Name() const { return Name; }
 #include "llvm/ProfileData/MIBEntryDef.inc"
#undef MIBEntryDef
   }
 
 private:
 #define MIBEntryDef(NameTag, Name, Type) Type Name = Type();
 #include "llvm/ProfileData/MIBEntryDef.inc"
 #undef MIBEntryDef
 };
 
+// Holds the memprof profile information for a function.
 struct MemProfRecord {
   // Describes a call frame for a dynamic allocation context. The contents of
   // the frame are populated by symbolizing the stack depot call frame from the
   // compiler runtime.
   struct Frame {
     // A uuid (uint64_t) identifying the function. It is obtained by
     // llvm::md5(FunctionName) which returns the lower 64 bits.
     GlobalValue::GUID Function;
static Frame deserialize(const unsigned char *Ptr) {
                    /*IsInlineFrame=*/I);
     }
 
     // Returns the size of the frame information.
     static constexpr size_t serializedSize() {
       return sizeof(Frame::Function) + sizeof(Frame::LineOffset) +
              sizeof(Frame::Column) + sizeof(Frame::IsInlineFrame);
     }
+
+    // Print the frame information in YAML format.
+    void printYAML(raw_ostream &OS) const {
+      OS << "      -\n"
+         << "        Function: " << Function << "\n"
+         << "        LineOffset: " << LineOffset << "\n"
+         << "        Column: " << Column << "\n"
+         << "        Inline: " << IsInlineFrame << "\n";
+    }
   };
 
+  struct AllocationInfo {
     // The dynamic calling context for the allocation.
     llvm::SmallVector<Frame> CallStack;
     // The statistics obtained from the runtime for the allocation.
     PortableMemInfoBlock Info;
 
-  void clear() {
-    CallStack.clear();
-    Info.clear();
+    AllocationInfo() = default;
+    AllocationInfo(ArrayRef<Frame> CS, const MemInfoBlock &MB)
+        : CallStack(CS.begin(), CS.end()), Info(MB) {}
+
+    void printYAML(raw_ostream &OS) const {
+      OS << "    -\n";
+      OS << "      Callstack:\n";
+      // TODO: Print out the frame on one line with to make it easier for deep
+      // callstacks once we have a test to check valid YAML is generated.
+      for (const auto &Frame : CallStack)
+        Frame.printYAML(OS);
+      Info.printYAML(OS);
     }
 
     size_t serializedSize() const {
       return sizeof(uint64_t) + // The number of frames to serialize.
              Frame::serializedSize() *
                  CallStack.size() + // The contents of the frames.
              PortableMemInfoBlock::serializedSize(); // The size of the payload.
     }
+  };
+
+  llvm::SmallVector<AllocationInfo> AllocSites;
+  llvm::SmallVector<llvm::SmallVector<Frame>> CallSites;
+
+  void clear() {
+    AllocSites.clear();
+    CallSites.clear();
+  }
+
+  void merge(const MemProfRecord &Other) {
+    // TODO: Filter out duplicates which may occur if multiple memprof profiles
+    // are merged together using llvm-profdata.
+    AllocSites.append(Other.AllocSites);
+    CallSites.append(Other.CallSites);
+  }
+
+  size_t serializedSize() const {
+    size_t Result = sizeof(Function::GUID);
+    for (const AllocationInfo &N : AllocSites)
+      Result += N.serializedSize();
+
+    // The number of callsites we have information for.
+    Result += sizeof(uint64_t);
+    for (const auto &Frames : CallSites) {
+      // The number of frames to serialize.
+      Result += sizeof(uint64_t);
+      for (const Frame &F : Frames)
+        Result += F.serializedSize();
+    }
+    return Result;
+  }
 
   // Prints out the contents of the memprof record in YAML.
   void print(llvm::raw_ostream &OS) const {
-    OS << "    Callstack:\n";
-    // TODO: Print out the frame on one line with to make it easier for deep
-    // callstacks once we have a test to check valid YAML is generated.
-    for (const auto &Frame : CallStack) {
-      OS << "    -\n"
-         << "      Function: " << Frame.Function << "\n"
-         << "      LineOffset: " << Frame.LineOffset << "\n"
-         << "      Column: " << Frame.Column << "\n"
-         << "      Inline: " << Frame.IsInlineFrame << "\n";
+    if (!AllocSites.empty()) {
+      OS << "    AllocSites:\n";
+      for (const AllocationInfo &N : AllocSites)
+        N.printYAML(OS);
     }
 
-    Info.printYAML(OS);
+    if (!CallSites.empty()) {
+      OS << "    CallSites:\n";
+      for (const auto &Frames : CallSites) {
+        for (const auto &F : Frames) {
+          OS << "    -\n";
+          F.printYAML(OS);
+        }
+      }
+    }
   }
 
   bool operator==(const MemProfRecord &Other) const {
-    if (Other.Info != Info)
+    if (Other.AllocSites.size() != AllocSites.size())
       return false;
 
-    if (Other.CallStack.size() != CallStack.size())
+    for (size_t I = 0; I < AllocSites.size(); I++) {
+      if (Other.AllocSites[I].Info != AllocSites[I].Info)
         return false;
 
-    for (size_t I = 0; I < Other.CallStack.size(); I++) {
-      if (Other.CallStack[I] != CallStack[I])
+      if (Other.AllocSites[I].CallStack.size() !=
+          AllocSites[I].CallStack.size())
         return false;
+
+      for (size_t J = 0; J < Other.AllocSites[I].CallStack.size(); J++) {
+        if (Other.AllocSites[I].CallStack[J] != AllocSites[I].CallStack[J])
+          return false;
+      }
     }
     return true;
   }
-};
 
   // Serializes the memprof records in \p Records to the ostream \p OS based on
   // the schema provided in \p Schema.
-void serializeRecords(const ArrayRef<MemProfRecord> Records,
-                      const MemProfSchema &Schema, raw_ostream &OS);
+  void serialize(const MemProfSchema &Schema, raw_ostream &OS);
 
-// Deserializes memprof records from the Buffer
-SmallVector<MemProfRecord, 4> deserializeRecords(const MemProfSchema &Schema,
-                                                 const unsigned char *Buffer);
+  // Deserializes memprof records from the Buffer.
+  static MemProfRecord deserialize(const MemProfSchema &Schema,
+};
 
 // Reads a memprof schema from a buffer. All entries in the buffer are
 // interpreted as uint64_t. The first entry in the buffer denotes the number of
 // ids in the schema. Subsequent entries are integers which map to memprof::Meta
 // enum class entries. After successfully reading the schema, the pointer is one
 // byte past the schema contents.
 Expected<MemProfSchema> readMemProfSchema(const unsigned char *&Buffer);
 
-using FunctionMemProfMap =
-    DenseMap<uint64_t, SmallVector<memprof::MemProfRecord, 4>>;
-
 /// Trait for lookups into the on-disk hash table for memprof format in the
 /// indexed profile.
 class MemProfRecordLookupTrait {
 public:
-  using data_type = ArrayRef<MemProfRecord>;
+  using data_type = const MemProfRecord &;
   using internal_key_type = uint64_t;
   using external_key_type = uint64_t;
   using hash_value_type = uint64_t;
   using offset_type = uint64_t;
 
   MemProfRecordLookupTrait() = delete;
   MemProfRecordLookupTrait(const MemProfSchema &S) : Schema(S) {}
 
public:
 
   uint64_t ReadKey(const unsigned char *D, offset_type /*Unused*/) {
     using namespace support;
     return endian::readNext<external_key_type, little, unaligned>(D);
   }
 
   data_type ReadData(uint64_t K, const unsigned char *D,
                      offset_type /*Unused*/) {
-    Records = deserializeRecords(Schema, D);
-    return Records;
+    Record = MemProfRecord::deserialize(Schema, D);
+    return Record;
   }
 
 private:
   // Holds the memprof schema used to deserialize records.
   MemProfSchema Schema;
   // Holds the records from one function deserialized from the indexed format.
-  llvm::SmallVector<MemProfRecord, 4> Records;
+  MemProfRecord Record;
 };
 
 class MemProfRecordWriterTrait {
 public:
   using key_type = uint64_t;
   using key_type_ref = uint64_t;
 
-  using data_type = ArrayRef<MemProfRecord>;
-  using data_type_ref = ArrayRef<MemProfRecord>;
+  using data_type = MemProfRecord;
+  using data_type_ref = MemProfRecord &;
 
   using hash_value_type = uint64_t;
   using offset_type = uint64_t;
 
   // Pointer to the memprof schema to use for the generator. Unlike the reader
   // we must use a default constructor with no params for the writer trait so we
   // have a public member which must be initialized by the user.
   MemProfSchema *Schema = nullptr;
 
   MemProfRecordWriterTrait() = default;
 
   static hash_value_type ComputeHash(key_type_ref K) { return K; }
 
   static std::pair<offset_type, offset_type>
   EmitKeyDataLength(raw_ostream &Out, key_type_ref K, data_type_ref V) {
     using namespace support;
 
     endian::Writer LE(Out, little);
-
     offset_type N = sizeof(K);
     LE.write<offset_type>(N);
-
-    offset_type M = 0;
-
-    M += sizeof(uint64_t);
-    for (const auto &Record : V) {
-      M += Record.serializedSize();
-    }
-
+    offset_type M = V.serializedSize();
     LE.write<offset_type>(M);
     return std::make_pair(N, M);
   }
 
   void EmitKey(raw_ostream &Out, key_type_ref K, offset_type /*Unused*/) {
     using namespace support;
     endian::Writer LE(Out, little);
     LE.write<uint64_t>(K);
   }
 
   void EmitData(raw_ostream &Out, key_type_ref /*Unused*/, data_type_ref V,
                 offset_type /*Unused*/) {
     assert(Schema != nullptr && "MemProf schema is not initialized!");
-    serializeRecords(V, *Schema, Out);
+    V.serialize(*Schema, Out);
   }
 };
 
 } // namespace memprof
 } // namespace llvm
 
 #endif // LLVM_PROFILEDATA_MEMPROF_H_

llvm/include/llvm/ProfileData/RawMemProfReader.h

 #ifndef LLVM_PROFILEDATA_RAWMEMPROFREADER_H_
 #define LLVM_PROFILEDATA_RAWMEMPROFREADER_H_
 //===- MemProfReader.h - Instrumented memory profiling reader ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains support for reading MemProf profiling data.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableModule.h"
 #include "llvm/DebugInfo/Symbolize/Symbolize.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/MemProf.h"
 #include "llvm/ProfileData/MemProfData.inc"
public:
   static bool hasFormat(const StringRef Path);
 
   // Create a RawMemProfReader after sanity checking the contents of the file at
   // \p Path. The binary from which the profile has been collected is specified
   // via a path in \p ProfiledBinary.
   static Expected<std::unique_ptr<RawMemProfReader>>
   create(const Twine &Path, const StringRef ProfiledBinary);
 
-  Error readNextRecord(MemProfRecord &Record);
-
-  using Iterator = InstrProfIterator<MemProfRecord, RawMemProfReader>;
+  using GuidMemProfRecordPair = std::pair<Function::GUID, MemProfRecord>;
+  using Iterator = InstrProfIterator<GuidMemProfRecordPair, RawMemProfReader>;
   Iterator end() { return Iterator(); }
   Iterator begin() {
-    Iter = ProfileData.begin();
+    Iter = FunctionProfileData.begin();
     return Iterator(this);
   }
 
+  Error readNextRecord(GuidMemProfRecordPair &GuidRecord);
+
   // The RawMemProfReader only holds memory profile information.
   InstrProfKind getProfileKind() const { return InstrProfKind::MemProf; }
 
   // Constructor for unittests only.
   RawMemProfReader(std::unique_ptr<llvm::symbolize::SymbolizableModule> Sym,
                    llvm::SmallVectorImpl<SegmentEntry> &Seg,
                    llvm::MapVector<uint64_t, MemInfoBlock> &Prof,
                    CallStackMap &SM)
       : Symbolizer(std::move(Sym)), SegmentInfo(Seg.begin(), Seg.end()),
-        ProfileData(Prof), StackMap(SM) {
+        CallstackProfileData(Prof), StackMap(SM) {
     // We don't call initialize here since there is no raw profile to read. The
     // test should pass in the raw profile as structured data.
 
     // If there is an error here then the mock symbolizer has not been
     // initialized properly.
     if (Error E = symbolizeAndFilterStackFrames())
       report_fatal_error(std::move(E));
+    if (Error E = mapRawProfileToRecords())
+      report_fatal_error(std::move(E));
   }
 
 private:
   RawMemProfReader(std::unique_ptr<MemoryBuffer> DataBuffer,
                    object::OwningBinary<object::Binary> &&Bin)
       : DataBuffer(std::move(DataBuffer)), Binary(std::move(Bin)) {}
   Error initialize();
   Error readRawProfile();
   // Symbolize and cache all the virtual addresses we encounter in the
   // callstacks from the raw profile. Also prune callstack frames which we can't
   // symbolize or those that belong to the runtime. For profile entries where
   // the entire callstack is pruned, we drop the entry from the profile.
   Error symbolizeAndFilterStackFrames();
+  // Construct memprof records for each function that and store it in the
+  // `FunctionProfileData` map. A function may have allocation profile data or
+  // callsite data or both.
+  Error mapRawProfileToRecords();
 
   object::SectionedAddress getModuleOffset(uint64_t VirtualAddress);
-  Error fillRecord(const uint64_t Id, const MemInfoBlock &MIB,
-                   MemProfRecord &Record);
   // Prints aggregate counts for each raw profile parsed from the DataBuffer in
   // YAML format.
   void printSummaries(raw_ostream &OS) const;
 
   std::unique_ptr<MemoryBuffer> DataBuffer;
   object::OwningBinary<object::Binary> Binary;
   std::unique_ptr<llvm::symbolize::SymbolizableModule> Symbolizer;
 
   // The contents of the raw profile.
   llvm::SmallVector<SegmentEntry, 16> SegmentInfo;
   // A map from callstack id (same as key in CallStackMap below) to the heap
   // information recorded for that allocation context.
-  llvm::MapVector<uint64_t, MemInfoBlock> ProfileData;
+  llvm::MapVector<uint64_t, MemInfoBlock> CallstackProfileData;
   CallStackMap StackMap;
 
   // Cached symbolization from PC to Frame.
   llvm::DenseMap<uint64_t, llvm::SmallVector<MemProfRecord::Frame>>
       SymbolizedFrame;
 
-  // Iterator to read from the ProfileData MapVector.
-  llvm::MapVector<uint64_t, MemInfoBlock>::iterator Iter = ProfileData.end();
+  llvm::MapVector<Function::GUID, MemProfRecord> FunctionProfileData;
+  llvm::MapVector<Function::GUID, MemProfRecord>::iterator Iter;
 };
 
 } // namespace memprof
 } // namespace llvm
 
 #endif // LLVM_PROFILEDATA_RAWMEMPROFREADER_H_

llvm/lib/ProfileData/InstrProfWriter.cpp

void InstrProfWriter::addRecord(StringRef Name, uint64_t Hash,
   } else {
     // We're updating a function we've seen before.
     Dest.merge(I, Weight, MapWarn);
   }
 
   Dest.sortValueData();
 }
 
-void InstrProfWriter::addRecord(const memprof::MemProfRecord &MR,
+void InstrProfWriter::addRecord(const Function::GUID Id,
+                                const memprof::MemProfRecord &Record,
                                 function_ref<void(Error)> Warn) {
-  // Use 0 as a sentinel value since its highly unlikely that the lower 64-bits
-  // of a 128 bit md5 hash will be all zeros.
-  // TODO: Move this Key frame detection to the contructor to avoid having to
-  // scan all the callstacks again when adding a new record.
-  uint64_t Key = 0;
-  for (auto Iter = MR.CallStack.rbegin(), End = MR.CallStack.rend();
-       Iter != End; Iter++) {
-    if (!Iter->IsInlineFrame) {
-      Key = Iter->Function;
-      break;
-    }
-  }
+  auto Result = MemProfData.insert({Id, Record});
+  if (!Result.second) {
+    memprof::MemProfRecord &Existing = Result.first->second;
+    Existing.merge(Record);
 
-  if (Key == 0) {
-    Warn(make_error<InstrProfError>(
-        instrprof_error::invalid_prof,
-        "could not determine leaf function for memprof record."));
-  }
-
-  MemProfData[Key].push_back(MR);
-}
-
 void InstrProfWriter::mergeRecordsFromWriter(InstrProfWriter &&IPW,
                                              function_ref<void(Error)> Warn) {
   for (auto &I : IPW.FunctionData)
     for (auto &Func : I.getValue())
       addRecord(I.getKey(), Func.first, std::move(Func.second), 1, Warn);
 
-  for (auto &I : IPW.MemProfData)
-    for (const auto &MR : I.second)
-      addRecord(MR, Warn);
+  for (auto &I : IPW.MemProfData) {
+    addRecord(I.first, I.second, Warn);
+  }
 }
 
 bool InstrProfWriter::shouldEncodeData(const ProfilingData &PD) {
   if (!Sparse)
     return true;
   for (const auto &Func : PD) {
     const InstrProfRecord &IPR = Func.second;
     if (llvm::any_of(IPR.Counts, [](uint64_t Count) { return Count > 0; }))
if (static_cast<bool>(ProfileKind & InstrProfKind::MemProf)) {
     for (const auto Id : Schema) {
       OS.write(static_cast<uint64_t>(Id));
     }
 
     auto MemProfWriter = std::make_unique<memprof::MemProfRecordWriterTrait>();
     MemProfWriter->Schema = &Schema;
     OnDiskChainedHashTableGenerator<memprof::MemProfRecordWriterTrait>
         MemProfGenerator;
-    for (const auto &I : MemProfData) {
-      // Insert the key (func hash) and value (vector of memprof records).
+    for (auto &I : MemProfData) {
+      // Insert the key (func hash) and value (memprof record).
       MemProfGenerator.insert(I.first, I.second);
     }
 
     uint64_t TableOffset = MemProfGenerator.Emit(OS.OS, *MemProfWriter);
     PatchItem PatchItems[] = {
         {MemProfSectionStart, &TableOffset, 1},
     };
     OS.patch(PatchItems, 1);

llvm/lib/ProfileData/MemProf.cpp

 #include "llvm/ProfileData/MemProf.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
 
 namespace llvm {
 namespace memprof {
 
-void serializeRecords(const ArrayRef<MemProfRecord> Records,
-                      const MemProfSchema &Schema, raw_ostream &OS) {
+void MemProfRecord::serialize(const MemProfSchema &Schema, raw_ostream &OS) {
   using namespace support;
 
   endian::Writer LE(OS, little);
 
-  LE.write<uint64_t>(Records.size());
-  for (const MemProfRecord &MR : Records) {
-    LE.write<uint64_t>(MR.CallStack.size());
-    for (const MemProfRecord::Frame &F : MR.CallStack) {
+  LE.write<uint64_t>(AllocSites.size());
+  for (const AllocationInfo &N : AllocSites) {
+    LE.write<uint64_t>(N.CallStack.size());
+    for (const Frame &F : N.CallStack)
       F.serialize(OS);
-    }
-    MR.Info.serialize(Schema, OS);
+    N.Info.serialize(Schema, OS);
+
+  // Related contexts.
+  LE.write<uint64_t>(CallSites.size());
+  for (const auto &Frames : CallSites) {
+    LE.write<uint64_t>(Frames.size());
+    for (const Frame &F : Frames)
+      F.serialize(OS);
   }
 }
 
-SmallVector<MemProfRecord, 4> deserializeRecords(const MemProfSchema &Schema,
-                                                 const unsigned char *Ptr) {
+MemProfRecord MemProfRecord::deserialize(const MemProfSchema &Schema,
   using namespace support;
 
-  SmallVector<MemProfRecord, 4> Records;
-  const uint64_t NumRecords =
-      endian::readNext<uint64_t, little, unaligned>(Ptr);
-  for (uint64_t I = 0; I < NumRecords; I++) {
-    MemProfRecord MR;
+  MemProfRecord Record;
+
+  // Read the meminfo nodes.
+  const uint64_t NumNodes = endian::readNext<uint64_t, little, unaligned>(Ptr);
+  for (uint64_t I = 0; I < NumNodes; I++) {
+    MemProfRecord::AllocationInfo Node;
     const uint64_t NumFrames =
         endian::readNext<uint64_t, little, unaligned>(Ptr);
     for (uint64_t J = 0; J < NumFrames; J++) {
       const auto F = MemProfRecord::Frame::deserialize(Ptr);
       Ptr += MemProfRecord::Frame::serializedSize();
-      MR.CallStack.push_back(F);
+      Node.CallStack.push_back(F);
     }
-    MR.Info.deserialize(Schema, Ptr);
+    Node.Info.deserialize(Schema, Ptr);
     Ptr += PortableMemInfoBlock::serializedSize();
-    Records.push_back(MR);
+    Record.AllocSites.push_back(Node);
+  }
+
+  // Read the callsite information.
+  const uint64_t NumCtxs = endian::readNext<uint64_t, little, unaligned>(Ptr);
+  for (uint64_t J = 0; J < NumCtxs; J++) {
+    const uint64_t NumFrames =
+        endian::readNext<uint64_t, little, unaligned>(Ptr);
+    llvm::SmallVector<Frame> Frames;
+    for (uint64_t K = 0; K < NumFrames; K++) {
+      const auto F = MemProfRecord::Frame::deserialize(Ptr);
+      Ptr += MemProfRecord::Frame::serializedSize();
+      Frames.push_back(F);
+    }
+    Record.CallSites.push_back(Frames);
   }
-  return Records;
+
+  return Record;
 }
 
 Expected<MemProfSchema> readMemProfSchema(const unsigned char *&Buffer) {
   using namespace support;
 
   const unsigned char *Ptr = Buffer;
   const uint64_t NumSchemaIds =
       endian::readNext<uint64_t, little, unaligned>(Ptr);

llvm/lib/ProfileData/RawMemProfReader.cpp

 void RawMemProfReader::printYAML(raw_ostream &OS) {
   OS << "MemprofProfile:\n";
   // TODO: Update printSummaries to print out the data after the profile has
   // been symbolized and pruned. We can parse some raw profile characteristics
   // from the data buffer for additional information.
   printSummaries(OS);
   // Print out the merged contents of the profiles.
   OS << "  Records:\n";
-  for (const auto &Record : *this) {
+  for (const auto &Entry : *this) {
     OS << "  -\n";
-    Record.print(OS);
+    OS << "    FunctionGUID: " << Entry.first << "\n";
+    Entry.second.print(OS);
   }
 }
 
 void RawMemProfReader::printSummaries(raw_ostream &OS) const {
   const char *Next = DataBuffer->getBufferStart();
   while (Next < DataBuffer->getBufferEnd()) {
     auto Summary = computeSummary(Next);
     OS << "  -\n";
auto SOFOr = symbolize::SymbolizableObjectFile::create(
       Object, std::move(Context), /*UntagAddresses=*/false);
   if (!SOFOr)
     return report(SOFOr.takeError(), FileName);
   Symbolizer = std::move(SOFOr.get());
 
   if (Error E = readRawProfile())
     return E;
 
-  return symbolizeAndFilterStackFrames();
+  if (Error E = symbolizeAndFilterStackFrames())
+    return E;
+
+  return mapRawProfileToRecords();
+}
+
+Error RawMemProfReader::mapRawProfileToRecords() {
+  // Convert the raw profile callstack data into memprof records. While doing so
+  // keep track of related contexts so that we can fill these in later.
+  llvm::DenseMap<Function::GUID, llvm::SetVector<uint64_t>>
+      PerFunctionCallSites;
+
+  for (const auto &Entry : CallstackProfileData) {
+    const uint64_t StackId = Entry.first;
+
+    auto It = StackMap.find(StackId);
+    if (It == StackMap.end())
+      return make_error<InstrProfError>(
+          instrprof_error::malformed,
+          "memprof callstack record does not contain id: " + Twine(StackId));
+
+    // Construct the symbolized callstack.
+    llvm::SmallVector<MemProfRecord::Frame> Callstack;
+    Callstack.reserve(It->getSecond().size());
+
+    Function::GUID LeafFunction = 0;
+    for (const uint64_t Address : It->getSecond()) {
+      assert(SymbolizedFrame.count(Address) > 0 &&
+             "Address not found in SymbolizedFrame map");
+      const SmallVector<MemProfRecord::Frame> &Frames =
+          SymbolizedFrame[Address];
+
+      // The leaf function is the first PC in the callstack so check if unset.
+      if (LeafFunction == 0) {
+        LeafFunction = Frames.front().Function;
+        assert(!Frames.front().IsInlineFrame &&
+               "The first frame should not be inlined");
+      } else {
+        // Record the function this address belongs to. We skip call sites in
+        // the leaf since that call is to the allocator itself.
+        PerFunctionCallSites[Frames.front().Function].insert(Address);
+      }
+      // Add all the frames to the current allocation callstack.
+      Callstack.append(Frames.begin(), Frames.end());
+    }
+
+    auto Result = FunctionProfileData.insert({LeafFunction, MemProfRecord()});
+    MemProfRecord &Record = Result.first->second;
+    Record.AllocSites.emplace_back(Callstack, Entry.second);
+  }
+
+  // Fill in the related contexts per function profile.
+  for (auto I = PerFunctionCallSites.begin(), E = PerFunctionCallSites.end();
+       I != E; I++) {
+    const Function::GUID Id = I->first;
+    // Some functions may have only callsite data and no allocation data. Here
+    // we insert a new entry for callsite data if we need to.
+    auto Result = FunctionProfileData.insert({Id, MemProfRecord()});
+    MemProfRecord &Record = Result.first->second;
+    for (const uint64_t Address : I->getSecond()) {
+      const SmallVector<MemProfRecord::Frame> &Frames =
+          SymbolizedFrame[Address];
+      Record.CallSites.push_back(Frames);
+    }
+  }
+
+  return Error::success();
 }
 
 Error RawMemProfReader::symbolizeAndFilterStackFrames() {
   // The specifier to use when symbolization is requested.
   const DILineInfoSpecifier Specifier(
       DILineInfoSpecifier::FileLineInfoKind::RawValue,
       DILineInfoSpecifier::FunctionNameKind::LinkageName);
 
CallStack.erase(std::remove_if(CallStack.begin(), CallStack.end(),
                     CallStack.end());
     if (CallStack.empty())
       EntriesToErase.push_back(Entry.getFirst());
   }
 
   // Drop the entries where the callstack is empty.
   for (const uint64_t Id : EntriesToErase) {
     StackMap.erase(Id);
-    ProfileData.erase(Id);
+    CallstackProfileData.erase(Id);
   }
 
   if (StackMap.empty())
     return make_error<InstrProfError>(
         instrprof_error::malformed,
         "no entries in callstack map after symbolization");
 
   return Error::success();
if (!SegmentInfo.empty() && SegmentInfo != Entries) {
           "memprof raw profile has different segment information");
     }
     SegmentInfo.assign(Entries.begin(), Entries.end());
 
     // Read in the MemInfoBlocks. Merge them based on stack id - we assume that
     // raw profiles in the same binary file are from the same process so the
     // stackdepot ids are the same.
     for (const auto &Value : readMemInfoBlocks(Next + Header->MIBOffset)) {
-      if (ProfileData.count(Value.first)) {
-        ProfileData[Value.first].Merge(Value.second);
+      if (CallstackProfileData.count(Value.first)) {
+        CallstackProfileData[Value.first].Merge(Value.second);
       } else {
-        ProfileData[Value.first] = Value.second;
+        CallstackProfileData[Value.first] = Value.second;
       }
     }
 
     // Read in the callstack for each ids. For multiple raw profiles in the same
     // file, we expect that the callstack is the same for a unique id.
     const CallStackMap CSM = readStackInfo(Next + Header->StackOffset);
     if (StackMap.empty()) {
       StackMap = CSM;
RawMemProfReader::getModuleOffset(const uint64_t VirtualAddress) {
   assert(ContainingSegment && "Could not find a segment entry");
   });
 
   // TODO: Compute the file offset based on the maps and program headers. For
   // now this only works for non PIE binaries.
   return object::SectionedAddress{VirtualAddress};
 }
 
-Error RawMemProfReader::fillRecord(const uint64_t Id, const MemInfoBlock &MIB,
-                                   MemProfRecord &Record) {
-  auto &CallStack = StackMap[Id];
-  for (const uint64_t Address : CallStack) {
-    assert(SymbolizedFrame.count(Address) &&
-           "Address not found in symbolized frame cache.");
-    Record.CallStack.append(SymbolizedFrame[Address]);
-  }
-  Record.Info = PortableMemInfoBlock(MIB);
-  return Error::success();
-}
-
-Error RawMemProfReader::readNextRecord(MemProfRecord &Record) {
-  if (ProfileData.empty())
+Error RawMemProfReader::readNextRecord(GuidMemProfRecordPair &GuidRecord) {
+  if (FunctionProfileData.empty())
     return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
 
-  if (Iter == ProfileData.end())
+  if (Iter == FunctionProfileData.end())
     return make_error<InstrProfError>(instrprof_error::eof);
 
-  Record.clear();
-  if (Error E = fillRecord(Iter->first, Iter->second, Record)) {
-    return E;
-  }
+  GuidRecord = {Iter->first, Iter->second};
   Iter++;
   return Error::success();
 }
 } // namespace memprof
 } // namespace llvm

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

 CHECK-NEXT:   Header:
 CHECK-NEXT:     Version: 1
 CHECK-NEXT:     TotalSizeBytes: 1016
 CHECK-NEXT:     NumSegments: 9
 CHECK-NEXT:     NumMibInfo: 3
 CHECK-NEXT:     NumStackOffsets: 3
 CHECK-NEXT:   Records:
 CHECK-NEXT:   -
-CHECK-NEXT:     Callstack:
-CHECK-NEXT:     -
-CHECK-NEXT:       Function: {{[0-9]+}}
-CHECK-NEXT:       LineOffset: 1
-CHECK-NEXT:       Column: 21
-CHECK-NEXT:       Inline: 0
-CHECK-NEXT:     MemInfoBlock:
-CHECK-NEXT:       AllocCount: 1
-CHECK-NEXT:       TotalAccessCount: 2
-CHECK-NEXT:       MinAccessCount: 2
-CHECK-NEXT:       MaxAccessCount: 2
-CHECK-NEXT:       TotalSize: 10
-CHECK-NEXT:       MinSize: 10
-CHECK-NEXT:       MaxSize: 10
-CHECK-NEXT:       AllocTimestamp: 986
-CHECK-NEXT:       DeallocTimestamp: 986
-CHECK-NEXT:       TotalLifetime: 0
-CHECK-NEXT:       MinLifetime: 0
-CHECK-NEXT:       MaxLifetime: 0
-CHECK-NEXT:       AllocCpuId: 56
-CHECK-NEXT:       DeallocCpuId: 56
-CHECK-NEXT:       NumMigratedCpu: 0
-CHECK-NEXT:       NumLifetimeOverlaps: 0
-CHECK-NEXT:       NumSameAllocCpu: 0
-CHECK-NEXT:       NumSameDeallocCpu: 0
-CHECK-NEXT:       DataTypeId: {{[0-9]+}}
-CHECK-NEXT:   -
-CHECK-NEXT:     Callstack:
-CHECK-NEXT:     -
-CHECK-NEXT:       Function: {{[0-9]+}}
-CHECK-NEXT:       LineOffset: 5
-CHECK-NEXT:       Column: 15
-CHECK-NEXT:       Inline: 0
-CHECK-NEXT:     MemInfoBlock:
-CHECK-NEXT:       AllocCount: 1
-CHECK-NEXT:       TotalAccessCount: 2
-CHECK-NEXT:       MinAccessCount: 2
-CHECK-NEXT:       MaxAccessCount: 2
-CHECK-NEXT:       TotalSize: 10
-CHECK-NEXT:       MinSize: 10
-CHECK-NEXT:       MaxSize: 10
-CHECK-NEXT:       AllocTimestamp: 987
-CHECK-NEXT:       DeallocTimestamp: 987
-CHECK-NEXT:       TotalLifetime: 0
-CHECK-NEXT:       MinLifetime: 0
-CHECK-NEXT:       MaxLifetime: 0
-CHECK-NEXT:       AllocCpuId: 56
-CHECK-NEXT:       DeallocCpuId: 56
-CHECK-NEXT:       NumMigratedCpu: 0
-CHECK-NEXT:       NumLifetimeOverlaps: 0
-CHECK-NEXT:       NumSameAllocCpu: 0
-CHECK-NEXT:       NumSameDeallocCpu: 0
-CHECK-NEXT:       DataTypeId: {{[0-9]+}}
+CHECK-NEXT:     FunctionGUID: {{[0-9]+}} 
+CHECK-NEXT:     AllocSites:
+CHECK-NEXT:     -
+CHECK-NEXT:       Callstack:
+CHECK-NEXT:       -
+CHECK-NEXT:         Function: {{[0-9]+}}
+CHECK-NEXT:         LineOffset: 1
+CHECK-NEXT:         Column: 21
+CHECK-NEXT:         Inline: 0
+CHECK-NEXT:       MemInfoBlock:
+CHECK-NEXT:         AllocCount: 1
+CHECK-NEXT:         TotalAccessCount: 2
+CHECK-NEXT:         MinAccessCount: 2
+CHECK-NEXT:         MaxAccessCount: 2
+CHECK-NEXT:         TotalSize: 10
+CHECK-NEXT:         MinSize: 10
+CHECK-NEXT:         MaxSize: 10
+CHECK-NEXT:         AllocTimestamp: 986
+CHECK-NEXT:         DeallocTimestamp: 986
+CHECK-NEXT:         TotalLifetime: 0
+CHECK-NEXT:         MinLifetime: 0
+CHECK-NEXT:         MaxLifetime: 0
+CHECK-NEXT:         AllocCpuId: 56
+CHECK-NEXT:         DeallocCpuId: 56
+CHECK-NEXT:         NumMigratedCpu: 0
+CHECK-NEXT:         NumLifetimeOverlaps: 0
+CHECK-NEXT:         NumSameAllocCpu: 0
+CHECK-NEXT:         NumSameDeallocCpu: 0
+CHECK-NEXT:         DataTypeId: {{[0-9]+}}
+CHECK-NEXT:     -
+CHECK-NEXT:       Callstack:
+CHECK-NEXT:       -
+CHECK-NEXT:         Function: {{[0-9]+}}
+CHECK-NEXT:         LineOffset: 5
+CHECK-NEXT:         Column: 15
+CHECK-NEXT:         Inline: 0
+CHECK-NEXT:       MemInfoBlock:
+CHECK-NEXT:         AllocCount: 1
+CHECK-NEXT:         TotalAccessCount: 2
+CHECK-NEXT:         MinAccessCount: 2
+CHECK-NEXT:         MaxAccessCount: 2
+CHECK-NEXT:         TotalSize: 10
+CHECK-NEXT:         MinSize: 10
+CHECK-NEXT:         MaxSize: 10
+CHECK-NEXT:         AllocTimestamp: 987
+CHECK-NEXT:         DeallocTimestamp: 987
+CHECK-NEXT:         TotalLifetime: 0
+CHECK-NEXT:         MinLifetime: 0
+CHECK-NEXT:         MaxLifetime: 0
+CHECK-NEXT:         AllocCpuId: 56
+CHECK-NEXT:         DeallocCpuId: 56
+CHECK-NEXT:         NumMigratedCpu: 0
+CHECK-NEXT:         NumLifetimeOverlaps: 0
+CHECK-NEXT:         NumSameAllocCpu: 0
+CHECK-NEXT:         NumSameDeallocCpu: 0
+CHECK-NEXT:         DataTypeId: {{[0-9]+}}

llvm/tools/llvm-profdata/llvm-profdata.cpp

if (Error E = WC->Writer.mergeProfileKind(Reader->getProfileKind())) {
           make_error<StringError>(
               "Cannot merge MemProf profile with Clang generated profile.",
               std::error_code()),
           Filename);
       return;
     }
 
     // Add the records into the writer context.
-    for (const memprof::MemProfRecord &MR : *Reader) {
-      WC->Writer.addRecord(MR, [&](Error E) {
+    for (auto I = Reader->begin(), E = Reader->end(); I != E; ++I) {
+      WC->Writer.addRecord(/*Id=*/I->first, /*Record=*/I->second, [&](Error E) {
         instrprof_error IPE = InstrProfError::take(std::move(E));
         WC->Errors.emplace_back(make_error<InstrProfError>(IPE), Filename);
       });
     }
     return;
   }
 
   auto ReaderOrErr = InstrProfReader::create(Input.Filename, Correlator);

llvm/unittests/ProfileData/InstrProfTest.cpp

 //===- unittest/ProfileData/InstrProfTest.cpp -------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/InstrProfWriter.h"
 #include "llvm/ProfileData/MemProf.h"
+#include "llvm/ProfileData/MemProfData.inc"
 #include "llvm/Support/Compression.h"
 #include "llvm/Testing/Support/Error.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gtest/gtest.h"
 #include <cstdarg>
 
 using namespace llvm;
 
TEST_F(InstrProfTest, test_writer_merge) {
 
   R = Reader->getInstrProfRecord("func2", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(2U, R->Counts.size());
   ASSERT_EQ(0U, R->Counts[0]);
   ASSERT_EQ(0U, R->Counts[1]);
 }
 
+using ::llvm::memprof::MemInfoBlock;
+using ::llvm::memprof::MemProfRecord;
+MemProfRecord makeRecord(std::initializer_list<MemProfRecord::Frame> Frames,
+                         const MemInfoBlock &Block = MemInfoBlock()) {
+  llvm::memprof::MemProfRecord MR;
+  MR.AllocSites.emplace_back(Frames, Block);
+  return MR;
+}
+
 TEST_F(InstrProfTest, test_memprof) {
   ASSERT_THAT_ERROR(Writer.mergeProfileKind(InstrProfKind::MemProf),
                     Succeeded());
-  llvm::memprof::MemProfRecord MR;
-  MR.CallStack.push_back({0x123, 1, 2, false});
-  MR.CallStack.push_back({0x345, 3, 4, true});
-  Writer.addRecord(MR, Err);
+
+  const MemProfRecord MR =
+      makeRecord({{0x123, 1, 2, false}, {0x345, 3, 4, true}});
+  Writer.addRecord(0x123, MR, Err);
 
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
   auto RecordsOr = Reader->getMemProfRecord(0x123);
   ASSERT_THAT_ERROR(RecordsOr.takeError(), Succeeded());
   const auto Records = RecordsOr.get();
   ASSERT_EQ(Records.size(), 1U);
   EXPECT_EQ(Records[0], MR);
 }
 
 TEST_F(InstrProfTest, test_memprof_merge) {
   Writer.addRecord({"func1", 0x1234, {42}}, Err);
 
   InstrProfWriter Writer2;
   ASSERT_THAT_ERROR(Writer2.mergeProfileKind(InstrProfKind::MemProf),
                     Succeeded());
 
-  llvm::memprof::MemProfRecord MR;
-  MR.CallStack.push_back({0x123, 1, 2, false});
-  MR.CallStack.push_back({0x345, 3, 4, true});
-  Writer2.addRecord(MR, Err);
+  const MemProfRecord MR =
+      makeRecord({{0x123, 1, 2, false}, {0x345, 3, 4, true}});
+  Writer2.addRecord(0x123, MR, Err);
 
   ASSERT_THAT_ERROR(Writer.mergeProfileKind(Writer2.getProfileKind()),
                     Succeeded());
   Writer.mergeRecordsFromWriter(std::move(Writer2), Err);
 
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
   Expected<InstrProfRecord> R = Reader->getInstrProfRecord("func1", 0x1234);
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   ASSERT_EQ(1U, R->Counts.size());
   ASSERT_EQ(42U, R->Counts[0]);
 
   auto RecordsOr = Reader->getMemProfRecord(0x123);
   ASSERT_THAT_ERROR(RecordsOr.takeError(), Succeeded());
   const auto Records = RecordsOr.get();
   ASSERT_EQ(Records.size(), 1U);
   EXPECT_EQ(Records[0], MR);
 }
 
-TEST_F(InstrProfTest, test_memprof_invalid_add_record) {
-  llvm::memprof::MemProfRecord MR;
-  // At least one of the frames should be a non-inline frame.
-  MR.CallStack.push_back({0x123, 1, 2, true});
-  MR.CallStack.push_back({0x345, 3, 4, true});
-
-  auto CheckErr = [](Error &&E) {
-    EXPECT_TRUE(ErrorEquals(instrprof_error::invalid_prof, std::move(E)));
-  };
-  Writer.addRecord(MR, CheckErr);
-}
+// TEST_F(InstrProfTest, test_memprof_invalid_add_record) {
+//   llvm::memprof::MemProfRecord MR;
+//   // At least one of the frames should be a non-inline frame.
+//   MR.CallStack.push_back({0x123, 1, 2, true});
+//   MR.CallStack.push_back({0x345, 3, 4, true});
+//
+//   auto CheckErr = [](Error &&E) {
+//     EXPECT_TRUE(ErrorEquals(instrprof_error::invalid_prof, std::move(E)));
+//   };
+//   Writer.addRecord(MR, CheckErr);
+// }
 
 static const char callee1[] = "callee1";
 static const char callee2[] = "callee2";
 static const char callee3[] = "callee3";
 static const char callee4[] = "callee4";
 static const char callee5[] = "callee5";
 static const char callee6[] = "callee6";
 

llvm/unittests/ProfileData/MemProfTest.cpp

TEST(MemProf, FillsValue) {
   EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x2000},
                                                 specifier(), false))
       .Times(1) // Only once since we cache the result for future lookups.
       .WillRepeatedly(Return(makeInliningInfo({
           {"foo", 10, 5, 30},
           {"bar", 201, 150, 20},
       })));
 
+  EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x4000},
+                                                specifier(), false))
+      .Times(1)
+      .WillRepeatedly(Return(makeInliningInfo({
+          {"xyz", 10, 5, 30},
+          {"abc", 10, 5, 30},
+      })));
+
   EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x6000},
                                                 specifier(), false))
       .Times(1)
       .WillRepeatedly(Return(makeInliningInfo({
           {"baz", 10, 5, 30},
           {"qux.llvm.12345", 75, 70, 10},
       })));
 
   CallStackMap CSM;
-  CSM[0x1] = {0x2000};
+  CSM[0x1] = {0x2000, 0x4000};
   CSM[0x2] = {0x6000, 0x2000};
 
   llvm::MapVector<uint64_t, MemInfoBlock> Prof;
   Prof[0x1].AllocCount = 1;
   Prof[0x2].AllocCount = 2;
 
   auto Seg = makeSegments();
 
   RawMemProfReader Reader(std::move(Symbolizer), Seg, Prof, CSM);
 
-  std::vector<MemProfRecord> Records;
-  for (const MemProfRecord &R : Reader) {
-    Records.push_back(R);
-  }
-  EXPECT_EQ(Records.size(), 2U);
-
-  EXPECT_EQ(Records[0].Info.getAllocCount(), 1U);
-  EXPECT_EQ(Records[1].Info.getAllocCount(), 2U);
-  EXPECT_THAT(Records[0].CallStack[0], FrameContains("foo", 5U, 30U, false));
-  EXPECT_THAT(Records[0].CallStack[1], FrameContains("bar", 51U, 20U, true));
-
-  EXPECT_THAT(Records[1].CallStack[0], FrameContains("baz", 5U, 30U, false));
-  EXPECT_THAT(Records[1].CallStack[1], FrameContains("qux", 5U, 10U, true));
-  EXPECT_THAT(Records[1].CallStack[2], FrameContains("foo", 5U, 30U, false));
-  EXPECT_THAT(Records[1].CallStack[3], FrameContains("bar", 51U, 20U, true));
+  llvm::SmallVector<MemProfRecord, 3> Records;
+  for (const auto &KeyRecordPair : Reader) {
+    Records.push_back(KeyRecordPair.second);
+  }
+  // The PCs in each callstack have the leaf function at index zero. Since
+  // memprof records are grouped by leaf function, here we expect 2 records for
+  // functions with allocation contexts. 1 additional record with only callsite
+  // data is expected for the first profile entry.
+  ASSERT_EQ(Records.size(), 3U);
+  ASSERT_EQ(Records[0].AllocSites.size(), 1U);
+  ASSERT_EQ(Records[1].AllocSites.size(), 1U);
+
+  EXPECT_EQ(Records[0].AllocSites[0].Info.getAllocCount(), 1U);
+  EXPECT_EQ(Records[1].AllocSites[0].Info.getAllocCount(), 2U);
+
+  EXPECT_THAT(Records[0].AllocSites[0].CallStack[0],
+              FrameContains("foo", 5U, 30U, false));
+  EXPECT_THAT(Records[0].AllocSites[0].CallStack[1],
+              FrameContains("bar", 51U, 20U, true));
+
+  EXPECT_THAT(Records[1].AllocSites[0].CallStack[0],
+              FrameContains("baz", 5U, 30U, false));
+  EXPECT_THAT(Records[1].AllocSites[0].CallStack[1],
+              FrameContains("qux", 5U, 10U, true));
+  EXPECT_THAT(Records[1].AllocSites[0].CallStack[2],
+              FrameContains("foo", 5U, 30U, false));
+  EXPECT_THAT(Records[1].AllocSites[0].CallStack[3],
+              FrameContains("bar", 51U, 20U, true));
+
+  // Check the callsite frames.
+  EXPECT_TRUE(Records[2].AllocSites.empty());
+  ASSERT_EQ(Records[2].CallSites.size(), 1U);
+  // We expect one actual frame and one inline frame.
+  ASSERT_EQ(Records[2].CallSites[0].size(), 2U);
+  EXPECT_THAT(Records[2].CallSites[0][0], FrameContains("xyz", 5U, 30U, false));
+  EXPECT_THAT(Records[2].CallSites[0][1], FrameContains("abc", 5U, 30U, true));
 }
 
 TEST(MemProf, PortableWrapper) {
   MemInfoBlock Info(/*size=*/16, /*access_count=*/7, /*alloc_timestamp=*/1000,
                     /*dealloc_timestamp=*/2000, /*alloc_cpu=*/3,
                     /*dealloc_cpu=*/4);
 
   const auto Schema = getFullSchema();
TEST(MemProf, PortableWrapper) {
   EXPECT_EQ(1UL, ReadBlock.getAllocCount());
   EXPECT_EQ(7ULL, ReadBlock.getTotalAccessCount());
   EXPECT_EQ(3UL, ReadBlock.getAllocCpuId());
 }
 
 TEST(MemProf, RecordSerializationRoundTrip) {
   const MemProfSchema Schema = getFullSchema();
 
-  llvm::SmallVector<MemProfRecord, 3> Records;
-  MemProfRecord MR;
-
   MemInfoBlock Info(/*size=*/16, /*access_count=*/7, /*alloc_timestamp=*/1000,
                     /*dealloc_timestamp=*/2000, /*alloc_cpu=*/3,
                     /*dealloc_cpu=*/4);
 
-  MR.Info = PortableMemInfoBlock(Info);
-  MR.CallStack.push_back({0x123, 1, 2, false});
-  MR.CallStack.push_back({0x345, 3, 4, false});
-  Records.push_back(MR);
-
-  MR.clear();
-  MR.Info = PortableMemInfoBlock(Info);
-  MR.CallStack.push_back({0x567, 5, 6, false});
-  MR.CallStack.push_back({0x789, 7, 8, false});
-  Records.push_back(MR);
+  llvm::SmallVector<llvm::SmallVector<MemProfRecord::Frame>> AllocCallStacks = {
+      {{0x123, 1, 2, false}, {0x345, 3, 4, false}},
+      {{0x123, 1, 2, false}, {0x567, 5, 6, false}}};
+
+  llvm::SmallVector<llvm::SmallVector<MemProfRecord::Frame>> CallSites = {
+      {{0x333, 1, 2, false}, {0x777, 3, 4, true}}};
+
+  MemProfRecord Record;
+  for (const auto &ACS : AllocCallStacks) {
+    // Use the same info block for both allocation sites.
+    Record.AllocSites.emplace_back(ACS, Info);
+  }
+  Record.CallSites.assign(CallSites);
 
   std::string Buffer;
   llvm::raw_string_ostream OS(Buffer);
-  serializeRecords(Records, Schema, OS);
+  Record.serialize(Schema, OS);
   OS.flush();
 
-  const llvm::SmallVector<MemProfRecord, 4> GotRecords = deserializeRecords(
+  const MemProfRecord GotRecord = MemProfRecord::deserialize(
       Schema, reinterpret_cast<const unsigned char *>(Buffer.data()));
 
-  ASSERT_TRUE(!GotRecords.empty());
-  EXPECT_EQ(GotRecords.size(), Records.size());
-  EXPECT_THAT(GotRecords[0], EqualsRecord(Records[0]));
-  EXPECT_THAT(GotRecords[1], EqualsRecord(Records[1]));
+  EXPECT_THAT(GotRecord, EqualsRecord(Record));
 }
 
 TEST(MemProf, SymbolizationFilter) {
   std::unique_ptr<MockSymbolizer> Symbolizer(new MockSymbolizer());
 
   EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x1000},
                                                 specifier(), false))
       .Times(1) // once since we don't lookup invalid PCs repeatedly.
TEST(MemProf, SymbolizationFilter) {
   llvm::MapVector<uint64_t, MemInfoBlock> Prof;
   Prof[0x1].AllocCount = 1;
   Prof[0x2].AllocCount = 1;
 
   auto Seg = makeSegments();
 
   RawMemProfReader Reader(std::move(Symbolizer), Seg, Prof, CSM);
 
-  std::vector<MemProfRecord> Records;
-  for (const MemProfRecord &R : Reader) {
-    Records.push_back(R);
+  llvm::SmallVector<MemProfRecord, 1> Records;
+  for (const auto &KeyRecordPair : Reader) {
+    Records.push_back(KeyRecordPair.second);
   }
+
   ASSERT_EQ(Records.size(), 1U);
-  ASSERT_EQ(Records[0].CallStack.size(), 1U);
-  EXPECT_THAT(Records[0].CallStack[0], FrameContains("foo", 5U, 30U, false));
+  ASSERT_EQ(Records[0].AllocSites.size(), 1U);
+  ASSERT_EQ(Records[0].AllocSites[0].CallStack.size(), 1U);
+  EXPECT_THAT(Records[0].AllocSites[0].CallStack[0],
+              FrameContains("foo", 5U, 30U, false));
 }
 } // namespace