Reland "[memprof] Store callsite metadata with memprof records."

This reverts commit f4b7944.

Reland with underlying msan issue fixed in D122260.

Author: Snehasish Kumar

llvm/include/llvm/ProfileData/InstrProfWriter.h

@@ -15,7 +15,9 @@
 #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"
@@ -41,7 +43,7 @@ class InstrProfWriter {
 
   // 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;
@@ -63,7 +65,8 @@ class InstrProfWriter {
     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.

llvm/include/llvm/ProfileData/MemProf.h

@@ -82,9 +82,9 @@ struct PortableMemInfoBlock {
 
   // 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
   }
@@ -133,6 +133,7 @@ struct PortableMemInfoBlock {
 #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
@@ -193,64 +194,152 @@ struct MemProfRecord {
       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";
+    }
   };
 
-  // The dynamic calling context for the allocation.
-  llvm::SmallVector<Frame> CallStack;
-  // The statistics obtained from the runtime for the allocation.
-  PortableMemInfoBlock Info;
+  struct AllocationInfo {
+    // The dynamic calling context for the allocation.
+    llvm::SmallVector<Frame> CallStack;
+    // The statistics obtained from the runtime for the allocation.
+    PortableMemInfoBlock Info;
+
+    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() {
-    CallStack.clear();
-    Info.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 {
-    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.
+    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);
+  // Serializes the memprof records in \p Records to the ostream \p OS based on
+  // the schema provided in \p Schema.
+  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,
+                                   const unsigned char *Buffer);
+
+  // 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
@@ -259,14 +348,11 @@ SmallVector<MemProfRecord, 4> deserializeRecords(const MemProfSchema &Schema,
 // 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;
@@ -297,24 +383,24 @@ class MemProfRecordLookupTrait {
 
   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;
@@ -333,17 +419,9 @@ class MemProfRecordWriterTrait {
     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);
   }
@@ -357,7 +435,7 @@ class MemProfRecordWriterTrait {
   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);
   }
 };
 

llvm/include/llvm/ProfileData/RawMemProfReader.h

@@ -14,9 +14,11 @@
 
 #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"
@@ -57,15 +59,16 @@ class RawMemProfReader {
   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; }
 
@@ -75,14 +78,16 @@ class RawMemProfReader {
                    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:
@@ -96,10 +101,12 @@ class RawMemProfReader {
   // 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;
@@ -112,15 +119,15 @@ class RawMemProfReader {
   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