[llvm-profgen] On-demand symbolization

[wlei]

Summary

Previously we do symbolization for all the functions and actually we only need the symbols that's hit by the samples.

This can significantly speed up the time for large size binary.

Optimization for per-inliner will come along with next patch.

[hoy]

Lgtm, thanks for the improvement.

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.h

174

We don't seem to use the size (value of the map)?

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.h

174

Oh, it's used in the next diff of preinliner optimization, tried to make a shortcut :)

[wenlei]

lgtm, thanks.

llvm/tools/llvm-profgen/ProfiledBinary.h

174

ok, thanks for clarification.

335

nit: emplace? Also see if we can save a lookup by following how such caching is handled in SampleProfileLoader::findFunctionSamples.

[wlei]

chang to use emplace and fix lint

llvm/tools/llvm-profgen/PerfReader.h

 
 // Sample counter with context to support context-sensitive profile
 using ContextSampleCounterMap =
     std::unordered_map<Hashable<ContextKey>, SampleCounter,
                        Hashable<ContextKey>::Hash, Hashable<ContextKey>::Equal>;
 
 struct FrameStack {
   SmallVector<uint64_t, 16> Stack;
-  const ProfiledBinary *Binary;
-  FrameStack(const ProfiledBinary *B) : Binary(B) {}
+  ProfiledBinary *Binary;
+  FrameStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
     Stack.push_back(Cur->Address);
     return true;
   }
 
   void popFrame() {
     if (!Stack.empty())
       Stack.pop_back();
   }
   std::shared_ptr<StringBasedCtxKey> getContextKey();
 };
 
 struct ProbeStack {
   SmallVector<const MCDecodedPseudoProbe *, 16> Stack;
-  const ProfiledBinary *Binary;
-  ProbeStack(const ProfiledBinary *B) : Binary(B) {}
+  ProfiledBinary *Binary;
+  ProbeStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
     const MCDecodedPseudoProbe *CallProbe =
         Binary->getCallProbeForAddr(Cur->Address);
     // We may not find a probe for a merged or external callsite.
     // Callsite merging may cause the loss of original probe IDs.
     // Cutting off the context from here since the inliner will
     // not know how to consume a context with unknown callsites.
     if (!CallProbe)

 LBR's source, which is the linear unwinding. As instruction from linear range
 can come from different function by inlining, linear unwinding will do the range
 splitting and record counters by the range with same inline context. Over those
 unwinding process we will record each call stack as context id and LBR/linear
 range as sample counter for further CS profile generation.
 */
 class VirtualUnwinder {
 public:
-  VirtualUnwinder(ContextSampleCounterMap *Counter, const ProfiledBinary *B)
+  VirtualUnwinder(ContextSampleCounterMap *Counter, ProfiledBinary *B)
       : CtxCounterMap(Counter), Binary(B) {}
   bool unwind(const PerfSample *Sample, uint64_t Repeat);
   std::set<uint64_t> &getUntrackedCallsites() { return UntrackedCallsites; }
 
 private:
   bool isCallState(UnwindState &State) const {
     // The tail call frame is always missing here in stack sample, we will
     // use a specific tail call tracker to infer it.
private:
 
   void recordRangeCount(uint64_t Start, uint64_t End, UnwindState &State,
                         uint64_t Repeat);
   void recordBranchCount(const LBREntry &Branch, UnwindState &State,
                          uint64_t Repeat);
 
   ContextSampleCounterMap *CtxCounterMap;
   // Profiled binary that current frame address belongs to
-  const ProfiledBinary *Binary;
+  ProfiledBinary *Binary;
   // Keep track of all untracked callsites
   std::set<uint64_t> UntrackedCallsites;
 };
 
 // Read perf trace to parse the events and samples.
 class PerfReaderBase {
 public:
   PerfReaderBase(ProfiledBinary *B) : Binary(B) {

llvm/tools/llvm-profgen/ProfiledBinary.h

class ProfiledBinary {
   // A list of text sections sorted by start RVA and size. Used to check
   // if a given RVA is a valid code address.
   std::set<std::pair<uint64_t, uint64_t>> TextSections;
   // An ordered map of mapping function's start offset to its name and
   // end offset.
   std::map<uint64_t, std::pair<std::string, uint64_t>> FuncStartOffsetMap;
   // Offset to context location map. Used to expand the context.
   std::unordered_map<uint64_t, SampleContextFrameVector> Offset2LocStackMap;
+
+  // Offset to instruction size map. Also used for quick offset lookup.
+  std::unordered_map<uint64_t, uint64_t> Offset2InstSizeMap;
+
   // An array of offsets of all instructions sorted in increasing order. The
   // sorting is needed to fast advance to the next forward/backward instruction.
   std::vector<uint64_t> CodeAddrs;
   // A set of call instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> CallAddrs;
   // A set of return instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> RetAddrs;
 
const std::vector<uint64_t> &getPreferredTextSegmentAddresses() const {
     return PreferredTextSegmentAddresses;
   }
   const std::vector<uint64_t> &getTextSegmentOffsets() const {
     return TextSegmentOffsets;
   }
 
   bool addressIsCode(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
-    return Offset2LocStackMap.find(Offset) != Offset2LocStackMap.end();
+    return Offset2InstSizeMap.find(Offset) != Offset2InstSizeMap.end();
   }
   bool addressIsCall(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     return CallAddrs.count(Offset);
   }
   bool addressIsReturn(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     return RetAddrs.count(Offset);
OffsetRange findFuncOffsetRange(uint64_t Offset) {
     I--;
     return {I->first, I->second.second};
   }
 
   uint32_t getFuncSizeForContext(SampleContext &Context) {
     return FuncSizeTracker.getFuncSizeForContext(Context);
   }
 
-  const SampleContextFrameVector &getFrameLocationStack(uint64_t Offset) const {
-    auto I = Offset2LocStackMap.find(Offset);
-    assert(I != Offset2LocStackMap.end() &&
-           "Can't find location for offset in the binary");
-    return I->second;
+  const SampleContextFrameVector &
+  getFrameLocationStack(uint64_t Offset, bool UseProbeDiscriminator = false) {
+    auto I = Offset2LocStackMap.emplace(Offset, SampleContextFrameVector());
+    if (I.second) {
+      InstructionPointer IP(this, Offset);
+      I.first->second = symbolize(IP, true, UseProbeDiscriminator);
+    }
+    return I.first->second;
   }
 
   Optional<SampleContextFrame> getInlineLeafFrameLoc(uint64_t Offset) {
     const auto &Stack = getFrameLocationStack(Offset);
     if (Stack.empty())
       return {};
     return Stack.back();
   }
 
   // Compare two addresses' inline context
-  bool inlineContextEqual(uint64_t Add1, uint64_t Add2) const;
+  bool inlineContextEqual(uint64_t Add1, uint64_t Add2);
 
   // Get the full context of the current stack with inline context filled in.
   // It will search the disassembling info stored in Offset2LocStackMap. This is
   // used as the key of function sample map
   SampleContextFrameVector
   getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
-                     bool &WasLeafInlined) const;
+                     bool &WasLeafInlined);
 
   const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
 
   void getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
                                 SampleContextFrameVector &InlineContextStack,
                                 bool IncludeLeaf = false) const {

llvm/tools/llvm-profgen/ProfiledBinary.cpp

void ProfiledBinary::load() {
 
   // Use function start and return address to infer prolog and epilog
   ProEpilogTracker.inferPrologOffsets(FuncStartOffsetMap);
   ProEpilogTracker.inferEpilogOffsets(RetAddrs);
 
   // TODO: decode other sections.
 }
 
-bool ProfiledBinary::inlineContextEqual(uint64_t Address1,
-                                        uint64_t Address2) const {
+bool ProfiledBinary::inlineContextEqual(uint64_t Address1, uint64_t Address2) {
   uint64_t Offset1 = virtualAddrToOffset(Address1);
   uint64_t Offset2 = virtualAddrToOffset(Address2);
   const SampleContextFrameVector &Context1 = getFrameLocationStack(Offset1);
   const SampleContextFrameVector &Context2 = getFrameLocationStack(Offset2);
   if (Context1.size() != Context2.size())
     return false;
   if (Context1.empty())
     return false;
   // The leaf frame contains location within the leaf, and it
   // needs to be remove that as it's not part of the calling context
   return std::equal(Context1.begin(), Context1.begin() + Context1.size() - 1,
                     Context2.begin(), Context2.begin() + Context2.size() - 1);
 }
 
 SampleContextFrameVector
 ProfiledBinary::getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
-                                   bool &WasLeafInlined) const {
+                                   bool &WasLeafInlined) {
   SampleContextFrameVector ContextVec;
   // Process from frame root to leaf
   for (auto Address : Stack) {
     uint64_t Offset = virtualAddrToOffset(Address);
     const SampleContextFrameVector &ExpandedContext =
         getFrameLocationStack(Offset);
     // An instruction without a valid debug line will be ignored by sample
     // processing
while (Offset < EndOffset) {
     }
 
     if (Disassembled) {
       const MCInstrDesc &MCDesc = MII->get(Inst.getOpcode());
       // Populate a vector of the symbolized callsite at this location
       // We don't need symbolized info for probe-based profile, just use an
       // empty stack as an entry to indicate a valid binary offset
       SampleContextFrameVector SymbolizedCallStack;
-      if (!UsePseudoProbes || TrackFuncContextSize) {
+      if (TrackFuncContextSize) {
         InstructionPointer IP(this, Offset);
         // TODO: reallocation of Offset2LocStackMap will lead to dangling
         // strings We need ProfiledBinary to owned these string.
         Offset2LocStackMap[Offset] = symbolize(IP, true, UsePseudoProbes);
         SampleContextFrameVector &SymbolizedCallStack =
             Offset2LocStackMap[Offset];
         // Record instruction size for the corresponding context
         if (TrackFuncContextSize && !SymbolizedCallStack.empty())
           FuncSizeTracker.addInstructionForContext(Offset2LocStackMap[Offset],
                                                    Size);
-      } else {
-        Offset2LocStackMap[Offset] = SampleContextFrameVector();
       }
+      // Record instruction size.
+      Offset2InstSizeMap[Offset] = Size;
 
       // Populate address maps.
       CodeAddrs.push_back(Offset);
       if (MCDesc.isCall())
         CallAddrs.insert(Offset);
       else if (MCDesc.isReturn())
         RetAddrs.insert(Offset);