[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

209

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

311

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

104

missing word(s) after "that"?

llvm/lib/ProfileData/RawMemProfReader.cpp

296

Document what the uint64_t values are.

319

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?

327

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?

303–304

Why is this test commented out now?

llvm/unittests/ProfileData/MemProfTest.cpp

188

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

296

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.

319

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

303–304

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

188

Sure, added the check below.

[tejohnson]

llvm/lib/ProfileData/RawMemProfReader.cpp

319

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

259

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

llvm/lib/ProfileData/RawMemProfReader.cpp

294

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

328

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 ]

340

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

172

Is this printing meant to be here?

178

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

236

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/IR/GlobalValue.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>
 
 namespace llvm {
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<GlobalValue::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 GlobalValue::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.
     }
 
+    bool operator==(const AllocationInfo &Other) const {
+      if (Other.Info != Info)
+        return false;
+
+      if (Other.CallStack.size() != CallStack.size())
+        return false;
+
+      for (size_t J = 0; J < Other.CallStack.size(); J++) {
+        if (Other.CallStack[J] != CallStack[J])
+          return false;
+      }
+      return true;
+    }
+
+    bool operator!=(const AllocationInfo &Other) const {
+      return !operator==(Other);
+    }
+  };
+
+  // Memory allocation sites in this function for which we have memory profiling
+  // data.
+  llvm::SmallVector<AllocationInfo> AllocSites;
+  // Holds call sites in this function which are part of some memory allocation
+  // context. We store this as a list of locations, each with its 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<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(GlobalValue::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())
+    if (Other.CallSites.size() != CallSites.size())
       return false;
 
-    for (size_t I = 0; I < Other.CallStack.size(); I++) {
-      if (Other.CallStack[I] != CallStack[I])
+    for (size_t I = 0; I < AllocSites.size(); I++) {
+      if (AllocSites[I] != Other.AllocSites[I])
+        return false;
+    }
+
+    for (size_t I = 0; I < CallSites.size(); I++) {
+      if (CallSites[I] != Other.CallSites[I])
         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,
 
+  // Returns the GUID for the function name after canonicalization. For memprof,
+  // we remove any .llvm suffix added by LTO. MemProfRecords are mapped to
+  // functions using this GUID.
+  static GlobalValue::GUID getGUID(const StringRef FunctionName);
+};
+
 // 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/IR/GlobalValue.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"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 
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<GlobalValue::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 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<GlobalValue::GUID, MemProfRecord> FunctionProfileData;
+  llvm::MapVector<GlobalValue::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/IR/Function.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);
   }
-  return Records;
+
+  // 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 Record;
+}
+
+GlobalValue::GUID MemProfRecord::getGUID(const StringRef FunctionName) {
+  const auto Pos = FunctionName.find(".llvm.");
+
+  // We use the function guid which we expect to be a uint64_t. At
+  // this time, it is the lower 64 bits of the md5 of the function
+  // name. Any suffix with .llvm. is trimmed since these are added by
+  // thinLTO global promotion. At the time the profile is consumed,
+  // these suffixes will not be present.
+  return Function::getGUID(FunctionName.take_front(Pos));
 }
 
 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

 //===- RawMemProfReader.cpp - Instrumented memory profiling reader --------===//
 //
 // 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 <algorithm>
 #include <cstdint>
 #include <type_traits>
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableModule.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableObjectFile.h"
-#include "llvm/IR/Function.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/MemProf.h"
 #include "llvm/ProfileData/MemProfData.inc"
 #include "llvm/ProfileData/RawMemProfReader.h"
 #include "llvm/Support/Endian.h"
if (I == To.end()) {
       // Check that the PCs are the same (in order).
       if (IdStack.second != I->second)
         return true;
     }
   }
   return false;
 }
 
