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

This patch adds support for optional memory profile information to be
included with and indexed profile. The indexed profile header adds a new
field which points to the offset of the memory profile section (if
present) in the indexed profile. For users who do not utilize this
feature the only overhead is a 64-bit offset in the header.

The memory profile section contains (1) profile metadata describing the
information recorded for each entry (2) an on-disk hashtable containing
the profile records indexed via llvm::md5(function_name). We chose to
introduce a separate hash table instead of the existing one since the
indexing for the instrumented fdo hash table is based on a CFG hash
which itself is perturbed by memprof instrumentation.

Differential Revision: https://reviews.llvm.org/D118653

Author: snehasish

compiler-rt/include/profile/InstrProfData.inc

@@ -650,7 +650,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
 /* Raw profile format version (start from 1). */
 #define INSTR_PROF_RAW_VERSION 8
 /* Indexed profile format version (start from 1). */
-#define INSTR_PROF_INDEX_VERSION 7
+#define INSTR_PROF_INDEX_VERSION 8
 /* Coverage mapping format version (start from 0). */
 #define INSTR_PROF_COVMAP_VERSION 5
 
@@ -662,6 +662,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
  * The 59th bit indicates whether to use debug info to correlate profiles.
  * The 60th bit indicates single byte coverage instrumentation.
  * The 61st bit indicates function entry instrumentation only.
+ * The 62nd bit indicates whether memory profile information is present.
  */
 #define VARIANT_MASKS_ALL 0xff00000000000000ULL
 #define GET_VERSION(V) ((V) & ~VARIANT_MASKS_ALL)
@@ -671,6 +672,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
 #define VARIANT_MASK_DBG_CORRELATE (0x1ULL << 59)
 #define VARIANT_MASK_BYTE_COVERAGE (0x1ULL << 60)
 #define VARIANT_MASK_FUNCTION_ENTRY_ONLY (0x1ULL << 61)
+#define VARIANT_MASK_MEMPROF (0x1ULL << 62)
 #define INSTR_PROF_RAW_VERSION_VAR __llvm_profile_raw_version
 #define INSTR_PROF_PROFILE_RUNTIME_VAR __llvm_profile_runtime
 #define INSTR_PROF_PROFILE_COUNTER_BIAS_VAR __llvm_profile_counter_bias

llvm/include/llvm/ProfileData/InstrProf.h

@@ -287,7 +287,8 @@ enum class InstrProfKind {
   CS = 0x8, // A context sensitive IR-level profile.
   SingleByteCoverage = 0x10, // Use single byte probes for coverage.
   FunctionEntryOnly = 0x20,  // Only instrument the function entry basic block.
-  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/FunctionEntryOnly)
+  MemProf = 0x40, // A memory profile collected using -fmemory-profile.
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/MemProf)
 };
 
 const std::error_category &instrprof_category();
@@ -1011,7 +1012,9 @@ enum ProfVersion {
   Version6 = 6,
   // An additional counter is added around logical operators.
   Version7 = 7,
-  // The current version is 7.
+  // An additional (optional) memory profile type is added.
+  Version8 = 8,
+  // The current version is 8.
   CurrentVersion = INSTR_PROF_INDEX_VERSION
 };
 const uint64_t Version = ProfVersion::CurrentVersion;
