[llvm-profgen] Fix inconsistent loading address issues

[wlei]

Summary

This is to fix two issues related with loading address:

1. When multiple MMAPs occur and their loading address are different, before it only used the first MMap as base address, all perf address after it used the wrong base address.

2. For pseudo probe profile, the address is always based on preferred loading address. If the base address is not equal to the preferred loading address, the pseudo probe address query will be wrong.

Solution: Instead of converting the address to offset lazily, right now all the address after parsing are converted on the fly based on preferred loading address in the parsing time. There is no "offset" used in profile generator any more.

[wenlei]

llvm/tools/llvm-profgen/PerfReader.cpp

1195

Where do we use getBaseAddress in probe queries? Can we use getPreferredBaseAddress directly there? Ideally there should be one definition for base address which is the actual executable segment load address, and changing it here makes it inconsistent.

llvm/tools/llvm-profgen/PerfReader.h

374

Perhaps AddrBasedCtxKey as name is just fine? This is in contrast with StringBasedCtxKey, so Addr as a general term isn't a big deal, I don't have strong opinion though..

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

1195

Agreed that setBaseAddress should be called once to be consistent. Pseudo probe decoding is based on the preferred address. I think a reasonable fix could be

1. make the decoding offset-based. This will require a change to the bolt code. base; or
2. do an offset to preferred addr translation for every probe look up in llvm-profgen. This is being done in some places, eg, extractPrefixContextStack.

#2 sounds more practical. WDYT?

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

1195

Hmm, I found it's very tricky case.

Where do we use getBaseAddress in probe queries?

We call it in Binary->offsetToVirtualAddr.

However, we call IP.advance() to iterate all the address in the range, in IP it will convert the Address back to offset using the base address. like Offset + preferred Address - Mmap loading address. That's another inconsistency.. Then we need to change code inside of the IP. but IP also used in PerfReader

llvm/tools/llvm-profgen/PerfReader.h

374

Sounds good.

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

1195

Or can we use virtual address(Offset + preferred Address) for all the code(PerfReader and ProfileGenerator)? to do this, we can convert the address at the very beginning of PerfReader. like Address = PhysicalAddress - Mmap loading address + PreferredAddress

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

1195

I see. This is complicated. How about we always use preferred address when doing this offset to virtual addr conversion, and we rebase all LBR and stack sample addresses to be on the preferred load address in the the perf reader ?

[wlei]

changed all offset usage to preferred address based virtual address

[hoy]

In D126827#3613023, @wlei wrote:

changed all offset usage to preferred address based virtual address

Thanks for the work! It should be more reliable and also looks cleaner.

llvm/tools/llvm-profgen/PerfReader.cpp

749

I'm a bit confused here. What is the semantics of using both offset and LoadableSegmentAsBase?

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

749

I think this is to be compatible to our internal old tool where some services use offset and the FirstLoadableAddress as the output of unsymbolized profile.

The original patch is https://reviews.llvm.org/D113727.

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

749

Do we still need this now? The offset here is computed by subtracting the runtime base address from preferred-based virtual address. This seems different from what previous computed, ie., subtracting the runtime base address from runtime virtual address.

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

749

It should be the same behavior as previous offset computation.

This is about how we define the "offset": Assuming MMAP loading address is missing or just equal to preferred-based-address, then the virtual address = offset + preferred-based-address. and for this two branches, if it's UseLoadableSegmentAsBase is False, then the offset is virtual address - preferred-based-address --> offset + preferred-based-address - preferred-based-address which is the original "offset". but when UseLoadableSegmentAsBase is True, the offset is virtual address - FirstLoadableAddress -->offset + preferred-based-address - FirstLoadableAddress
so the offset is not the offset as we defined before. However, even it's "wrong", if depends on the source of input, our internal tool which just produced this "wrong" offset, so llvm-profgen also need to use the same way to recover it.

Yeah, but if we won't use the old tool anymore, I think we can remove all the UseLoadableSegmentAsBase code.

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

749

so the offset is not the offset as we defined before. However, even it's "wrong",

Could the "wrong" offsets now cause trouble for pseudo probes which always have the "correct" offset? It should work fine previously?

