Reland "[memprof] Extend the index prof format to include memory profiles."

This reverts commit 0f73fb18ca333e38cdb9ffa701a8db026c56041d.

Use llvm/Profile/MIBEntryDef.inc instead of relative path.

Generated the raw profile data with `-mllvm
-enable-name-compression=false` so that builbots where the reader is
built without zlib do not fail.

Also updated the test build instructions.

6 files changed, 226 insertions, 11 deletions

diff --git a/llvm/lib/ProfileData/CMakeLists.txt b/llvm/lib/ProfileData/CMakeLists.txt
index 2749119f72d9..486c45d0dff5 100644
--- a/llvm/lib/ProfileData/CMakeLists.txt
+++ b/llvm/lib/ProfileData/CMakeLists.txt
@@ -4,6 +4,7 @@ add_llvm_component_library(LLVMProfileData
   InstrProfCorrelator.cpp
   InstrProfReader.cpp
   InstrProfWriter.cpp
+  MemProf.cpp
   ProfileSummaryBuilder.cpp
   SampleProf.cpp
   SampleProfReader.cpp
diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 6bdca486205e..924125eec8e1 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -1345,8 +1345,15 @@ Expected<Header> Header::readFromBuffer(const unsigned char *Buffer) {
     return make_error<InstrProfError>(instrprof_error::unsupported_version);
 
   switch (GET_VERSION(FormatVersion)) {
-  // When a new field is added in the header add a case statement here to
-  // populate it.
+    // When a new field is added in the header add a case statement here to
+    // populate it.
+    static_assert(
+        IndexedInstrProf::ProfVersion::CurrentVersion == Version8,
+        "Please update the reading code below if a new field has been added, "
+        "if not add a case statement to fall through to the latest version.");
+  case 8ull:
+    H.MemProfOffset = read(Buffer, offsetOf(&Header::MemProfOffset));
+    LLVM_FALLTHROUGH;
   default:
     H.HashType = read(Buffer, offsetOf(&Header::HashType));
     H.HashOffset = read(Buffer, offsetOf(&Header::HashOffset));
@@ -1357,9 +1364,15 @@ Expected<Header> Header::readFromBuffer(const unsigned char *Buffer) {
 
 size_t Header::size() const {
   switch (GET_VERSION(Version)) {
-  // When a new field is added to the header add a case statement here to
-  // compute the size as offset of the new field + size of the new field. This
-  // relies on the field being added to the end of the list.
+    // When a new field is added to the header add a case statement here to
+    // compute the size as offset of the new field + size of the new field. This
+    // relies on the field being added to the end of the list.
+    static_assert(IndexedInstrProf::ProfVersion::CurrentVersion == Version8,
+                  "Please update the size computation below if a new field has "
+                  "been added to the header, if not add a case statement to "
+                  "fall through to the latest version.");
+  case 8ull:
+    return offsetOf(&Header::MemProfOffset) + sizeof(Header::MemProfOffset);
   default:
     return offsetOf(&Header::HashOffset) + sizeof(Header::HashOffset);
   }
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index c45b8cde2cda..39ba004ee879 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -19,7 +19,9 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/InstrProf.h"
+#include "llvm/ProfileData/MemProf.h"
 #include "llvm/ProfileData/ProfileCommon.h"
+#include "llvm/ProfileData/RawMemProfReader.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
@@ -57,6 +59,9 @@ static InstrProfKind getProfileKindFromVersion(uint64_t Version) {
   if (Version & VARIANT_MASK_FUNCTION_ENTRY_ONLY) {
     ProfileKind |= InstrProfKind::FunctionEntryOnly;
   }
+  if (Version & VARIANT_MASK_MEMPROF) {
+    ProfileKind |= InstrProfKind::MemProf;
+  }
   return ProfileKind;
 }
 
@@ -955,10 +960,35 @@ Error IndexedInstrProfReader::readHeader() {
 
   uint64_t HashOffset = endian::byte_swap<uint64_t, little>(Header->HashOffset);
 
-  // The rest of the file is an on disk hash table.
+  // The hash table with profile counts comes next.
   auto IndexPtr = std::make_unique<InstrProfReaderIndex<OnDiskHashTableImplV3>>(
       Start + HashOffset, Cur, Start, HashType, Header->Version);
 
+  // The MemProfOffset field in the header is only valid when the format version
+  // is higher than 8 (when it was introduced).
+  if (GET_VERSION(Header->Version) >= 8 &&
+      Header->Version & VARIANT_MASK_MEMPROF) {
+    uint64_t MemProfOffset =
+        endian::byte_swap<uint64_t, little>(Header->MemProfOffset);
+
+    const unsigned char *Ptr = Start + MemProfOffset;
+    // The value returned from Generator.Emit.
+    const uint64_t TableOffset =
+        support::endian::readNext<uint64_t, little, unaligned>(Ptr);
+
+    // Read the schema.
+    auto SchemaOr = memprof::readMemProfSchema(Ptr);
+    if (!SchemaOr)
+      return SchemaOr.takeError();
+    Schema = SchemaOr.get();
+
+    // Now initialize the table reader with a pointer into data buffer.
+    MemProfTable.reset(MemProfHashTable::Create(
+        /*Buckets=*/Start + TableOffset,
+        /*Payload=*/Ptr,
+        /*Base=*/Start, memprof::MemProfRecordLookupTrait(Schema)));
+  }
+
   // Load the remapping table now if requested.
   if (RemappingBuffer) {
     Remapper = std::make_unique<
@@ -1003,6 +1033,17 @@ IndexedInstrProfReader::getInstrProfRecord(StringRef FuncName,
   return error(instrprof_error::hash_mismatch);
 }
 
+Expected<ArrayRef<memprof::MemProfRecord>>
+IndexedInstrProfReader::getMemProfRecord(uint64_t FuncNameHash) {
+  auto Iter = MemProfTable->find(FuncNameHash);
+  if (Iter == MemProfTable->end())
+    // TODO: Add memprof specific errors.
+    return make_error<InstrProfError>(instrprof_error::hash_mismatch,
+                                      "memprof record not found for hash " +
+                                          Twine(FuncNameHash));
+  return *Iter;
+}
+
 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
                                                 uint64_t FuncHash,
                                                 std::vector<uint64_t> &Counts) {
diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index ebf89317d585..4c974f402d2b 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/InstrProf.h"
+#include "llvm/ProfileData/MemProf.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
@@ -63,11 +64,16 @@ public:
 
     if (IsFDOStream) {
       raw_fd_ostream &FDOStream = static_cast<raw_fd_ostream &>(OS);
+      const uint64_t LastPos = FDOStream.tell();
       for (int K = 0; K < NItems; K++) {
         FDOStream.seek(P[K].Pos);
         for (int I = 0; I < P[K].N; I++)
           write(P[K].D[I]);
       }
+      // Reset the stream to the last position after patching so that users
+      // don't accidentally overwrite data. This makes it consistent with
+      // the string stream below which replaces the data directly.
+      FDOStream.seek(LastPos);
     } else {
       raw_string_ostream &SOStream = static_cast<raw_string_ostream &>(OS);
       std::string &Data = SOStream.str(); // with flush
@@ -248,11 +254,39 @@ void InstrProfWriter::addRecord(StringRef Name, uint64_t Hash,
   Dest.sortValueData();
 }
 
+void InstrProfWriter::addRecord(const memprof::MemProfRecord &MR,
+                                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;
+    }
+  }
+
+  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);
 }
 
 bool InstrProfWriter::shouldEncodeData(const ProfilingData &PD) {
@@ -297,6 +331,7 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   for (const auto &I : FunctionData)
     if (shouldEncodeData(I.getValue()))
       Generator.insert(I.getKey(), &I.getValue());
+
   // Write the header.
   IndexedInstrProf::Header Header;
   Header.Magic = IndexedInstrProf::Magic;
@@ -311,16 +346,18 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
     Header.Version |= VARIANT_MASK_BYTE_COVERAGE;
   if (static_cast<bool>(ProfileKind & InstrProfKind::FunctionEntryOnly))
     Header.Version |= VARIANT_MASK_FUNCTION_ENTRY_ONLY;
+  if (static_cast<bool>(ProfileKind & InstrProfKind::MemProf))
+    Header.Version |= VARIANT_MASK_MEMPROF;
 
   Header.Unused = 0;
   Header.HashType = static_cast<uint64_t>(IndexedInstrProf::HashType);
   Header.HashOffset = 0;
+  Header.MemProfOffset = 0;
   int N = sizeof(IndexedInstrProf::Header) / sizeof(uint64_t);
 
-  // Only write out all the fields except 'HashOffset'. We need
-  // to remember the offset of that field to allow back patching
-  // later.
-  for (int I = 0; I < N - 1; I++)
+  // Only write out all the fields except 'HashOffset' and 'MemProfOffset'. We
+  // need to remember the offset of these fields to allow back patching later.
+  for (int I = 0; I < N - 2; I++)
     OS.write(reinterpret_cast<uint64_t *>(&Header)[I]);
 
   // Save the location of Header.HashOffset field in \c OS.
@@ -328,6 +365,13 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   // Reserve the space for HashOffset field.
   OS.write(0);
 
+  // Save the location of MemProf profile data. This is stored in two parts as
+  // the schema and as a separate on-disk chained hashtable.
+  uint64_t MemProfSectionOffset = OS.tell();
+  // Reserve space for the MemProf table field to be patched later if this
+  // profile contains memory profile information.
+  OS.write(0);
+
   // Reserve space to write profile summary data.
   uint32_t NumEntries = ProfileSummaryBuilder::DefaultCutoffs.size();
   uint32_t SummarySize = Summary::getSize(Summary::NumKinds, NumEntries);
@@ -347,6 +391,42 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   // Write the hash table.
   uint64_t HashTableStart = Generator.Emit(OS.OS, *InfoObj);
 
+  // Write the MemProf profile data if we have it. This includes a simple schema
+  // with the format described below followed by the hashtable:
+  // uint64_t Offset = MemProfGenerator.Emit
+  // uint64_t Num schema entries
+  // uint64_t Schema entry 0
+  // uint64_t Schema entry 1
+  // ....
+  // uint64_t Schema entry N - 1
+  // OnDiskChainedHashTable MemProfFunctionData
+  uint64_t MemProfSectionStart = 0;
+  if (static_cast<bool>(ProfileKind & InstrProfKind::MemProf)) {
+    MemProfSectionStart = OS.tell();
+    OS.write(0ULL); // Reserve space for the offset.
+
+    auto Schema = memprof::PortableMemInfoBlock::getSchema();
+    OS.write(static_cast<uint64_t>(Schema.size()));
+    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).
+      MemProfGenerator.insert(I.first, I.second);
+    }
+
+    uint64_t TableOffset = MemProfGenerator.Emit(OS.OS, *MemProfWriter);
+    PatchItem PatchItems[] = {
+        {MemProfSectionStart, &TableOffset, 1},
+    };
+    OS.patch(PatchItems, 1);
+  }
+
   // Allocate space for data to be serialized out.
   std::unique_ptr<IndexedInstrProf::Summary> TheSummary =
       IndexedInstrProf::allocSummary(SummarySize);
@@ -369,6 +449,8 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   PatchItem PatchItems[] = {
       // Patch the Header.HashOffset field.
       {HashTableStartFieldOffset, &HashTableStart, 1},
+      // Patch the Header.MemProfOffset (=0 for profiles without MemProf data).
+      {MemProfSectionOffset, &MemProfSectionStart, 1},
       // Patch the summary data.
       {SummaryOffset, reinterpret_cast<uint64_t *>(TheSummary.get()),
        (int)(SummarySize / sizeof(uint64_t))},
diff --git a/llvm/lib/ProfileData/MemProf.cpp b/llvm/lib/ProfileData/MemProf.cpp
new file mode 100644
index 000000000000..6a9b69ff6cff
--- /dev/null
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -0,0 +1,73 @@
+#include "llvm/ProfileData/MemProf.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) {
+  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) {
+      F.write(OS);
+    }
+    MR.Info.serialize(Schema, OS);
+  }
+}
+
+SmallVector<MemProfRecord, 4> deserializeRecords(const MemProfSchema &Schema,
+                                                 const unsigned char *Ptr) {
+  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;
+    const uint64_t NumFrames =
+        endian::readNext<uint64_t, little, unaligned>(Ptr);
+    for (uint64_t J = 0; J < NumFrames; J++) {
+      const auto F = *reinterpret_cast<const MemProfRecord::Frame *>(Ptr);
+      Ptr += sizeof(MemProfRecord::Frame);
+      MR.CallStack.push_back(F);
+    }
+    MR.Info.deserialize(Schema, Ptr);
+    Ptr += PortableMemInfoBlock::serializedSize();
+    Records.push_back(MR);
+  }
+  return Records;
+}
+
+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);
+  if (NumSchemaIds > static_cast<uint64_t>(Meta::Size)) {
+    return make_error<InstrProfError>(instrprof_error::malformed,
+                                      "memprof schema invalid");
+  }
+
+  MemProfSchema Result;
+  for (size_t I = 0; I < NumSchemaIds; I++) {
+    const uint64_t Tag = endian::readNext<uint64_t, little, unaligned>(Ptr);
+    if (Tag >= static_cast<uint64_t>(Meta::Size)) {
+      return make_error<InstrProfError>(instrprof_error::malformed,
+                                        "memprof schema invalid");
+    }
+    Result.push_back(static_cast<Meta>(Tag));
+  }
+  // Advace the buffer to one past the schema if we succeeded.
+  Buffer = Ptr;
+  return Result;
+}
+
+} // namespace memprof
+} // namespace llvm
diff --git a/llvm/lib/ProfileData/RawMemProfReader.cpp b/llvm/lib/ProfileData/RawMemProfReader.cpp
index 43ef7c947366..9bcba2a2b04e 100644
--- a/llvm/lib/ProfileData/RawMemProfReader.cpp
+++ b/llvm/lib/ProfileData/RawMemProfReader.cpp
@@ -362,7 +362,12 @@ Error RawMemProfReader::fillRecord(const uint64_t Id, const MemInfoBlock &MIB,
     for (size_t I = 0; I < DI.getNumberOfFrames(); I++) {
       const auto &Frame = DI.getFrame(I);
       Record.CallStack.emplace_back(
-          std::to_string(llvm::MD5Hash(trimSuffix(Frame.FunctionName))),
+          // 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)),
           Frame.Line - Frame.StartLine, Frame.Column,
           // Only the first entry is not an inlined location.
           I != 0);