[memprof] Symbolize and cache stack frames.

[snehasish]

Summary

Currently, symbolization of stack frames occurs on demand when the instrprof writer
iterates over all the records in the raw memprof reader. With this
change we symbolize and cache the frames immediately after reading the
raw profiles. For a large internal binary this results in a runtime
reduction of ~50% (2m -> 48s) when merging a memprof raw profile with a
raw instr profile to generate an indexed profile. This change also makes
it simpler in the future to generate additional calling context
metadata to attach to each memprof record.

[davidxl]

llvm/lib/ProfileData/RawMemProfReader.cpp

354

Using a map from addr to Frames for caching can also avoid redundant symbolization computation. Is there an advantage of doing eager symbolization?

[tejohnson]

What is the memory impact of caching? Hopefully not too onerous since this is a big speedup!

llvm/include/llvm/ProfileData/RawMemProfReader.h

81

I take it initialize() not get called in this case? Maybe add a comment.

llvm/lib/ProfileData/RawMemProfReader.cpp

354

I suspect because it makes it easier for adding additional records, such as identifying interior call stack nodes within their functions (we need to mark these with metadata as well).

[snehasish]

Address comment.

[snehasish]

Updated the patch with comments. PTAL, thanks!

In D120430#3355934, @tejohnson wrote:

What is the memory impact of caching? Hopefully not too onerous since this is a big speedup!

There is hardly any change in peak memory consumption, valgrind massif shows 8.618G (now) vs 8.615G (before). In this patch we have introduced a layer of indirection to where the indexing (keys in map of PC->Frame) account for the small increase in memory.

llvm/lib/ProfileData/RawMemProfReader.cpp

354

Caching the frames outside the local context allows us to decouple the record generation from the symbolization and makes it easier to extend for pruning (D120860) and marking interior callstack nodes (as Teresa noted). Note that the symbolization here isn't eager in the sense that more work is performed, the number of unique addresses symbolized remains the same before and after this patch.

[tejohnson]

lgtm

[davidxl]

lgtm

llvm/include/llvm/ProfileData/MemProf.h

struct Frame {
     // 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);
     }
   };
 
   // The dynamic calling context for the allocation.
-  std::vector<Frame> CallStack;
+  llvm::SmallVector<Frame> CallStack;
   // The statistics obtained from the runtime for the allocation.
   PortableMemInfoBlock Info;
 
   void clear() {
     CallStack.clear();
     Info.clear();
   }
 

llvm/include/llvm/ProfileData/RawMemProfReader.h

 
 #include <cstddef>
 
 namespace llvm {
 namespace memprof {
 
 // Map from id (recorded from sanitizer stack depot) to virtual addresses for
 // each program counter address in the callstack.
-using CallStackMap = llvm::DenseMap<uint64_t, llvm::SmallVector<uint64_t, 32>>;
+using CallStackMap = llvm::DenseMap<uint64_t, llvm::SmallVector<uint64_t>>;
 
 class RawMemProfReader {
 public:
   RawMemProfReader(std::unique_ptr<MemoryBuffer> DataBuffer)
       : DataBuffer(std::move(DataBuffer)) {}
   RawMemProfReader(const RawMemProfReader &) = delete;
   RawMemProfReader &operator=(const RawMemProfReader &) = delete;
 
public:
   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) {}
+        ProfileData(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 = symbolizeStackFrames())
+      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();
+  Error symbolizeStackFrames();
 
   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;
   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();
 };
 
 } // namespace memprof
 } // namespace llvm
 
 #endif // LLVM_PROFILEDATA_RAWMEMPROFREADER_H_

llvm/lib/ProfileData/RawMemProfReader.cpp