Yeah, but if we won't use the old tool anymore, I think we can remove all the UseLoadableSegmentAsBase code.

Agreed.

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

749

I looked a bit deeper. The "wrong" offsets should work. getFirstLoadableAddress stands for the preferred first load address, not the runtime one where I got the impression from its definition.

// The runtime base address that the first loadabe segment is loaded at.
uint64_t FirstLoadableAddress = 0;

Then we no longer need to remove the feature in this change.

[wenlei]

Agreed that having everything canonicalized to use preferred load address as base is clean and practical. Thanks for making the changes. LGTM with a nit.

llvm/tools/llvm-profgen/ProfiledBinary.h

353 ↗

(On Diff #440317)

nit: canonicalizeVirtualAddress is probably a better name.

[wlei]

rebase and rename convertAddress --> canonicalizeVirtualAddress

llvm/test/tools/llvm-profgen/Inputs/noinline-cs-pseudoprobe.perfscript

-PERF_RECORD_MMAP2 1243676/1243676: [0x201000(0x1000) @ 0 00:1d 224517108 1044165]: r-xp /home/noinline-cs-pseudoprobe.perfbin
+PERF_RECORD_MMAP2 1243676/1243676: [0x301000(0x1000) @ 0 00:1d 224517108 1044165]: r-xp /home/noinline-cs-pseudoprobe.perfbin
 
-            20179e
-            2017f9
+            30179e
+            3017f9
       7f83e84e7793
   5541f689495641d7
- 0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0
+ 0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0
 
-            2017c4
-            2017f9
+            3017c4
+            3017f9
       7f83e84e7793
   5541f689495641d7
- 0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0
+ 0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0
 
-            2017c4
-            2017f9
+            3017c4
+            3017f9
       7f83e84e7793
   5541f689495641d7
- 0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0
+ 0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0  0x3017bf/0x301760/P/-/-/0  0x3017cf/0x30179e/P/-/-/0  0x30177f/0x3017c4/P/-/-/0

llvm/tools/llvm-profgen/PerfReader.h

bool isEqual(const ContextKey *K) const override {
     return Context == Other->Context;
   }
 
   void genHashCode() override {
     HashCode = hash_value(SampleContextFrames(Context));
   }
 };
 
-// Address-based context id
-struct AddrBasedCtxKey : public ContextKey {
+// Offset-based context id
+struct OffsetBasedCtxKey : public ContextKey {
   SmallVector<uint64_t, 16> Context;
 
   bool WasLeafInlined;
-  AddrBasedCtxKey() : ContextKey(CK_AddrBased), WasLeafInlined(false){};
+  OffsetBasedCtxKey() : ContextKey(CK_AddrBased), WasLeafInlined(false){};
   static bool classof(const ContextKey *K) {
     return K->getKind() == CK_AddrBased;
   }
 
   bool isEqual(const ContextKey *K) const override {
-    const AddrBasedCtxKey *Other = dyn_cast<AddrBasedCtxKey>(K);
+    const OffsetBasedCtxKey *Other = dyn_cast<OffsetBasedCtxKey>(K);
     return Context == Other->Context;
   }
 
   void genHashCode() override {
     HashCode = hash_combine_range(Context.begin(), Context.end());
   }
 };
 

 
 struct FrameStack {
   SmallVector<uint64_t, 16> Stack;
   ProfiledBinary *Binary;
   FrameStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
     assert(!Cur->isExternalFrame() &&
            "External frame's not expected for context stack.");
-    Stack.push_back(Cur->Address);
+    Stack.push_back(Binary->virtualAddrToOffset(Cur->Address));
     return true;
   }
 
   void popFrame() {
     if (!Stack.empty())
       Stack.pop_back();
   }
   std::shared_ptr<StringBasedCtxKey> getContextKey();
 };
 
 struct AddressStack {
   SmallVector<uint64_t, 16> Stack;
   ProfiledBinary *Binary;
   AddressStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
     assert(!Cur->isExternalFrame() &&
            "External frame's not expected for context stack.");
-    Stack.push_back(Cur->Address);
+    Stack.push_back(Binary->virtualAddrToOffset(Cur->Address));
     return true;
   }
 
   void popFrame() {
     if (!Stack.empty())
       Stack.pop_back();
   }
