[memprof] Support symbolization of PIE binaries.

[snehasish]

Summary

Support symolization of PIE binaries in memprof. We assume that the
profiled binary has one executable text segment for simplicity. Update
the memprof-pic test to now expect the same output as the memprof-basic test.

[tejohnson]

llvm/lib/ProfileData/RawMemProfReader.cpp

283

Is this limitation going to affect us?

291

Is it expected that p_vaddr will not be page boundary aligned already?

297

I think it might be clearer to report this error in the above loop. Otherwise when reading that code it is a little confusing why PreferredTextSegmentAddress will be overwritten for each executable segment. Similar suggestion when setting ProfiledTextSegment{Start,End} in setupForSymbolization

581

I'm a little confused on the difference between PreferredTextSegmentAddress and ProfiledTextSegmentStart, and the arithmetic being done here. Can you clarify?

[snehasish]

Address comments.

[snehasish]

PTAL, thanks!

llvm/lib/ProfileData/RawMemProfReader.cpp

283

No, the systems we care about default to 4K. Notable exceptions include newer versions of Darwin which use 16K. See a similar assumption made in llvm-profgen [1].

https://github.com/llvm/llvm-project/blob/main/llvm/tools/llvm-profgen/ProfiledBinary.cpp#L316

291

It most likely will be page (or 2M hugepage aligned) already. The default is to align by 4K [1] and this should almost always be a no-op since even 2M aligned is also 4K aligned.

[1] https://github.com/llvm/llvm-project/blob/main/lld/ELF/Target.h#L112

581

For non-PIE binaries:
PreferredTextSegmentAddress is non-zero, e.g 0x200000 or 0x400000. In this case the loader will put the segment at the virtual address requested by the binary and the ProfiledTextSegmentStart will be the same as PreferredTextSegmentAddress. Thus this arithmetic is a no-op.

For PIE binaries:
PreferredTextSegmentAddress is zero. The loader is free to place the segment at any address (respecting alignment specified by the header). Then the instruction offset in the binary is computed as the virtual address minus the start address of the ProfiledTextSegment. This assumes p_offset = 0 for the executable segment (added an assertion).

Added comments to the code to clarify.

[tejohnson]

lgtm. Couple suggestions for asserts if you think they make sense.

llvm/lib/ProfileData/RawMemProfReader.cpp

291

Consider making this an assert that it is already page aligned?

581

Thanks, the new comment helps. Perhaps add an assert somewhere like:
assert(PreferredTextSegmentAddress == 0 || (PreferredTextSegmentAddress == ProfiledTextSegmentStart));

[snehasish]

Address comments.

[snehasish]

Thanks for the review!

[snehasish]

Rebase, remove var only used in assert.

llvm/include/llvm/ProfileData/RawMemProfReader.h

 
 private:
   RawMemProfReader(object::OwningBinary<object::Binary> &&Bin, bool KeepName)
       : Binary(std::move(Bin)), KeepSymbolName(KeepName) {}
   // Initializes the RawMemProfReader with the contents in `DataBuffer`.
   Error initialize(std::unique_ptr<MemoryBuffer> DataBuffer);
   // Read and parse the contents of the `DataBuffer` as a binary format profile.
   Error readRawProfile(std::unique_ptr<MemoryBuffer> DataBuffer);
+  // Initialize the segment mapping information for symbolization.
+  Error setupForSymbolization();
   // 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();
 
   // A helper method to extract the frame from the IdToFrame map.
   const Frame &idToFrame(const FrameId Id) const {
     auto It = IdToFrame.find(Id);
     assert(It != IdToFrame.end() && "Id not found in map.");
     return It->getSecond();
   }
 
   object::SectionedAddress getModuleOffset(uint64_t VirtualAddress);
 
+  // The profiled binary.
   object::OwningBinary<object::Binary> Binary;
+  // A symbolizer to translate virtual addresses to code locations.
   std::unique_ptr<llvm::symbolize::SymbolizableModule> Symbolizer;