-StringRef trimSuffix(const StringRef Name) {
-  const auto Pos = Name.find(".llvm.");
-  return Name.take_front(Pos);
-}
-
 Error report(Error E, const StringRef Context) {
   return joinErrors(createStringError(inconvertibleErrorCode(), Context),
                     std::move(E));
 }
 
 bool isRuntimePath(const StringRef Path) {
   return StringRef(llvm::sys::path::convert_to_slash(Path))
       .contains("memprof/memprof_");

 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() {
+  // Hold a mapping from function to each callsite location we encounter within
+  // it that is part of some dynamic allocation context. The location is stored
+  // as a pointer to a symbolized list of inline frames.
+  using LocationPtr = const llvm::SmallVector<MemProfRecord::Frame> *;
+  llvm::DenseMap<GlobalValue::GUID, llvm::SetVector<LocationPtr>>
+      PerFunctionCallSites;
+
+  // 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.
+  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());
+
+    llvm::ArrayRef<uint64_t> Addresses = It->getSecond();
+    for (size_t I = 0; I < Addresses.size(); I++) {
+      const uint64_t Address = Addresses[I];
+      assert(SymbolizedFrame.count(Address) > 0 &&
+             "Address not found in SymbolizedFrame map");
+      const SmallVector<MemProfRecord::Frame> &Frames =
+          SymbolizedFrame[Address];
+
+      assert(!Frames.back().IsInlineFrame &&
+             "The last frame should not be inlined");
+
+      // Record the callsites for each function. Skip the first frame of the
+      // first address since it is the allocation site itself that is recorded
+      // as an alloc site.
+      for (size_t J = 0; J < Frames.size(); J++) {
+        if (I == 0 && J == 0)
+          continue;
+        // We attach the entire bottom-up frame here for the callsite even
+        // though we only need the frames up to and including the frame for
+        // Frames[J].Function. This will enable better deduplication for
+        // compression in the future.
+        PerFunctionCallSites[Frames[J].Function].insert(&Frames);
+      }
+
+      // Add all the frames to the current allocation callstack.
+      Callstack.append(Frames.begin(), Frames.end());
+    }
+
+    // We attach the memprof record to each function bottom-up including the
+    // first non-inline frame.
+    for (size_t I = 0; /*Break out using the condition below*/; I++) {
+      auto Result =
+          FunctionProfileData.insert({Callstack[I].Function, MemProfRecord()});
+      MemProfRecord &Record = Result.first->second;
+      Record.AllocSites.emplace_back(Callstack, Entry.second);
+
+      if (!Callstack[I].IsInlineFrame)
+        break;
+    }
+  }
+
+  // Fill in the related callsites per function.
+  for (auto I = PerFunctionCallSites.begin(), E = PerFunctionCallSites.end();
+       I != E; I++) {
+    const GlobalValue::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 (LocationPtr Loc : I->getSecond()) {
+      Record.CallSites.push_back(*Loc);
+    }
+  }
+
+  return Error::success();
 }
 
 Error RawMemProfReader::symbolizeAndFilterStackFrames() {
   // The specifier to use when symbolization is requested.
   const DILineInfoSpecifier Specifier(
       DILineInfoSpecifier::FileLineInfoKind::RawValue,
       DILineInfoSpecifier::FunctionNameKind::LinkageName);
 
for (const uint64_t VAddr : Entry.getSecond()) {
       }
 
       for (size_t I = 0, NumFrames = DI.getNumberOfFrames(); I < NumFrames;
            I++) {
         const auto &Frame = DI.getFrame(I);
         LLVM_DEBUG(
             // Print out the name to guid mapping for debugging.
             llvm::dbgs() << "FunctionName: " << Frame.FunctionName << " GUID: "
-                         << Function::getGUID(trimSuffix(Frame.FunctionName))
+                         << MemProfRecord::getGUID(Frame.FunctionName)
                          << "\n";);
         SymbolizedFrame[VAddr].emplace_back(
-            // We use the function guid which we expect to be a uint64_t. At
-            // this time, it is the lower 64 bits of the md5 of the function
-            // name. Any suffix with .llvm. is trimmed since these are added by
-            // thinLTO global promotion. At the time the profile is consumed,
-            // these suffixes will not be present.
-            Function::getGUID(trimSuffix(Frame.FunctionName)),
+            MemProfRecord::getGUID(Frame.FunctionName),
             Frame.Line - Frame.StartLine, Frame.Column,
             // Only the last entry is not an inlined location.
             I != NumFrames - 1);
       }
     }
 
     auto &CallStack = Entry.getSecond();
     CallStack.erase(std::remove_if(CallStack.begin(), CallStack.end(),
                                    [&AllVAddrsToDiscard](const uint64_t A) {
                                      return AllVAddrsToDiscard.contains(A);
                                    }),
                     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

 virtual addresses for the callstack, we do not expect the raw binary profile to
 be deterministic. The summary should be deterministic apart from changes to
 the shared libraries linked in which could change the number of segments
 recorded.
 
 ```
 clang -fuse-ld=lld -Wl,--no-rosegment -gmlt -fdebug-info-for-profiling \
       -fmemory-profile -mno-omit-leaf-frame-pointer -fno-omit-frame-pointer \
-      -fno-optimize-sibling-calls -m64 -Wl,-build-id source.c -o basic.memprofexe 
+      -fno-optimize-sibling-calls -m64 -Wl,-build-id source.c -o basic.memprofexe
 
 env MEMPROF_OPTIONS=log_path=stdout ./rawprofile.out > basic.memprofraw
 ```
 
 RUN: llvm-profdata show --memory %p/Inputs/basic.memprofraw --profiled-binary %p/Inputs/basic.memprofexe -o - | FileCheck %s
 
 We expect 2 MIB entries, 1 each for the malloc calls in the program. Any
 additional allocations which do not originate from the main binary are pruned.
 
 CHECK: MemprofProfile:
 CHECK-NEXT:   -
 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/test/tools/llvm-profdata/memprof-inline.test

 FunctionName: qux GUID: 15505678318020221912
 FunctionName: foo GUID: 6699318081062747564
 FunctionName: bar GUID: 16434608426314478903
 FunctionName: main GUID: 15822663052811949562
 
 [..omit output here which is checked below..]
 ```
 
-RUN: llvm-profdata show --memory %p/Inputs/inline.memprofraw --profiled-binary %p/Inputs/memprof-inline.exe
+RUN: llvm-profdata show --memory %p/Inputs/inline.memprofraw --profiled-binary %p/Inputs/inline.memprofexe | FileCheck %s
 
 CHECK:  MemprofProfile:
 CHECK-NEXT:  -
 CHECK-NEXT:  Header:
 CHECK-NEXT:    Version: 1
 CHECK-NEXT:    TotalSizeBytes: 880
 CHECK-NEXT:    NumSegments: 9
 CHECK-NEXT:    NumMibInfo: 2
 CHECK-NEXT:    NumStackOffsets: 2
 CHECK-NEXT:  Records:
 CHECK-NEXT:  -
-CHECK-NEXT:    Callstack:
-CHECK-NEXT:    -
-CHECK-NEXT:      Function: 15505678318020221912
-CHECK-NEXT:      LineOffset: 1
-CHECK-NEXT:      Column: 15
-CHECK-NEXT:      Inline: 0
-CHECK-NEXT:    -
-CHECK-NEXT:      Function: 6699318081062747564
-CHECK-NEXT:      LineOffset: 0
-CHECK-NEXT:      Column: 18
-CHECK-NEXT:      Inline: 1
-CHECK-NEXT:    -
-CHECK-NEXT:      Function: 16434608426314478903
-CHECK-NEXT:      LineOffset: 0
-CHECK-NEXT:      Column: 19
-CHECK-NEXT:      Inline: 0
-CHECK-NEXT:    -
-CHECK-NEXT:      Function: 15822663052811949562
-CHECK-NEXT:      LineOffset: 1
-CHECK-NEXT:      Column: 3
-CHECK-NEXT:      Inline: 0
-CHECK-NEXT:    MemInfoBlock:
-CHECK-NEXT:      AllocCount: 1
-CHECK-NEXT:      TotalAccessCount: 1
-CHECK-NEXT:      MinAccessCount: 1
-CHECK-NEXT:      MaxAccessCount: 1
-CHECK-NEXT:      TotalSize: 1
-CHECK-NEXT:      MinSize: 1
-CHECK-NEXT:      MaxSize: 1
-CHECK-NEXT:      AllocTimestamp: 894
-CHECK-NEXT:      DeallocTimestamp: 894
-CHECK-NEXT:      TotalLifetime: 0
-CHECK-NEXT:      MinLifetime: 0
-CHECK-NEXT:      MaxLifetime: 0
-CHECK-NEXT:      AllocCpuId: 23
-CHECK-NEXT:      DeallocCpuId: 23
-CHECK-NEXT:      NumMigratedCpu: 0
-CHECK-NEXT:      NumLifetimeOverlaps: 0
-CHECK-NEXT:      NumSameAllocCpu: 0
-CHECK-NEXT:      NumSameDeallocCpu: 0
-CHECK-NEXT:      DataTypeId: {{[0-9]+}}
-
+CHECK-NEXT:    FunctionGUID: 15505678318020221912
+CHECK-NEXT:    AllocSites:
+CHECK-NEXT:    -
+CHECK-NEXT:      Callstack:
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15505678318020221912
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 15
+CHECK-NEXT:        Inline: 1
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 6699318081062747564
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 18
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 16434608426314478903
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 19
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15822663052811949562
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 3
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      MemInfoBlock:
+CHECK-NEXT:        AllocCount: 1
+CHECK-NEXT:        TotalAccessCount: 1
+CHECK-NEXT:        MinAccessCount: 1
+CHECK-NEXT:        MaxAccessCount: 1
+CHECK-NEXT:        TotalSize: 1
+CHECK-NEXT:        MinSize: 1
+CHECK-NEXT:        MaxSize: 1
+CHECK-NEXT:        AllocTimestamp: 894
+CHECK-NEXT:        DeallocTimestamp: 894
+CHECK-NEXT:        TotalLifetime: 0
+CHECK-NEXT:        MinLifetime: 0
+CHECK-NEXT:        MaxLifetime: 0
+CHECK-NEXT:        AllocCpuId: 23
+CHECK-NEXT:        DeallocCpuId: 23
+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:    FunctionGUID: 6699318081062747564
+CHECK-NEXT:    AllocSites:
+CHECK-NEXT:    -
+CHECK-NEXT:      Callstack:
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15505678318020221912
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 15
+CHECK-NEXT:        Inline: 1
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 6699318081062747564
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 18
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 16434608426314478903
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 19
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15822663052811949562
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 3
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:      MemInfoBlock:
+CHECK-NEXT:        AllocCount: 1
+CHECK-NEXT:        TotalAccessCount: 1
+CHECK-NEXT:        MinAccessCount: 1
+CHECK-NEXT:        MaxAccessCount: 1
+CHECK-NEXT:        TotalSize: 1
+CHECK-NEXT:        MinSize: 1
+CHECK-NEXT:        MaxSize: 1
+CHECK-NEXT:        AllocTimestamp: 894
+CHECK-NEXT:        DeallocTimestamp: 894
+CHECK-NEXT:        TotalLifetime: 0
+CHECK-NEXT:        MinLifetime: 0
+CHECK-NEXT:        MaxLifetime: 0
+CHECK-NEXT:        AllocCpuId: 23
+CHECK-NEXT:        DeallocCpuId: 23
+CHECK-NEXT:        NumMigratedCpu: 0
+CHECK-NEXT:        NumLifetimeOverlaps: 0
+CHECK-NEXT:        NumSameAllocCpu: 0
+CHECK-NEXT:        NumSameDeallocCpu: 0
+CHECK-NEXT:        DataTypeId: {{[0-9]+}}
+CHECK-NEXT:    CallSites:
+CHECK-NEXT:    -
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15505678318020221912
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 15
+CHECK-NEXT:        Inline: 1
+CHECK-NEXT:    -
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 6699318081062747564
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 18
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:  -
+CHECK-NEXT:    FunctionGUID: 15822663052811949562
+CHECK-NEXT:    CallSites:
+CHECK-NEXT:    -
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 15822663052811949562
+CHECK-NEXT:        LineOffset: 1
+CHECK-NEXT:        Column: 3
+CHECK-NEXT:        Inline: 0
+CHECK-NEXT:  -
+CHECK-NEXT:    FunctionGUID: 16434608426314478903
+CHECK-NEXT:    CallSites:
+CHECK-NEXT:    -
+CHECK-NEXT:      -
+CHECK-NEXT:        Function: 16434608426314478903
+CHECK-NEXT:        LineOffset: 0
+CHECK-NEXT:        Column: 19
+CHECK-NEXT:        Inline: 0

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<std::initializer_list<MemProfRecord::Frame>>
+               AllocFrames,
+           std::initializer_list<std::initializer_list<MemProfRecord::Frame>>
+               CallSiteFrames,
+           const MemInfoBlock &Block = MemInfoBlock()) {
+  llvm::memprof::MemProfRecord MR;
+  for (const auto &Frames : AllocFrames)
+    MR.AllocSites.emplace_back(Frames, Block);
+  for (const auto &Frames : CallSiteFrames)
+    MR.CallSites.push_back(Frames);
+  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(
+      /*AllocFrames=*/
+      {
+          {{0x123, 1, 2, false}, {0x345, 3, 4, true}},
+          {{0x125, 5, 6, false}, {0x567, 7, 8, true}},
+      },
+      /*CallSiteFrames=*/{
+          {{0x124, 5, 6, false}, {0x789, 8, 9, true}},
+      });
+  Writer.addRecord(/*Id=*/0x9999, MR, Err);
 
   auto Profile = Writer.writeBuffer();
   readProfile(std::move(Profile));
 
-  auto RecordsOr = Reader->getMemProfRecord(0x123);
+  auto RecordsOr = Reader->getMemProfRecord(0x9999);
   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(
+      /*AllocFrames=*/
+      {
+          {{0x123, 1, 2, false}, {0x345, 3, 4, true}},
+          {{0x125, 5, 6, false}, {0x567, 7, 8, true}},
+      },
+      /*CallSiteFrames=*/{
+          {{0x124, 5, 6, false}, {0x789, 8, 9, true}},
+      });
+  Writer2.addRecord(/*Id=*/0x9999, 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);
+  auto RecordsOr = Reader->getMemProfRecord(0x9999);
   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);
-}
-
 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";
 
 TEST_P(MaybeSparseInstrProfTest, get_icall_data_read_write) {
   NamedInstrProfRecord Record1("caller", 0x1234, {1, 2});

llvm/unittests/ProfileData/MemProfTest.cpp

 #include "llvm/ProfileData/MemProf.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/DebugInfo/DIContext.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableModule.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/Value.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/MemProfData.inc"
 #include "llvm/ProfileData/RawMemProfReader.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gmock/gmock.h"

 #include "llvm/ProfileData/MIBEntryDef.inc"
 #undef MIBEntryDef
   return Schema;
 }
 
 TEST(MemProf, FillsValue) {
   std::unique_ptr<MockSymbolizer> Symbolizer(new MockSymbolizer());
 
+  EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x1000},
+                                                specifier(), false))
+      .Times(1) // Only once since we remember invalid PCs.
+      .WillRepeatedly(Return(makeInliningInfo({
+          {"new", 70, 57, 3, "memprof/memprof_new_delete.cpp"},
+      })));
+
   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{0x6000},
+  EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x3000},
                                                 specifier(), false))
       .Times(1)
       .WillRepeatedly(Return(makeInliningInfo({
-          {"baz", 10, 5, 30},
-          {"qux.llvm.12345", 75, 70, 10},
+          {"xyz", 10, 5, 30},
+          {"abc", 10, 5, 30},
       })));
 
   CallStackMap CSM;
-  CSM[0x1] = {0x2000};
-  CSM[0x2] = {0x6000, 0x2000};
+  CSM[0x1] = {0x1000, 0x2000, 0x3000};
 
   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, true));
-  EXPECT_THAT(Records[0].CallStack[1], FrameContains("bar", 51U, 20U, false));
-
-  EXPECT_THAT(Records[1].CallStack[0], FrameContains("baz", 5U, 30U, true));
-  EXPECT_THAT(Records[1].CallStack[1], FrameContains("qux", 5U, 10U, false));
-  EXPECT_THAT(Records[1].CallStack[2], FrameContains("foo", 5U, 30U, true));
-  EXPECT_THAT(Records[1].CallStack[3], FrameContains("bar", 51U, 20U, false));
+  llvm::DenseMap<llvm::GlobalValue::GUID, MemProfRecord> Records;
+  for (const auto &Pair : Reader) {
+    Records.insert({Pair.first, Pair.second});
+  }
+
+  // Mock program psuedocode and expected memprof record contents.
+  //
+  //                              AllocSite       CallSite
+  // inline foo() { new(); }         Y               N
+  // bar() { foo(); }                Y               Y
+  // inline xyz() { bar(); }         N               Y
+  // abc() { xyz(); }                N               Y
+
+  // We expect 4 records. We attach alloc site data to foo and bar, i.e.
+  // all frames bottom up until we find a non-inline frame. We attach call site
+  // data to bar, xyz and abc.
+  ASSERT_EQ(Records.size(), 4U);
+
+  // Check the memprof record for foo.
+  const llvm::GlobalValue::GUID FooId = MemProfRecord::getGUID("foo");
+  ASSERT_EQ(Records.count(FooId), 1U);
+  const MemProfRecord &Foo = Records[FooId];
+  ASSERT_EQ(Foo.AllocSites.size(), 1U);
+  EXPECT_EQ(Foo.AllocSites[0].Info.getAllocCount(), 1U);
+  EXPECT_THAT(Foo.AllocSites[0].CallStack[0],
+              FrameContains("foo", 5U, 30U, true));
+  EXPECT_THAT(Foo.AllocSites[0].CallStack[1],
+              FrameContains("bar", 51U, 20U, false));
+  EXPECT_THAT(Foo.AllocSites[0].CallStack[2],
+              FrameContains("xyz", 5U, 30U, true));
+  EXPECT_THAT(Foo.AllocSites[0].CallStack[3],
+              FrameContains("abc", 5U, 30U, false));
+  EXPECT_TRUE(Foo.CallSites.empty());
+
+  // Check the memprof record for bar.
+  const llvm::GlobalValue::GUID BarId = MemProfRecord::getGUID("bar");
+  ASSERT_EQ(Records.count(BarId), 1U);
+  const MemProfRecord &Bar = Records[BarId];
+  ASSERT_EQ(Bar.AllocSites.size(), 1U);
+  EXPECT_EQ(Bar.AllocSites[0].Info.getAllocCount(), 1U);
+  EXPECT_THAT(Bar.AllocSites[0].CallStack[0],
+              FrameContains("foo", 5U, 30U, true));
+  EXPECT_THAT(Bar.AllocSites[0].CallStack[1],
+              FrameContains("bar", 51U, 20U, false));
+  EXPECT_THAT(Bar.AllocSites[0].CallStack[2],
+              FrameContains("xyz", 5U, 30U, true));
+  EXPECT_THAT(Bar.AllocSites[0].CallStack[3],
+              FrameContains("abc", 5U, 30U, false));
+
+  ASSERT_EQ(Bar.CallSites.size(), 1U);
+  ASSERT_EQ(Bar.CallSites[0].size(), 2U);
+  EXPECT_THAT(Bar.CallSites[0][0], FrameContains("foo", 5U, 30U, true));
+  EXPECT_THAT(Bar.CallSites[0][1], FrameContains("bar", 51U, 20U, false));
+
+  // Check the memprof record for xyz.
+  const llvm::GlobalValue::GUID XyzId = MemProfRecord::getGUID("xyz");
+  ASSERT_EQ(Records.count(XyzId), 1U);
+  const MemProfRecord &Xyz = Records[XyzId];
+  ASSERT_EQ(Xyz.CallSites.size(), 1U);
+  ASSERT_EQ(Xyz.CallSites[0].size(), 2U);
+  // Expect the entire frame even though in practice we only need the first
+  // entry here.
+  EXPECT_THAT(Xyz.CallSites[0][0], FrameContains("xyz", 5U, 30U, true));
+  EXPECT_THAT(Xyz.CallSites[0][1], FrameContains("abc", 5U, 30U, false));
+
+  // Check the memprof record for abc.
+  const llvm::GlobalValue::GUID AbcId = MemProfRecord::getGUID("abc");
+  ASSERT_EQ(Records.count(AbcId), 1U);
+  const MemProfRecord &Abc = Records[AbcId];
+  EXPECT_TRUE(Abc.AllocSites.empty());
+  ASSERT_EQ(Abc.CallSites.size(), 1U);
+  ASSERT_EQ(Abc.CallSites[0].size(), 2U);
+  EXPECT_THAT(Abc.CallSites[0][0], FrameContains("xyz", 5U, 30U, true));
+  EXPECT_THAT(Abc.CallSites[0][1], FrameContains("abc", 5U, 30U, false));
 }
 
 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