@@ -1028,6 +1031,7 @@ struct Header {
   uint64_t Unused; // Becomes unused since version 4
   uint64_t HashType;
   uint64_t HashOffset;
+  uint64_t MemProfOffset;
   // New fields should only be added at the end to ensure that the size
   // computation is correct. The methods below need to be updated to ensure that
   // the new field is read correctly.

llvm/include/llvm/ProfileData/InstrProfData.inc

@@ -650,7 +650,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
 /* Raw profile format version (start from 1). */
 #define INSTR_PROF_RAW_VERSION 8
 /* Indexed profile format version (start from 1). */
-#define INSTR_PROF_INDEX_VERSION 7
+#define INSTR_PROF_INDEX_VERSION 8
 /* Coverage mapping format version (start from 0). */
 #define INSTR_PROF_COVMAP_VERSION 5
 
@@ -662,6 +662,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
  * The 59th bit indicates whether to use debug info to correlate profiles.
  * The 60th bit indicates single byte coverage instrumentation.
  * The 61st bit indicates function entry instrumentation only.
+ * The 62nd bit indicates whether memory profile information is present.
  */
 #define VARIANT_MASKS_ALL 0xff00000000000000ULL
 #define GET_VERSION(V) ((V) & ~VARIANT_MASKS_ALL)
@@ -671,6 +672,7 @@ serializeValueProfDataFrom(ValueProfRecordClosure *Closure,
 #define VARIANT_MASK_DBG_CORRELATE (0x1ULL << 59)
 #define VARIANT_MASK_BYTE_COVERAGE (0x1ULL << 60)
 #define VARIANT_MASK_FUNCTION_ENTRY_ONLY (0x1ULL << 61)
+#define VARIANT_MASK_MEMPROF (0x1ULL << 62)
 #define INSTR_PROF_RAW_VERSION_VAR __llvm_profile_raw_version
 #define INSTR_PROF_PROFILE_RUNTIME_VAR __llvm_profile_runtime
 #define INSTR_PROF_PROFILE_COUNTER_BIAS_VAR __llvm_profile_counter_bias

llvm/include/llvm/ProfileData/InstrProfReader.h

@@ -19,6 +19,7 @@
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/InstrProfCorrelator.h"
+#include "llvm/ProfileData/MemProf.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/LineIterator.h"
@@ -471,6 +472,9 @@ struct InstrProfReaderIndexBase {
 using OnDiskHashTableImplV3 =
     OnDiskIterableChainedHashTable<InstrProfLookupTrait>;
 
+using MemProfHashTable =
+    OnDiskIterableChainedHashTable<memprof::MemProfRecordLookupTrait>;
+
 template <typename HashTableImpl>
 class InstrProfReaderItaniumRemapper;
 
@@ -556,6 +560,11 @@ class IndexedInstrProfReader : public InstrProfReader {
   std::unique_ptr<ProfileSummary> Summary;
   /// Context sensitive profile summary data.
   std::unique_ptr<ProfileSummary> CS_Summary;
+  /// MemProf profile schema (if available).
+  memprof::MemProfSchema Schema;
+  /// MemProf profile data on-disk indexed via llvm::md5(FunctionName).
+  std::unique_ptr<MemProfHashTable> MemProfTable;
+
   // Index to the current record in the record array.
   unsigned RecordIndex;
 
@@ -609,6 +618,11 @@ class IndexedInstrProfReader : public InstrProfReader {
   Expected<InstrProfRecord> getInstrProfRecord(StringRef FuncName,
                                                uint64_t FuncHash);
 
+  /// Return the memprof records for the function identified by
+  /// llvm::md5(Name).
+  Expected<ArrayRef<memprof::MemProfRecord>>
+  getMemProfRecord(uint64_t FuncNameHash);
+
   /// Fill Counts with the profile data for the given function name.
   Error getFunctionCounts(StringRef FuncName, uint64_t FuncHash,
                           std::vector<uint64_t> &Counts);

llvm/include/llvm/ProfileData/InstrProfWriter.h

@@ -17,6 +17,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ProfileData/InstrProf.h"
+#include "llvm/ProfileData/MemProf.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -37,6 +38,11 @@ class InstrProfWriter {
 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;
+
   // An enum describing the attributes of the profile.
   InstrProfKind ProfileKind = InstrProfKind::Unknown;
   // Use raw pointer here for the incomplete type object.
@@ -57,6 +63,9 @@ class InstrProfWriter {
     addRecord(std::move(I), 1, Warn);
   }
 
+  void addRecord(const ::llvm::memprof::MemProfRecord &MR,
+                 function_ref<void(Error)> Warn);
+
   /// Merge existing function counts from the given writer.
   void mergeRecordsFromWriter(InstrProfWriter &&IPW,
                               function_ref<void(Error)> Warn);
@@ -112,6 +121,8 @@ class InstrProfWriter {
     return Error::success();
   }
 
+  InstrProfKind getProfileKind() const { return ProfileKind; }
+
   // Internal interface for testing purpose only.
   void setValueProfDataEndianness(support::endianness Endianness);
   void setOutputSparse(bool Sparse);

llvm/include/llvm/ProfileData/MemProf.h

@@ -5,6 +5,7 @@
 #include <string>
 #include <vector>
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ProfileData/MemProfData.inc"
 #include "llvm/ProfileData/ProfileCommon.h"
@@ -134,25 +135,65 @@ struct PortableMemInfoBlock {
 };
 
 struct MemProfRecord {
-  struct Frame {
-    std::string Function;
+  // 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.
+  PACKED(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;
+    // The source line offset of the call from the beginning of parent function.
     uint32_t LineOffset;
+    // The source column number of the call to help distinguish multiple calls
+    // on the same line.
     uint32_t Column;
+    // Whether the current frame is inlined.
     bool IsInlineFrame;
 
-    Frame(std::string Str, uint32_t Off, uint32_t Col, bool Inline)
-        : Function(std::move(Str)), LineOffset(Off), Column(Col),
-          IsInlineFrame(Inline) {}
-  };
+    Frame(uint64_t Hash, uint32_t Off, uint32_t Col, bool Inline)
+        : Function(Hash), LineOffset(Off), Column(Col), IsInlineFrame(Inline) {}
 
+    bool operator==(const Frame &Other) const {
+      return Other.Function == Function && Other.LineOffset == LineOffset &&
+             Other.Column == Column && Other.IsInlineFrame == IsInlineFrame;
+    }
+
+    bool operator!=(const Frame &Other) const { return !operator==(Other); }
+
+    // Write the contents of the frame to the ostream \p OS.
+    void write(raw_ostream & OS) const {
+      using namespace support;
+
+      endian::Writer LE(OS, little);
+
+      // If the type of the GlobalValue::GUID changes, then we need to update
+      // the reader and the writer.
+      static_assert(std::is_same<GlobalValue::GUID, uint64_t>::value,
+                    "Expect GUID to be uint64_t.");
+      LE.write<uint64_t>(Function);
+
+      LE.write<uint32_t>(LineOffset);
+      LE.write<uint32_t>(Column);
+      LE.write<bool>(IsInlineFrame);
+    }
+  });
+
+  // The dynamic calling context for the allocation.
   std::vector<Frame> CallStack;
+  // The statistics obtained from the runtime for the allocation.
   PortableMemInfoBlock Info;
 
   void clear() {
     CallStack.clear();
     Info.clear();
   }
 
+  size_t serializedSize() const {
+    return sizeof(uint64_t) + // The number of frames to serialize.
+           sizeof(Frame) * CallStack.size() + // The contents of the frames.
+           PortableMemInfoBlock::serializedSize(); // The size of the payload.
+  }
+
   // Prints out the contents of the memprof record in YAML.
   void print(llvm::raw_ostream &OS) const {
     OS << "    Callstack:\n";
@@ -168,6 +209,138 @@ struct MemProfRecord {
 
     Info.printYAML(OS);
   }
+
+  bool operator==(const MemProfRecord &Other) const {
+    if (Other.Info != Info)
+      return false;
+
+    if (Other.CallStack.size() != CallStack.size())
+      return false;
+
+    for (size_t I = 0; I < Other.CallStack.size(); I++) {
+      if (Other.CallStack[I] != CallStack[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);
+
+// Deserializes memprof records from the Buffer
+SmallVector<MemProfRecord, 4> deserializeRecords(const MemProfSchema &Schema,
+                                                 const unsigned char *Buffer);
+
+// 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 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) {}
+
+  static bool EqualKey(uint64_t A, uint64_t B) { return A == B; }
+  static uint64_t GetInternalKey(uint64_t K) { return K; }
+  static uint64_t GetExternalKey(uint64_t K) { return K; }
+
+  hash_value_type ComputeHash(uint64_t K) { return K; }
+
+  static std::pair<offset_type, offset_type>
+  ReadKeyDataLength(const unsigned char *&D) {
+    using namespace support;
+
+    offset_type KeyLen = endian::readNext<offset_type, little, unaligned>(D);
+    offset_type DataLen = endian::readNext<offset_type, little, unaligned>(D);
+    return std::make_pair(KeyLen, DataLen);
+  }
+
+  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;
+  }
+
+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;
+};
+
+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 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();
+    }
+
+    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);
+  }
 };
 
 } // namespace memprof

llvm/include/llvm/ProfileData/MemProfData.inc

@@ -1,5 +1,5 @@
-#ifndef LLVM_PROFILEDATA_MEMPROFDATA_INC
-#define LLVM_PROFILEDATA_MEMPROFDATA_INC
+#ifndef MEMPROF_DATA_INC
+#define MEMPROF_DATA_INC
 /*===-- MemProfData.inc - MemProf profiling runtime structures -*- C++ -*-=== *\
 |*
 |* Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.

llvm/include/llvm/ProfileData/RawMemProfReader.h

@@ -66,6 +66,9 @@ class RawMemProfReader {
     return Iterator(this);
   }
 
+  // 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,