[llvm-profgen] On-demand track optimized-away inlinees for preinliner.

[hoy]

Summary

Tracking optimized-away inlinees based on all probes in a binary is expansive in terms of memory usage I'm making the tracking on-demand based on profiled functions only. This saves about 10% memory overall for a medium-sized benchmark.

Before:

note: After parsePerfTraces
note: Thu Jan 27 18:42:09 2022
note: VM: 8.68 GB   RSS: 8.39 GB
note: After computeSizeForProfiledFunctions
note: Thu Jan 27 18:42:41 2022
note: **VM: 10.63 GB   RSS: 10.20 GB**
note: After generateProbeBasedProfile
note: Thu Jan 27 18:45:49 2022
note: VM: 25.00 GB   RSS: 24.95 GB
note: After postProcessProfiles
note: Thu Jan 27 18:49:29 2022
note: VM: 26.34 GB   RSS: 26.27 GB

After:

note: After parsePerfTraces
note: Fri Jan 28 12:04:49 2022
note: VM: 8.68 GB   RSS: 7.65 GB
note: After computeSizeForProfiledFunctions
note: Fri Jan 28 12:05:26 2022
note: **VM: 8.68 GB   RSS: 8.42 GB**
note: After generateProbeBasedProfile
note: Fri Jan 28 12:08:03 2022
note: VM: 22.93 GB   RSS: 22.89 GB
note: After postProcessProfiles
note: Fri Jan 28 12:11:30 2022
note: VM: 24.27 GB   RSS: 24.22 GB

This should be a no-diff change in terms of profile quality.

[hoy]

Updating D118515: [llvm-profgen] On-demand track optimized-away inlinees for preinliner.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

608–609

It is weird to have profile generator flush symbolizer of ProfileBinary, how much saving is this comparing to on-demand context size tracking?

I assume the symbolization here is only needed for dwarf base profile, can we free symbolizer for probe profile generation right after ProfiledBinary::load?

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

Actually ProbeDecoder.getDummyInlineRoot().getChildren() is already a map, wondering can we make it possible to look up by function name without building an extra map?

I can see that the key is Guid, ProbeId pair - we can get Guid from names, but ProbeId under dummy root is index. What was the reason for top level nodes to have different probe Id instead of a dummy probe Id? For top level, we don't expect same name to appear more than once.

llvm/tools/llvm-profgen/ProfiledBinary.h

483

If this invalidates symbolizer, i.e making it not functional, we should also reset the pointer to null.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

608–609

The symbolizer consumes quite some memory and it's needed in both probe and dwarf case to calculate code sizes for inlinees. Perhaps I should just reset the symbolizer pointer instead of flushing it.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

Top level nodes should have a dummy probe id, i.e, 0. The key in ProbeDecoder.getDummyInlineRoot().getChildren() is the guid. and we need a guid to look up a node in the children map. Unfortunately, there isn't a func name to guid map available currently. What we have is the GUID2FuncDescMap. So I'm building a name to top level node map. Alternatively we can build a name to guid map here but that'll incur two hash lookups to get the node. Since either map is only used by llvm-profgen, I'm not placing the map to MCPseudoProbe which is also shared by the Bolt probe decorder.

[hoy]

Updating D118515: [llvm-profgen] On-demand track optimized-away inlinees for preinliner.

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

Unfortunately, there isn't a func name to guid map available currently

I thought function guid can be computed from names?

Top level nodes should have a dummy probe id, i.e, 0

if i read the code correctly in MCPseudoProbeDecoder::buildAddress2ProbeMap, that's not the case:

while (Data < End) {
    if (Root == Cur) {
      // Use a sequential id for top level inliner.
      Index = Root->getChildren().size();
    }

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

Ah, you are right on both. Now I completely remember why using guid based lookup didn't work. The key of the children map is <caller's guid, callsite probe id> and I was using callee's guid to lookup.

// A DFS-based decoding
  while (Data < End) {
    if (Root == Cur) {
      // Use a sequential id for top level inliner.
      Index = Root->getChildren().size();
    } else {
      ...
    }
    // Switch/add to a new tree node(inlinee)
    Cur = Cur->getOrAddNode(std::make_tuple(Cur->Guid, Index));

We are now building callee name to callee node map here.

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

I guess it then goes back to my original question.. why the index for top level node needs to be non-zero? (I remember it was all zero initially, and then changed to use index). But there's no actual call site calling into top level frames, so the probe id isn't meaningful. If we have 0 probe id for them, name/guid would be sufficient for looking up the probe frame.

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

351

The sequential id was made intentionally in https://reviews.llvm.org/D100235 for the use of reporting zero samples towards non-executed probes in a frame. Using zero index caused all top-level inliners to share the same probe inline frame.

Even with the zero index, guid based hash lookup still won't work. Note that the guid part of the key of the Children map is the caller guid, which is the guid of the dummy root, which is zero.

[wenlei]

lgtm, thanks.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

608–609

Ok, sounds good. It works for saving memory, but just wanted to note that this is hacky to flush symbolizer owned by ProfiledBinary at a random place by its user and yet APIs of ProfiledBinary depend on the availability of symbolizer..

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (Binary->usePseudoProbes()) {
     generateProbeBasedProfile();
   } else {
     generateLineNumBasedProfile();
   }
   postProcessProfiles();
 }
 
 void CSProfileGenerator::computeSizeForProfiledFunctions() {
-  // Hash map to deduplicate the function range and the item is a pair of
-  // function start and end offset.
-  std::unordered_map<uint64_t, uint64_t> AggregatedRanges;
+  std::unordered_set<const BinaryFunction *> ProfiledFunctions;
+
   // Go through all the ranges in the CS counters, use the start of the range to
-  // look up the function it belongs and record the function range.
+  // look up the function it belongs and record the function.
   for (const auto &CI : SampleCounters) {
     for (const auto &Item : CI.second.RangeCounter) {
       // FIXME: Filter the bogus crossing function range.
       uint64_t StartOffset = Item.first.first;
-      // Note that a function can be spilt into multiple ranges, so get all
-      // ranges of the function.
-      for (const auto &Range : Binary->getRangesForOffset(StartOffset))
-        AggregatedRanges[Range.first] = Range.second;
+      if (FuncRange *FRange = Binary->findFuncRangeForOffset(StartOffset))
+        ProfiledFunctions.insert(FRange->Func);
     }
   }
 
-  for (const auto &I : AggregatedRanges) {
-    uint64_t StartOffset = I.first;
-    uint64_t EndOffset = I.second;
-    Binary->computeInlinedContextSizeForRange(StartOffset, EndOffset);
-  }
+  for (auto *Func : ProfiledFunctions)
+    Binary->computeInlinedContextSizeForFunc(Func);
+
+  // Flush the symbolizer to save memory.
+  Binary->flushSymbolizer();
 }
 
 void CSProfileGenerator::generateLineNumBasedProfile() {
   for (const auto &CI : SampleCounters) {
     const auto *CtxKey = cast<StringBasedCtxKey>(CI.first.getPtr());
 
     // Get or create function profile for the range
     FunctionSamples &FunctionProfile =

llvm/tools/llvm-profgen/ProfiledBinary.h

public:
   // for the given context, try to retrieve the size of that function from
   // closest matching context.
   uint32_t getFuncSizeForContext(const SampleContext &Context);
 
   // For inlinees that are full optimized away, we can establish zero size using
   // their remaining probes.
   void trackInlineesOptimizedAway(MCPseudoProbeDecoder &ProbeDecoder);
 
-  void dump() { RootContext.dumpTree(); }
-
-private:
   using ProbeFrameStack = SmallVector<std::pair<StringRef, uint32_t>>;
   void trackInlineesOptimizedAway(MCPseudoProbeDecoder &ProbeDecoder,
                               MCDecodedPseudoProbeInlineTree &ProbeNode,
                               ProbeFrameStack &Context);
 
+  void dump() { RootContext.dumpTree(); }
+
+private:
   // Root node for context trie tree, node that this is a reverse context trie
   // with callee as parent and caller as child. This way we can traverse from
   // root to find the best/longest matching context if an exact match does not
   // exist. It gives us the best possible estimate for function's post-inline,
   // post-optimization byte size.
   ContextTrieNode RootContext;
 };
 
class ProfiledBinary {
   std::unordered_set<std::string> NameStrings;
 
   // A collection of functions to print disassembly for.
   StringSet<> DisassembleFunctionSet;
 
   // Pseudo probe decoder
   MCPseudoProbeDecoder ProbeDecoder;
 
+  // Function name to probe frame map for top-level outlined functions.
+  StringMap<MCDecodedPseudoProbeInlineTree *> TopLevelProbeFrameMap;
+
   bool UsePseudoProbes = false;
 
   bool UseFSDiscriminator = false;
 
   // Whether we need to symbolize all instructions to get function context size.
   bool TrackFuncContextSize = false;
 
   // Indicate if the base loading address is parsed from the mmap event or uses
public:
 
   Optional<SampleContextFrame> getInlineLeafFrameLoc(uint64_t Offset) {
     const auto &Stack = getFrameLocationStack(Offset);
     if (Stack.empty())
       return {};
     return Stack.back();
   }
 
+  void flushSymbolizer() { Symbolizer->flush(); }
+
   // Compare two addresses' inline context
   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);
   // Go through instructions among the given range and record its size for the
   // inline context.
   void computeInlinedContextSizeForRange(uint64_t StartOffset,
                                          uint64_t EndOffset);
 
+  void computeInlinedContextSizeForFunc(const BinaryFunction *Func);
+
   const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
 
   void getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
                                 SampleContextFrameVector &InlineContextStack,
                                 bool IncludeLeaf = false) const {
     SmallVector<MCPseduoProbeFrameLocation, 16> ProbeInlineContext;

llvm/tools/llvm-profgen/ProfiledBinary.cpp

if (!DebugBinaryPath.empty()) {
     loadSymbolsFromDWARF(*dyn_cast<ObjectFile>(DebugPath.getBinary()));
   } else {
     loadSymbolsFromDWARF(*dyn_cast<ObjectFile>(&ExeBinary));
   }
 
   // Disassemble the text sections.
   disassemble(Obj);
 
-  // Track size for optimized inlinees when probe is available
-  if (UsePseudoProbes && TrackFuncContextSize)
-    FuncSizeTracker.trackInlineesOptimizedAway(ProbeDecoder);
-
   // Use function start and return address to infer prolog and epilog
   ProEpilogTracker.inferPrologOffsets(StartOffset2FuncRangeMap);
   ProEpilogTracker.inferEpilogOffsets(RetOffsets);
 
   warnNoFuncEntry();
 
   // TODO: decode other sections.
 }
if (SectionName == ".pseudo_probe_desc") {
               reinterpret_cast<const uint8_t *>(Contents.data()),
               Contents.size()))
         exitWithError("Pseudo Probe decoder fail in .pseudo_probe section");
       // set UsePseudoProbes flag, used for PerfReader
       UsePseudoProbes = true;
     }
   }
 
