[CSSPGO] Use probe inline tree to track zero size fully optimized context for pre-inliner

[wenlei]

Summary

This is a follow up diff for BinarySizeContextTracker to track zero size for fully optimized inlinee. When an inlinee is fully optimized away, we won't be able to get its size through symbolizing instructions, hence we will treat the corresponding context size as unknown. However by traversing the inlined probe forest, we know what're original inlinees regardless of optimization. If a context show up in inlined probes, but not during symbolization, we know that it's fully optimized away hence its size is zero instead of unknown. It should provide more accurate size cost estimation for pre-inliner to make better inline decisions in llvm-profgen.

[hoy]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

84

Use Optional for this field?

[wenlei]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

84

Sounds good, changed.

[wenlei]

Address comments

[wenlei]

update

[hoy]

lgtm

[wmi]

LGTM.

llvm/tools/llvm-profgen/ProfiledBinary.h

123

trackOptimizedInlinees --> trackInlineesOptimizedAway?

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.h

123

Sounds good, renamed.

[wenlei]

Update

llvm/include/llvm/MC/MCPseudoProbe.h

DerivedProbeInlineTreeType *getOrAddNode(const InlineSite &Site) {
     Ret.first->second->Parent = this;
     return Ret.first->second.get();
   };
 };
 
 // A Tri-tree based data structure to group probes by inline stack.
 // A tree is allocated for a standalone .text section. A fake
 // instance is created as the root of a tree.