+  // The preferred load address of the executable segment.
+  uint64_t PreferredTextSegmentAddress = 0;
+  // The base address of the text segment in the process during profiling.
+  uint64_t ProfiledTextSegmentStart = 0;
+  // The limit address of the text segment in the process during profiling.
+  uint64_t ProfiledTextSegmentEnd = 0;
+
+  // The memory mapped segment information for all executable segments in the
+  // profiled binary (filtered from the raw profile using the build id).
+  llvm::SmallVector<SegmentEntry, 2> SegmentInfo;
 
-  // 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> CallstackProfileData;
   CallStackMap StackMap;
 
   // Cached symbolization from PC to Frame.
   llvm::DenseMap<uint64_t, llvm::SmallVector<FrameId>> SymbolizedFrame;
   llvm::DenseMap<FrameId, Frame> IdToFrame;

llvm/lib/ProfileData/RawMemProfReader.cpp

 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableModule.h"
 #include "llvm/DebugInfo/Symbolize/SymbolizableObjectFile.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/BuildID.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/Debug.h"
 #include "llvm/Support/Endian.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/Path.h"
 
 #define DEBUG_TYPE "memprof"
 
 namespace llvm {
 namespace memprof {
 namespace {
 template <class T = uint64_t> inline T alignedRead(const char *Ptr) {
Error RawMemProfReader::initialize(std::unique_ptr<MemoryBuffer> DataBuffer) {
   }
 
   // Check whether the profiled binary was built with position independent code
   // (PIC). For now we provide a error message until symbolization support
   // is added for pic.
   auto* Elf64LEObject = llvm::cast<llvm::object::ELF64LEObjectFile>(ElfObject);
   const llvm::object::ELF64LEFile& ElfFile = Elf64LEObject->getELFFile();
   auto PHdrsOr = ElfFile.program_headers();
-  if(!PHdrsOr) 
-    return report(make_error<StringError>(Twine("Could not read program headers: "),
-                                          inconvertibleErrorCode()),
-                  FileName);
-  auto FirstLoadHeader = PHdrsOr->begin();
-  while (FirstLoadHeader->p_type != llvm::ELF::PT_LOAD)
-    ++FirstLoadHeader;
-  if(FirstLoadHeader->p_vaddr == 0)
-    return report(make_error<StringError>(Twine("Unsupported position independent code"),
-                                          inconvertibleErrorCode()),
-                  FileName);
+  if (!PHdrsOr)
+    return report(
+        make_error<StringError>(Twine("Could not read program headers: "),
+
+  int NumExecutableSegments = 0;
+  for (const auto &Phdr : *PHdrsOr) {
+    if (Phdr.p_type == ELF::PT_LOAD) {
+      if (Phdr.p_flags & ELF::PF_X) {
+        // We assume only one text segment in the main binary for simplicity and
+        // reduce the overhead of checking multiple ranges during symbolization.
+        if (++NumExecutableSegments > 1) {
+          return report(
+              make_error<StringError>(
+                  "Expect only one executable load segment in the binary",
+        }
+        // Segment will always be loaded at a page boundary, expect it to be
+        // aligned already. Assume 4K pagesize for the machine from which the
+        // profile has been collected. This should be fine for now, in case we
+        // want to support other pagesizes it can be recorded in the raw profile
+        // during collection.
+        PreferredTextSegmentAddress = Phdr.p_vaddr;
+        assert(Phdr.p_vaddr == (Phdr.p_vaddr & ~(0x1000 - 1U)) &&
+               "Expect p_vaddr to always be page aligned");
+        assert(Phdr.p_offset == 0 && "Expect p_offset = 0 for symbolization.");
+      }
+    }
+  }
 
   auto Triple = ElfObject->makeTriple();
   if (!Triple.isX86())
     return report(make_error<StringError>(Twine("Unsupported target: ") +
                                               Triple.getArchName(),
                                           inconvertibleErrorCode()),
                   FileName);
 
   auto *Object = cast<object::ObjectFile>(Binary.getBinary());
   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());
 
+  // Process the raw profile.
   if (Error E = readRawProfile(std::move(DataBuffer)))
     return E;
 
+  if (Error E = setupForSymbolization())
+    return E;
+
   if (Error E = symbolizeAndFilterStackFrames())
     return E;
 
   return mapRawProfileToRecords();
 }
 
+Error RawMemProfReader::setupForSymbolization() {
+  auto *Object = cast<object::ObjectFile>(Binary.getBinary());
+  auto BuildIdOr = object::getBuildID(Object);
+  if (!BuildIdOr.has_value())
+    return make_error<StringError>(Twine("No build id found in binary ") +
+                                       Binary.getBinary()->getFileName(),
+                                   inconvertibleErrorCode());
+  llvm::ArrayRef<uint8_t> BinaryId = BuildIdOr.value();
+
+  int NumMatched = 0;
+  for (const auto &Entry : SegmentInfo) {
+    llvm::ArrayRef<uint8_t> SegmentId(Entry.BuildId, Entry.BuildIdSize);
+    if (BinaryId == SegmentId) {
+      // We assume only one text segment in the main binary for simplicity and
+      // reduce the overhead of checking multiple ranges during symbolization.
+      if (++NumMatched > 1) {
+        return make_error<StringError>(
+            "We expect only one executable segment in the profiled binary",
+            inconvertibleErrorCode());
+      }
+      ProfiledTextSegmentStart = Entry.Start;
+      ProfiledTextSegmentEnd = Entry.End;
+    }
+  }
+  assert(NumMatched != 0 && "No matching executable segments in segment info.");
+  assert(PreferredTextSegmentAddress == 0 ||
+         (PreferredTextSegmentAddress == ProfiledTextSegmentStart) &&
+             "Expect text segment address to be 0 or equal to profiled text "
+             "segment start.");
+  return Error::success();
+}
+
 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<FrameId> *;
   llvm::MapVector<GlobalValue::GUID, llvm::SetVector<LocationPtr>>
       PerFunctionCallSites;
 
while (Next < DataBuffer->getBufferEnd()) {
     Next += Header->TotalSize;
   }
 
   return Error::success();
 }
 
 object::SectionedAddress
 RawMemProfReader::getModuleOffset(const uint64_t VirtualAddress) {
-  LLVM_DEBUG({
-  SegmentEntry *ContainingSegment = nullptr;
-  for (auto &SE : SegmentInfo) {
-    if (VirtualAddress > SE.Start && VirtualAddress <= SE.End) {
-      ContainingSegment = &SE;
-    }
-  }
-
-  // Ensure that the virtual address is valid.
-  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.
+  if (VirtualAddress > ProfiledTextSegmentStart &&
+      VirtualAddress <= ProfiledTextSegmentEnd) {
+    // For PIE binaries, the preferred address is zero and we adjust the virtual
+    // address by start of the profiled segment assuming that the offset of the
+    // segment in the binary is zero. For non-PIE binaries the preferred and
+    // profiled segment addresses should be equal and this is a no-op.
+    const uint64_t AdjustedAddress =
+        VirtualAddress + PreferredTextSegmentAddress - ProfiledTextSegmentStart;
+    return object::SectionedAddress{AdjustedAddress};
+  }
+  // Addresses which do not originate from the profiled text segment in the
+  // binary are not adjusted. These will fail symbolization and be filtered out
+  // during processing.
   return object::SectionedAddress{VirtualAddress};
 }
 
 Error RawMemProfReader::readNextRecord(GuidMemProfRecordPair &GuidRecord) {
   if (FunctionProfileData.empty())
     return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
 
   if (Iter == FunctionProfileData.end())

llvm/test/tools/llvm-profdata/memprof-pic.test

 REQUIRES: x86_64-linux
 
 Since the profile contains 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.
 
 To update the inputs used below run Inputs/update_memprof_inputs.sh /path/to/updated/clang
-RUN: not llvm-profdata show --memory %p/Inputs/pic.memprofraw --profiled-binary %p/Inputs/pic.memprofexe -o - 2>&1 | FileCheck %s
+RUN: llvm-profdata show --memory %p/Inputs/pic.memprofraw --profiled-binary %p/Inputs/pic.memprofexe -o - 2>&1 | FileCheck %s
 
-CHECK: Unsupported position independent code
+CHECK:  MemprofProfile:
+CHECK-NEXT:   Summary:
+CHECK-NEXT:     Version: 3
+CHECK-NEXT:     NumSegments: {{[0-9]+}}
+CHECK-NEXT:     NumMibInfo: 2
+CHECK-NEXT:     NumAllocFunctions: 1
+CHECK-NEXT:     NumStackOffsets: 2
+CHECK-NEXT:   Segments:
+CHECK-NEXT:   -
+CHECK-NEXT:     BuildId: {{[[:xdigit:]]+}}
+CHECK-NEXT:     Start: 0x{{[[:xdigit:]]+}}
+CHECK-NEXT:     End: 0x{{[[:xdigit:]]+}}
+CHECK-NEXT:     Offset: 0x{{[[:xdigit:]]+}}
+CHECK-NEXT:   -
+
+CHECK:   Records:
+CHECK-NEXT:   -
+CHECK-NEXT:     FunctionGUID: {{[0-9]+}}
+CHECK-NEXT:     AllocSites:
+CHECK-NEXT:     -
+CHECK-NEXT:       Callstack:
+CHECK-NEXT:       -
+CHECK-NEXT:         Function: {{[0-9]+}}
+CHECK-NEXT:         SymbolName: main
+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: {{[0-9]+}}
+CHECK-NEXT:         DeallocTimestamp: {{[0-9]+}}
+CHECK-NEXT:         TotalLifetime: 0
+CHECK-NEXT:         MinLifetime: 0
+CHECK-NEXT:         MaxLifetime: 0
+CHECK-NEXT:         AllocCpuId: {{[0-9]+}}
+CHECK-NEXT:         DeallocCpuId: {{[0-9]+}}
+CHECK-NEXT:         NumMigratedCpu: 0
+CHECK-NEXT:         NumLifetimeOverlaps: 0
+CHECK-NEXT:         NumSameAllocCpu: 0
+CHECK-NEXT:         NumSameDeallocCpu: 0
+CHECK-NEXT:         DataTypeId: {{[0-9]+}}
+CHECK-NEXT:         TotalAccessDensity: 20
+CHECK-NEXT:         MinAccessDensity: 20
+CHECK-NEXT:         MaxAccessDensity: 20
+CHECK-NEXT:         TotalLifetimeAccessDensity: 20000
+CHECK-NEXT:         MinLifetimeAccessDensity: 20000
+CHECK-NEXT:         MaxLifetimeAccessDensity: 20000
+CHECK-NEXT:     -
+CHECK-NEXT:       Callstack:
+CHECK-NEXT:       -
+CHECK-NEXT:         Function: {{[0-9]+}}
+CHECK-NEXT:         SymbolName: main
+CHECK-NEXT:         LineOffset: 4
+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: {{[0-9]+}}
+CHECK-NEXT:         DeallocTimestamp: {{[0-9]+}}
+CHECK-NEXT:         TotalLifetime: 0
+CHECK-NEXT:         MinLifetime: 0
+CHECK-NEXT:         MaxLifetime: 0
+CHECK-NEXT:         AllocCpuId: {{[0-9]+}}
+CHECK-NEXT:         DeallocCpuId: {{[0-9]+}}
+CHECK-NEXT:         NumMigratedCpu: 0
+CHECK-NEXT:         NumLifetimeOverlaps: 0
+CHECK-NEXT:         NumSameAllocCpu: 0
+CHECK-NEXT:         NumSameDeallocCpu: 0
+CHECK-NEXT:         DataTypeId: {{[0-9]+}}
+CHECK-NEXT:         TotalAccessDensity: 20
+CHECK-NEXT:         MinAccessDensity: 20
+CHECK-NEXT:         MaxAccessDensity: 20
+CHECK-NEXT:         TotalLifetimeAccessDensity: 20000
+CHECK-NEXT:         MinLifetimeAccessDensity: 20000
+CHECK-NEXT:         MaxLifetimeAccessDensity: 20000