CallStackMap readStackInfo(const char *Ptr) {
   const uint64_t NumItemsToRead =
       endian::readNext<uint64_t, little, unaligned>(Ptr);
   CallStackMap Items;
 
   for (uint64_t I = 0; I < NumItemsToRead; I++) {
     const uint64_t StackId = endian::readNext<uint64_t, little, unaligned>(Ptr);
     const uint64_t NumPCs = endian::readNext<uint64_t, little, unaligned>(Ptr);
 
-    SmallVector<uint64_t, 32> CallStack;
+    SmallVector<uint64_t> CallStack;
     for (uint64_t J = 0; J < NumPCs; J++) {
       CallStack.push_back(endian::readNext<uint64_t, little, unaligned>(Ptr));
     }
 
     Items[StackId] = CallStack;
   }
   return Items;
 }
std::unique_ptr<DIContext> Context = DWARFContext::create(
       *Object, DWARFContext::ProcessDebugRelocations::Process);
 
   auto SOFOr = symbolize::SymbolizableObjectFile::create(
       Object, std::move(Context), /*UntagAddresses=*/false);
   if (!SOFOr)
     return report(SOFOr.takeError(), FileName);
   Symbolizer = std::move(SOFOr.get());
 
-  return readRawProfile();
+  if (Error E = readRawProfile())
+    return E;
+
+  return symbolizeStackFrames();
+}
+
+Error RawMemProfReader::symbolizeStackFrames() {
+  // The specifier to use when symbolization is requested.
+  const DILineInfoSpecifier Specifier(
+      DILineInfoSpecifier::FileLineInfoKind::RawValue,
+      DILineInfoSpecifier::FunctionNameKind::LinkageName);
+
+  for (const auto &Entry : StackMap) {
+    for (const uint64_t VAddr : Entry.getSecond()) {
+      // Check if we have already symbolized and cached the result.
+      if (SymbolizedFrame.count(VAddr) > 0)
+        continue;
+
+      Expected<DIInliningInfo> DIOr = Symbolizer->symbolizeInlinedCode(
+          getModuleOffset(VAddr), Specifier, /*UseSymbolTable=*/false);
+      if (!DIOr)
+        return DIOr.takeError();
+      DIInliningInfo DI = DIOr.get();
+
+      for (size_t I = 0; I < DI.getNumberOfFrames(); I++) {
+        const auto &Frame = DI.getFrame(I);
+        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)),
+            Frame.Line - Frame.StartLine, Frame.Column,
+            // Only the first entry is not an inlined location.
+            I != 0);
+      }
+    }
+  }
+  return Error::success();
 }
 
 Error RawMemProfReader::readRawProfile() {
   const char *Next = DataBuffer->getBufferStart();
 
   while (Next < DataBuffer->getBufferEnd()) {
     auto *Header = reinterpret_cast<const memprof::Header *>(Next);
 
RawMemProfReader::getModuleOffset(const uint64_t VirtualAddress) {
   // 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];
-  DILineInfoSpecifier Specifier(
-      DILineInfoSpecifier::FileLineInfoKind::RawValue,
-      DILineInfoSpecifier::FunctionNameKind::LinkageName);
   for (const uint64_t Address : CallStack) {
-    Expected<DIInliningInfo> DIOr = Symbolizer->symbolizeInlinedCode(
-        getModuleOffset(Address), Specifier, /*UseSymbolTable=*/false);
-
-    if (!DIOr)
-      return DIOr.takeError();
-    DIInliningInfo DI = DIOr.get();
-
-    for (size_t I = 0; I < DI.getNumberOfFrames(); I++) {
-      const auto &Frame = DI.getFrame(I);
-      Record.CallStack.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)),
-          Frame.Line - Frame.StartLine, Frame.Column,
-          // Only the first entry is not an inlined location.
-          I != 0);
-    }
+    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())
     return make_error<InstrProfError>(instrprof_error::empty_raw_profile);

llvm/unittests/ProfileData/MemProfTest.cpp

#undef MIBEntryDef
   return Schema;
 }
 
 TEST(MemProf, FillsValue) {
   std::unique_ptr<MockSymbolizer> Symbolizer(new MockSymbolizer());
 
   EXPECT_CALL(*Symbolizer, symbolizeInlinedCode(SectionedAddress{0x2000},
                                                 specifier(), false))
-      .Times(2)
+      .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},
                                                 specifier(), false))
       .Times(1)