-// A real instance of this class is created for each function, either an
-// unlined function that has code in .text section or an inlined function.
-
+// A real instance of this class is created for each function, either a
+// not inlined function that has code in .text section or an inlined function.
 class MCPseudoProbeInlineTree
     : public MCPseudoProbeInlineTreeBase<MCPseudoProbe,
                                          MCPseudoProbeInlineTree> {
 public:
   MCPseudoProbeInlineTree() = default;
   MCPseudoProbeInlineTree(uint64_t Guid) { this->Guid = Guid; }
   MCPseudoProbeInlineTree(const InlineSite &Site) {
     this->Guid = std::get<0>(Site);

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

 
 // Internal trie tree representation used for tracking context tree and sample
 // profiles. The path from root node to a given node represents the context of
 // that nodes' profile.
 class ContextTrieNode {
 public:
   ContextTrieNode(ContextTrieNode *Parent = nullptr,
                   StringRef FName = StringRef(),
-                  FunctionSamples *FSamples = nullptr, uint32_t FSize = 0,
+                  FunctionSamples *FSamples = nullptr,
                   LineLocation CallLoc = {0, 0})
       : ParentContext(Parent), FuncName(FName), FuncSamples(FSamples),
-        FuncSize(FSize), CallSiteLoc(CallLoc){};
+        CallSiteLoc(CallLoc){};
   ContextTrieNode *getChildContext(const LineLocation &CallSite,
                                    StringRef ChildName);
   ContextTrieNode *getHottestChildContext(const LineLocation &CallSite);
   ContextTrieNode *getOrCreateChildContext(const LineLocation &CallSite,
                                            StringRef ChildName,
                                            bool AllowCreate = true);
 
   ContextTrieNode &moveToChildContext(const LineLocation &CallSite,
                                       ContextTrieNode &&NodeToMove,
                                       StringRef ContextStrToRemove,
                                       bool DeleteNode = true);
   void removeChildContext(const LineLocation &CallSite, StringRef ChildName);
   std::map<uint32_t, ContextTrieNode> &getAllChildContext();
   StringRef getFuncName() const;
   FunctionSamples *getFunctionSamples() const;
   void setFunctionSamples(FunctionSamples *FSamples);
-  uint32_t getFunctionSize() const;
-  void setFunctionSize(uint32_t FSize);
+  Optional<uint32_t> getFunctionSize() const;
+  void addFunctionSize(uint32_t FSize);
   LineLocation getCallSiteLoc() const;
   ContextTrieNode *getParentContext() const;
   void setParentContext(ContextTrieNode *Parent);
   void dumpNode();
   void dumpTree();
 
 private:
   static uint32_t nodeHash(StringRef ChildName, const LineLocation &Callsite);
 
   // Map line+discriminator location to child context
   std::map<uint32_t, ContextTrieNode> AllChildContext;
 
   // Link to parent context node
   ContextTrieNode *ParentContext;
 
   // Function name for current context
   StringRef FuncName;
 
   // Function Samples for current context
   FunctionSamples *FuncSamples;
 
   // Function size for current context
-  uint32_t FuncSize;
+  Optional<uint32_t> FuncSize;
 
   // Callsite location in parent context
   LineLocation CallSiteLoc;
 };
 
 // Profile tracker that manages profiles and its associated context. It
 // provides interfaces used by sample profile loader to query context profile or
 // base profile for given function or location; it also manages context tree

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

 FunctionSamples *ContextTrieNode::getFunctionSamples() const {
   return FuncSamples;
 }
 
 void ContextTrieNode::setFunctionSamples(FunctionSamples *FSamples) {
   FuncSamples = FSamples;
 }
 
-uint32_t ContextTrieNode::getFunctionSize() const { return FuncSize; }
+Optional<uint32_t> ContextTrieNode::getFunctionSize() const { return FuncSize; }
 
-void ContextTrieNode::setFunctionSize(uint32_t FSize) { FuncSize = FSize; }
+void ContextTrieNode::addFunctionSize(uint32_t FSize) {
+  if (!FuncSize.hasValue())
+    FuncSize = 0;
+
+  FuncSize = FuncSize.getValue() + FSize;
+}
 
 LineLocation ContextTrieNode::getCallSiteLoc() const { return CallSiteLoc; }
 
 ContextTrieNode *ContextTrieNode::getParentContext() const {
   return ParentContext;
 }
 
 void ContextTrieNode::setParentContext(ContextTrieNode *Parent) {
if (It != AllChildContext.end()) {
     assert(It->second.getFuncName() == CalleeName &&
            "Hash collision for child context node");
     return &It->second;
   }
 
   if (!AllowCreate)
     return nullptr;
 
-  AllChildContext[Hash] =
-      ContextTrieNode(this, CalleeName, nullptr, 0, CallSite);
+  AllChildContext[Hash] = ContextTrieNode(this, CalleeName, nullptr, CallSite);
   return &AllChildContext[Hash];
 }
 
 // Profiler tracker than manages profiles and its associated context
 SampleContextTracker::SampleContextTracker(
     StringMap<FunctionSamples> &Profiles) {
   for (auto &FuncSample : Profiles) {
     FunctionSamples *FSamples = &FuncSample.second;

llvm/test/tools/llvm-profgen/cs-preinline-cost.test

 ; REQUIRES: asserts
 ; Test default using size of profile as a proxy
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/cs-preinline-cost.perfscript --binary=%S/Inputs/cs-preinline-cost.perfbin --csspgo-preinliner --debug-only=cs-preinliner --output=/dev/null 2>&1 | FileCheck %s --check-prefix=CHECK-DEFAULT
 
 ; Test use-context-cost-for-preinliner using inlinee's byte size as context-sensitive inline cost
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/cs-preinline-cost.perfscript --binary=%S/Inputs/cs-preinline-cost.perfbin --csspgo-preinliner --debug-only=cs-preinliner --use-context-cost-for-preinliner --output=/dev/null 2>&1 | FileCheck %s --check-prefix=CHECK-CSCOST
 
 CHECK-DEFAULT:      Process main for context-sensitive pre-inlining (pre-inline size: 9, size limit: 108)
 CHECK-DEFAULT-NEXT:   Inlined context profile for: main:9 @ _Z3fooi (callee size: 2, call count:545)
 CHECK-DEFAULT-NEXT:   Inlined context profile for: main:7 @ _Z3fooi (callee size: 14, call count:545)
 CHECK-DEFAULT-NEXT:   Inlined context profile for: main:8 @ _Z3fooi (callee size: 4, call count:544)
 
 CHECK-CSCOST:      Process main for context-sensitive pre-inlining (pre-inline size: 69, size limit: 828)
-CHECK-CSCOST-NEXT:   Inlined context profile for: main:9 @ _Z3fooi (callee size: 264, call count:545)
+; This inlinee is fully optimized away, make sure we have the correct zero size for that context even if the size is
+; not available through symbolization.
+CHECK-CSCOST-NEXT:   Inlined context profile for: main:9 @ _Z3fooi (callee size: 0, call count:545)
 CHECK-CSCOST-NEXT:   Inlined context profile for: main:7 @ _Z3fooi (callee size: 279, call count:545)
 CHECK-CSCOST-NEXT:   Inlined context profile for: main:8 @ _Z3fooi (callee size: 44, call count:544)
 
 ; binary is built with the source below using the following command line:
 ;   clang -O3 -g -fpseudo-probe-for-profiling -fexperimental-new-pass-manager test.cpp
 ;
 ;#include <stdio.h>
 ;

llvm/tools/llvm-profgen/ProfiledBinary.h

public:
   void addInstructionForContext(const FrameLocationStack &Context,
                                 uint32_t InstrSize);
 
   // Get function size with a specific context. When there's no exact match
   // 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);
+
   // 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;
 };
 

llvm/tools/llvm-profgen/ProfiledBinary.cpp

void BinarySizeContextTracker::addInstructionForContext(
   bool IsLeaf = true;
   for (const auto &Callsite : reverse(Context)) {
     StringRef CallerName = Callsite.first;
     LineLocation CallsiteLoc = IsLeaf ? LineLocation(0, 0) : Callsite.second;
     CurNode = CurNode->getOrCreateChildContext(CallsiteLoc, CallerName);
     IsLeaf = false;
   }
 
-  CurNode->setFunctionSize(CurNode->getFunctionSize() + InstrSize);
+  CurNode->addFunctionSize(InstrSize);
 }
 
 uint32_t
 BinarySizeContextTracker::getFuncSizeForContext(const SampleContext &Context) {
   ContextTrieNode *CurrNode = &RootContext;
   ContextTrieNode *PrevNode = nullptr;
   StringRef ContextRemain = Context;
   StringRef ChildContext;
   StringRef CallerName;
-  uint32_t Size = 0;
+  Optional<uint32_t> Size;
 
   // Start from top-level context-less function, travese down the reverse
   // context trie to find the best/longest match for given context, then
   // retrieve the size.
   while (CurrNode && !ContextRemain.empty()) {
     // rsplit so we process from leaf function to callers (added to context).
     auto ContextSplit = SampleContext::rsplitContextString(ContextRemain);
     ChildContext = ContextSplit.second;
     ContextRemain = ContextSplit.first;
     LineLocation CallSiteLoc(0, 0);
     SampleContext::decodeContextString(ChildContext, CallerName, CallSiteLoc);
     PrevNode = CurrNode;
     CurrNode = CurrNode->getChildContext(CallSiteLoc, CallerName);
-    if (CurrNode && CurrNode->getFunctionSize())
-      Size = CurrNode->getFunctionSize();
+    if (CurrNode && CurrNode->getFunctionSize().hasValue())
+      Size = CurrNode->getFunctionSize().getValue();
   }
 
   // If we traversed all nodes along the path of the context and haven't
   // found a size yet, pivot to look for size from sibling nodes, i.e size
   // of inlinee under different context.
-  if (!Size) {
+  if (!Size.hasValue()) {
     if (!CurrNode)
       CurrNode = PrevNode;
-    while (!Size && CurrNode) {
+    while (!Size.hasValue() && CurrNode &&
+           !CurrNode->getAllChildContext().empty()) {
       CurrNode = &CurrNode->getAllChildContext().begin()->second;
-      Size = CurrNode->getFunctionSize();
+      if (CurrNode->getFunctionSize().hasValue())
+        Size = CurrNode->getFunctionSize().getValue();
     }
   }
 
-  return Size;
+  assert(Size.hasValue() && "We should at least find one context size.");
+  return Size.getValue();
+}
+
+void BinarySizeContextTracker::trackInlineesOptimizedAway(
+    MCPseudoProbeDecoder &ProbeDecoder) {
+  ProbeFrameStack ProbeContext;
+  for (const auto &Child : ProbeDecoder.getDummyInlineRoot().getChildren())
+    trackInlineesOptimizedAway(ProbeDecoder, *Child.second.get(), ProbeContext);
+}
+
+void BinarySizeContextTracker::trackInlineesOptimizedAway(
+    MCPseudoProbeDecoder &ProbeDecoder,
+    MCDecodedPseudoProbeInlineTree &ProbeNode, ProbeFrameStack &ProbeContext) {
+  StringRef FuncName =
+      ProbeDecoder.getFuncDescForGUID(ProbeNode.Guid)->FuncName;
+  ProbeContext.emplace_back(FuncName, 0);
+
+  // This ProbeContext has a probe, so it has code before inlining and
+  // optimization. Make sure we mark its size as known.
+  if (!ProbeNode.getProbes().empty()) {
+    ContextTrieNode *SizeContext = &RootContext;
+    for (auto &ProbeFrame : reverse(ProbeContext)) {
+      StringRef CallerName = ProbeFrame.first;
+      LineLocation CallsiteLoc(ProbeFrame.second, 0);
+      SizeContext =
+          SizeContext->getOrCreateChildContext(CallsiteLoc, CallerName);
+    }
+    // Add 0 size to make known.
+    SizeContext->addFunctionSize(0);
+  }
+
+  // DFS down the probe inline tree
+  for (const auto &ChildNode : ProbeNode.getChildren()) {
+    InlineSite Location = ChildNode.first;
+    ProbeContext.back().second = std::get<1>(Location);
+    trackInlineesOptimizedAway(ProbeDecoder, *ChildNode.second.get(), ProbeContext);
+  }
+
+  ProbeContext.pop_back();
 }
 
 void ProfiledBinary::load() {
   // Attempt to open the binary.
   OwningBinary<Binary> OBinary = unwrapOrError(createBinary(Path), Path);
   Binary &Binary = *OBinary.getBinary();
 
   auto *Obj = dyn_cast<ELFObjectFileBase>(&Binary);
void ProfiledBinary::load() {
   setPreferredTextSegmentAddresses(Obj);
 
   // Decode pseudo probe related section
   decodePseudoProbe(Obj);
 
   // 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(FuncStartAddrMap);
   ProEpilogTracker.inferEpilogOffsets(RetAddrs);
 
   // TODO: decode other sections.
 }
 
 bool ProfiledBinary::inlineContextEqual(uint64_t Address1,