-  std::shared_ptr<AddrBasedCtxKey> getContextKey();
+  std::shared_ptr<OffsetBasedCtxKey> getContextKey();
 };
 
 /*
 As in hybrid sample we have a group of LBRs and the most recent sampling call
 stack, we can walk through those LBRs to infer more call stacks which would be
 used as context for profile. VirtualUnwinder is the class to do the call stack
 unwinding based on LBR state. Two types of unwinding are processd here:
 1) LBR unwinding and 2) linear range unwinding.

llvm/tools/llvm-profgen/PerfReader.cpp

 
 std::shared_ptr<StringBasedCtxKey> FrameStack::getContextKey() {
   std::shared_ptr<StringBasedCtxKey> KeyStr =
       std::make_shared<StringBasedCtxKey>();
   KeyStr->Context = Binary->getExpandedContext(Stack, KeyStr->WasLeafInlined);
   return KeyStr;
 }
 
-std::shared_ptr<AddrBasedCtxKey> AddressStack::getContextKey() {
-  std::shared_ptr<AddrBasedCtxKey> KeyStr = std::make_shared<AddrBasedCtxKey>();
+std::shared_ptr<OffsetBasedCtxKey> AddressStack::getContextKey() {
+  std::shared_ptr<OffsetBasedCtxKey> KeyStr =
+      std::make_shared<OffsetBasedCtxKey>();
   KeyStr->Context = Stack;
   CSProfileGenerator::compressRecursionContext<uint64_t>(KeyStr->Context);
   CSProfileGenerator::trimContext<uint64_t>(KeyStr->Context);
   return KeyStr;
 }
 
 template <typename T>
 void VirtualUnwinder::collectSamplesFromFrame(UnwindState::ProfiledFrame *Cur,
if (Event.Offset == Binary->getTextSegmentOffset()) {
     }
   }
 }
 
 static std::string getContextKeyStr(ContextKey *K,
                                     const ProfiledBinary *Binary) {
   if (const auto *CtxKey = dyn_cast<StringBasedCtxKey>(K)) {
     return SampleContext::getContextString(CtxKey->Context);
-  } else if (const auto *CtxKey = dyn_cast<AddrBasedCtxKey>(K)) {
+  } else if (const auto *CtxKey = dyn_cast<OffsetBasedCtxKey>(K)) {
     std::ostringstream OContextStr;
     for (uint32_t I = 0; I < CtxKey->Context.size(); I++) {
       if (OContextStr.str().size())
         OContextStr << " @ ";
-      OContextStr << "0x"
-                  << to_hexString(
-                         Binary->virtualAddrToOffset(CtxKey->Context[I]),
-                         false);
+      OContextStr << "0x" << to_hexString(CtxKey->Context[I], false);
     }
     return OContextStr.str();
   } else {
     llvm_unreachable("unexpected key type");
   }
 }
 
 void HybridPerfReader::unwindSamples() {
for (auto &I : Counter) {
         Start = Binary->offsetToVirtualAddr(Start);
         End = Binary->offsetToVirtualAddr(End);
       }
 
       if (UseOffset && UseLoadableSegmentAsBase) {
         Start -= Binary->getFirstLoadableAddress();
         End -= Binary->getFirstLoadableAddress();
       }
 
       OS.indent(Indent);
       OS << Twine::utohexstr(Start) << Separator << Twine::utohexstr(End) << ":"
          << I.second << "\n";
     }
   };
 
   for (auto &CI : OrderedCounters) {
     uint32_t Indent = 0;
emitWarningSummary(NumLeadingOutgoingLBR, NumTotalSample,
                      "of samples have leading external LBR.");
 
   // Generate unsymbolized profile.
   warnTruncatedStack();
   warnInvalidRange();
   generateUnsymbolizedProfile();
   AggregatedSamples.clear();
 
+  // Base address can be loaded from MMP event, but pseudo probe address is
+  // always based on preferred loading address. Recover the base address to
+  // preferred loading address in order to match pseduo probe address correctly.
+  Binary->setBaseAddress(Binary->getPreferredBaseAddress());
+
   if (SkipSymbolization)
     writeUnsymbolizedProfile(OutputFilename);
 }
 
 } // end namespace sampleprof
 } // end namespace llvm

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 
 void ProfileGeneratorBase::collectProfiledFunctions() {
   std::unordered_set<const BinaryFunction *> ProfiledFunctions;
   if (SampleCounters) {
     // Go through all the stacks, ranges and branches in sample counters, use
     // the start of the range to look up the function it belongs and record the
     // function.
     for (const auto &CI : *SampleCounters) {
-      if (const auto *CtxKey = dyn_cast<AddrBasedCtxKey>(CI.first.getPtr())) {
-        for (auto Addr : CtxKey->Context) {
-          if (FuncRange *FRange = Binary->findFuncRangeForOffset(
-                  Binary->virtualAddrToOffset(Addr)))
+      if (const auto *CtxKey = dyn_cast<OffsetBasedCtxKey>(CI.first.getPtr())) {
+        for (auto Offset : CtxKey->Context) {
+          if (FuncRange *FRange = Binary->findFuncRangeForOffset(Offset))
             ProfiledFunctions.insert(FRange->Func);
         }
       }
 
       for (auto Item : CI.second.RangeCounter) {
         uint64_t StartOffset = Item.first.first;
         if (FuncRange *FRange = Binary->findFuncRangeForOffset(StartOffset))
           ProfiledFunctions.insert(FRange->Func);
do {
         for (const auto &Probe : It->second) {
           ProbeCounter[&Probe] += Count;
         }
       }
     } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
-static void
-extractPrefixContextStack(SampleContextFrameVector &ContextStack,
-                          const SmallVectorImpl<uint64_t> &Addresses,
+static void extractPrefixContextStack(SampleContextFrameVector &ContextStack,
+                                      const SmallVectorImpl<uint64_t> &Offsets,
                                       ProfiledBinary *Binary) {
   SmallVector<const MCDecodedPseudoProbe *, 16> Probes;
-  for (auto Addr : reverse(Addresses)) {
-    const MCDecodedPseudoProbe *CallProbe = Binary->getCallProbeForAddr(Addr);
+  for (auto Offset : reverse(Offsets)) {
+    const MCDecodedPseudoProbe *CallProbe =
+        Binary->getCallProbeForAddr(Binary->offsetToVirtualAddr(Offset));
     // These could be the cases when a probe is not found at a calliste. Cutting
     // off the context from here since the inliner will not know how to consume
     // a context with unknown callsites.
     // 1. for functions that are not sampled when
     // --decode-probe-for-profiled-functions-only is on.
     // 2. for a merged callsite. Callsite merging may cause the loss of original
     // probe IDs.
     // 3. for an external callsite.
for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStack, true);
   }
 }
 
 void CSProfileGenerator::generateProbeBasedProfile() {
   // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;
   for (const auto &CI : *SampleCounters) {
-    const AddrBasedCtxKey *CtxKey =
-        dyn_cast<AddrBasedCtxKey>(CI.first.getPtr());
+    const OffsetBasedCtxKey *CtxKey =
+        dyn_cast<OffsetBasedCtxKey>(CI.first.getPtr());
     SampleContextFrameVector ContextStack;
     extractPrefixContextStack(ContextStack, CtxKey->Context, Binary);
     // Fill in function body samples from probes, also infer caller's samples
     // from callee's probe
     populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStack);
     // Fill in boundary samples for a call probe
     populateBoundarySamplesWithProbes(CI.second.BranchCounter, ContextStack);
   }

llvm/tools/llvm-profgen/ProfiledBinary.cpp

 
 SampleContextFrameVector
 ProfiledBinary::getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
                                    bool &WasLeafInlined) {
   SampleContextFrameVector ContextVec;
   if (Stack.empty())
     return ContextVec;
   // Process from frame root to leaf
-  for (auto Address : Stack) {
-    uint64_t Offset = virtualAddrToOffset(Address);
+  for (auto Offset : Stack) {
     const SampleContextFrameVector &ExpandedContext =
         getFrameLocationStack(Offset);
     // An instruction without a valid debug line will be ignored by sample
     // processing
     if (ExpandedContext.empty())
       return SampleContextFrameVector();
     // Set WasLeafInlined to the size of inlined frame count for the last
     // address which is leaf