[llvm-profgen] Do not cache the frame location stack during computing inlined context size

[wlei]

Summary

In computeInlinedContextSizeForRange, the offset of range is only used one time, there is no need to cache the frame location stack.
Measured on one internal service binary, this can save 2GB memory usage and reduce a small run time (avoid one hash search).

[hoy]

Nice catch!

llvm/tools/llvm-profgen/ProfiledBinary.h

496–502

How about extend this function with a parameter for caching or not?

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.h

496–502

It seems the return types of the two functions should be different, the one without caching should return the value type instead of the reference as we don't store it to a map.

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.h

496–502

I see. Then it makes sense to leave them separate.

How about flip the names, i.e

getFrameLocationStackWithoutCaching -> getFrameLocationStack

getFrameLocationStack -> getCachedFrameLocation

Caching is expansive so we want users to be explicitly aware of it. Also perhaps in the future we want to use two containers for cached contexts, one for UseProbeDiscriminator==true and one for UseProbeDiscriminator==false.

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.h

496–502

Renamed, thanks for the suggestion!

[wlei]

renaming according to reviewer's suggestion

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.h

502

nit: getCachedFrameLocation -> getCachedFrameLocationStack

Sorry for missing it previously.

[wlei]

Updating D128859: [llvm-profgen] Do not cache the frame location stack during computing inlined context size

[wenlei]

Good find, thanks for the memory improvement!

llvm/tools/llvm-profgen/ProfileGenerator.cpp

for (const auto &Range : preprocessRangeCounter(RangeCounter)) {
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
     if (IP.Address > RangeEnd)
       continue;
 
     do {
       uint64_t Offset = Binary->virtualAddrToOffset(IP.Address);
-      const SampleContextFrameVector &FrameVec =
+      const SampleContextFrameVector FrameVec =
           Binary->getFrameLocationStack(Offset);
       if (!FrameVec.empty()) {
         // FIXME: As accumulating total count per instruction caused some
         // regression, we changed to accumulate total count per byte as a
         // workaround. Tuning hotness threshold on the compiler side might be
         // necessary in the future.
         FunctionSamples &FunctionProfile = getLeafProfileAndAddTotalSamples(
             FrameVec, Count * Binary->getInstSize(Offset));
for (const auto &Entry : BranchCounters) {
     uint64_t Count = Entry.second;
     assert(Count != 0 && "Unexpected zero weight branch");
 
     StringRef CalleeName = getCalleeNameForOffset(TargetOffset);
     if (CalleeName.size() == 0)
       continue;
     // Record called target sample and its count.
     const SampleContextFrameVector &FrameVec =
-        Binary->getFrameLocationStack(SourceOffset);
+        Binary->getCachedFrameLocation(SourceOffset);
     if (!FrameVec.empty()) {
       FunctionSamples &FunctionProfile =
           getLeafProfileAndAddTotalSamples(FrameVec, 0);
       FunctionProfile.addCalledTargetSamples(
           FrameVec.back().Location.LineOffset,
           getBaseDiscriminator(FrameVec.back().Location.Discriminator),
           CalleeName, Count);
     }

llvm/tools/llvm-profgen/ProfiledBinary.h

public:
 
   uint32_t getFuncSizeForContext(const ContextTrieNode *ContextNode) {
     return FuncSizeTracker.getFuncSizeForContext(ContextNode);
   }
 
   // Load the symbols from debug table and populate into symbol list.
   void populateSymbolListFromDWARF(ProfileSymbolList &SymbolList);
 
-  const SampleContextFrameVector &
+  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 symbolize(IP, true, UseProbeDiscriminator);
+
+  const SampleContextFrameVector &
+  getCachedFrameLocation(uint64_t Offset, bool UseProbeDiscriminator = false) {
+    auto I = Offset2LocStackMap.emplace(Offset, SampleContextFrameVector());
+    if (I.second)
+      I.first->second = getFrameLocationStack(Offset, UseProbeDiscriminator);
     return I.first->second;
   }
 
   Optional<SampleContextFrame> getInlineLeafFrameLoc(uint64_t Offset) {
-    const auto &Stack = getFrameLocationStack(Offset);
+    const auto &Stack = getCachedFrameLocation(Offset);
     if (Stack.empty())
       return {};
     return Stack.back();
   }
 
   void flushSymbolizer() { Symbolizer.reset(); }
 
   // Compare two addresses' inline context

llvm/tools/llvm-profgen/ProfiledBinary.cpp

void ProfiledBinary::load() {
   warnNoFuncEntry();
 
   // TODO: decode other sections.
 }
 
 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);
+  const SampleContextFrameVector &Context1 = getCachedFrameLocation(Offset1);
+  const SampleContextFrameVector &Context2 = getCachedFrameLocation(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) {
   SampleContextFrameVector ContextVec;
   if (Stack.empty())
     return ContextVec;
   // Process from frame root to leaf
   for (auto Address : Stack) {
     uint64_t Offset = virtualAddrToOffset(Address);
     const SampleContextFrameVector &ExpandedContext =
-        getFrameLocationStack(Offset);
+        getCachedFrameLocation(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
     WasLeafInlined = (ExpandedContext.size() > 1);
     ContextVec.append(ExpandedContext);
if (IP.Address != RangeBegin)
     WithColor::warning() << "Invalid start instruction at "
                          << format("%8" PRIx64, RangeBegin) << "\n";
 
   if (IP.Address >= RangeEnd)
     return;
 
   do {
     uint64_t Offset = virtualAddrToOffset(IP.Address);
-    const SampleContextFrameVector &SymbolizedCallStack =
+    const SampleContextFrameVector SymbolizedCallStack =
         getFrameLocationStack(Offset, UsePseudoProbes);
     uint64_t Size = Offset2InstSizeMap[Offset];
 
     // Record instruction size for the corresponding context
     FuncSizeTracker.addInstructionForContext(SymbolizedCallStack, Size);
 
   } while (IP.advance() && IP.Address < RangeEnd);
 }