+  // Build TopLevelProbeFrameMap to track size for optimized inlinees when probe
+  // is available
+  if (UsePseudoProbes && TrackFuncContextSize) {
+    for (const auto &Child : ProbeDecoder.getDummyInlineRoot().getChildren()) {
+      auto *Frame = Child.second.get();
+      StringRef FuncName =
+          ProbeDecoder.getFuncDescForGUID(Frame->Guid)->FuncName;
+      TopLevelProbeFrameMap[FuncName] = Frame;
+    }
+  }
+
   if (ShowPseudoProbe)
     ProbeDecoder.printGUID2FuncDescMap(outs());
 }
 
 void ProfiledBinary::setIsFuncEntry(uint64_t Offset, StringRef RangeSymName) {
   // Note that the start offset of each ELF section can be a non-function
   // symbol, we need to binary search for the start of a real function range.
   auto *FuncRange = findFuncRangeForOffset(Offset);
do {
     uint64_t Size = Offset2InstSizeMap[Offset];
 
     // Record instruction size for the corresponding context
     FuncSizeTracker.addInstructionForContext(SymbolizedCallStack, Size);
 
   } while (IP.advance() && IP.Address < RangeEnd);
 }
 
+void ProfiledBinary::computeInlinedContextSizeForFunc(
+    const BinaryFunction *Func) {
+  // Note that a function can be spilt into multiple ranges, so compute for all
+  // ranges of the function.
+  for (const auto &Range : Func->Ranges)
+    computeInlinedContextSizeForRange(Range.first, Range.second);
+
+  // Track optimized-away inlinee for probed binary. A function inlined and then
+  // optimized away should still have their probes left over in places.
+  if (usePseudoProbes()) {
+    auto I = TopLevelProbeFrameMap.find(Func->FuncName);
+    if (I != TopLevelProbeFrameMap.end()) {
+      BinarySizeContextTracker::ProbeFrameStack ProbeContext;
+      FuncSizeTracker.trackInlineesOptimizedAway(ProbeDecoder, *I->second,
+                                                 ProbeContext);
+    }
+  }
+}
+
 InstructionPointer::InstructionPointer(const ProfiledBinary *Binary,
                                        uint64_t Address, bool RoundToNext)
     : Binary(Binary), Address(Address) {
   Index = Binary->getIndexForAddr(Address);
   if (RoundToNext) {
     // we might get address which is not the code
     // it should round to the next valid address
     if (Index >= Binary->getCodeOffsetsSize())