1 //===- SampleContextTracker.cpp - Context-sensitive Profile Tracker -------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the SampleContextTracker used by CSSPGO.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/Transforms/IPO/SampleContextTracker.h"
14 #include "llvm/ADT/StringMap.h"
15 #include "llvm/ADT/StringRef.h"
16 #include "llvm/IR/DebugInfoMetadata.h"
17 #include "llvm/IR/Instructions.h"
18 #include "llvm/ProfileData/SampleProf.h"
19 #include <map>
20 #include <queue>
21 #include <vector>
22 
23 using namespace llvm;
24 using namespace sampleprof;
25 
26 #define DEBUG_TYPE "sample-context-tracker"
27 
28 namespace llvm {
29 
30 ContextTrieNode *ContextTrieNode::getChildContext(const LineLocation &CallSite,
31                                                   StringRef CalleeName) {
32   if (CalleeName.empty())
33     return getHottestChildContext(CallSite);
34 
35   uint32_t Hash = nodeHash(CalleeName, CallSite);
36   auto It = AllChildContext.find(Hash);
37   if (It != AllChildContext.end())
38     return &It->second;
39   return nullptr;
40 }
41 
42 ContextTrieNode *
43 ContextTrieNode::getHottestChildContext(const LineLocation &CallSite) {
44   // CSFDO-TODO: This could be slow, change AllChildContext so we can
45   // do point look up for child node by call site alone.
46   // Retrieve the child node with max count for indirect call
47   ContextTrieNode *ChildNodeRet = nullptr;
48   uint64_t MaxCalleeSamples = 0;
49   for (auto &It : AllChildContext) {
50     ContextTrieNode &ChildNode = It.second;
51     if (ChildNode.CallSiteLoc != CallSite)
52       continue;
53     FunctionSamples *Samples = ChildNode.getFunctionSamples();
54     if (!Samples)
55       continue;
56     if (Samples->getTotalSamples() > MaxCalleeSamples) {
57       ChildNodeRet = &ChildNode;
58       MaxCalleeSamples = Samples->getTotalSamples();
59     }
60   }
61 
62   return ChildNodeRet;
63 }
64 
65 ContextTrieNode &ContextTrieNode::moveToChildContext(
66     const LineLocation &CallSite, ContextTrieNode &&NodeToMove,
67     StringRef ContextStrToRemove, bool DeleteNode) {
68   uint32_t Hash = nodeHash(NodeToMove.getFuncName(), CallSite);
69   assert(!AllChildContext.count(Hash) && "Node to remove must exist");
70   LineLocation OldCallSite = NodeToMove.CallSiteLoc;
71   ContextTrieNode &OldParentContext = *NodeToMove.getParentContext();
72   AllChildContext[Hash] = NodeToMove;
73   ContextTrieNode &NewNode = AllChildContext[Hash];
74   NewNode.CallSiteLoc = CallSite;
75 
76   // Walk through nodes in the moved the subtree, and update
77   // FunctionSamples' context as for the context promotion.
78   // We also need to set new parant link for all children.
79   std::queue<ContextTrieNode *> NodeToUpdate;
80   NewNode.setParentContext(this);
81   NodeToUpdate.push(&NewNode);
82 
83   while (!NodeToUpdate.empty()) {
84     ContextTrieNode *Node = NodeToUpdate.front();
85     NodeToUpdate.pop();
86     FunctionSamples *FSamples = Node->getFunctionSamples();
87 
88     if (FSamples) {
89       FSamples->getContext().promoteOnPath(ContextStrToRemove);
90       FSamples->getContext().setState(SyntheticContext);
91       LLVM_DEBUG(dbgs() << "  Context promoted to: " << FSamples->getContext()
92                         << "\n");
93     }
94 
95     for (auto &It : Node->getAllChildContext()) {
96       ContextTrieNode *ChildNode = &It.second;
97       ChildNode->setParentContext(Node);
98       NodeToUpdate.push(ChildNode);
99     }
100   }
101 
102   // Original context no longer needed, destroy if requested.
103   if (DeleteNode)
104     OldParentContext.removeChildContext(OldCallSite, NewNode.getFuncName());
105 
106   return NewNode;
107 }
108 
109 void ContextTrieNode::removeChildContext(const LineLocation &CallSite,
110                                          StringRef CalleeName) {
111   uint32_t Hash = nodeHash(CalleeName, CallSite);
112   // Note this essentially calls dtor and destroys that child context
113   AllChildContext.erase(Hash);
114 }
115 
116 std::map<uint32_t, ContextTrieNode> &ContextTrieNode::getAllChildContext() {
117   return AllChildContext;
118 }
119 
120 StringRef ContextTrieNode::getFuncName() const { return FuncName; }
121 
122 FunctionSamples *ContextTrieNode::getFunctionSamples() const {
123   return FuncSamples;
124 }
125 
126 void ContextTrieNode::setFunctionSamples(FunctionSamples *FSamples) {
127   FuncSamples = FSamples;
128 }
129 
130 LineLocation ContextTrieNode::getCallSiteLoc() const { return CallSiteLoc; }
131 
132 ContextTrieNode *ContextTrieNode::getParentContext() const {
133   return ParentContext;
134 }
135 
136 void ContextTrieNode::setParentContext(ContextTrieNode *Parent) {
137   ParentContext = Parent;
138 }
139 
140 void ContextTrieNode::dump() {
141   dbgs() << "Node: " << FuncName << "\n"
142          << "  Callsite: " << CallSiteLoc << "\n"
143          << "  Children:\n";
144 
145   for (auto &It : AllChildContext) {
146     dbgs() << "    Node: " << It.second.getFuncName() << "\n";
147   }
148 }
149 
150 uint32_t ContextTrieNode::nodeHash(StringRef ChildName,
151                                    const LineLocation &Callsite) {
152   // We still use child's name for child hash, this is
153   // because for children of root node, we don't have
154   // different line/discriminator, and we'll rely on name
155   // to differentiate children.
156   uint32_t NameHash = std::hash<std::string>{}(ChildName.str());
157   uint32_t LocId = (Callsite.LineOffset << 16) | Callsite.Discriminator;
158   return NameHash + (LocId << 5) + LocId;
159 }
160 
161 ContextTrieNode *ContextTrieNode::getOrCreateChildContext(
162     const LineLocation &CallSite, StringRef CalleeName, bool AllowCreate) {
163   uint32_t Hash = nodeHash(CalleeName, CallSite);
164   auto It = AllChildContext.find(Hash);
165   if (It != AllChildContext.end()) {
166     assert(It->second.getFuncName() == CalleeName &&
167            "Hash collision for child context node");
168     return &It->second;
169   }
170 
171   if (!AllowCreate)
172     return nullptr;
173 
174   AllChildContext[Hash] = ContextTrieNode(this, CalleeName, nullptr, CallSite);
175   return &AllChildContext[Hash];
176 }
177 
178 // Profiler tracker than manages profiles and its associated context
179 SampleContextTracker::SampleContextTracker(
180     StringMap<FunctionSamples> &Profiles) {
181   for (auto &FuncSample : Profiles) {
182     FunctionSamples *FSamples = &FuncSample.second;
183     SampleContext Context(FuncSample.first(), RawContext);
184     LLVM_DEBUG(dbgs() << "Tracking Context for function: " << Context << "\n");
185     if (!Context.isBaseContext())
186       FuncToCtxtProfiles[Context.getNameWithoutContext()].push_back(FSamples);
187     ContextTrieNode *NewNode = getOrCreateContextPath(Context, true);
188     assert(!NewNode->getFunctionSamples() &&
189            "New node can't have sample profile");
190     NewNode->setFunctionSamples(FSamples);
191   }
192 }
193 
194 FunctionSamples *
195 SampleContextTracker::getCalleeContextSamplesFor(const CallBase &Inst,
196                                                  StringRef CalleeName) {
197   LLVM_DEBUG(dbgs() << "Getting callee context for instr: " << Inst << "\n");
198   DILocation *DIL = Inst.getDebugLoc();
199   if (!DIL)
200     return nullptr;
201 
202   CalleeName = FunctionSamples::getCanonicalFnName(CalleeName);
203 
204   // For indirect call, CalleeName will be empty, in which case the context
205   // profile for callee with largest total samples will be returned.
206   ContextTrieNode *CalleeContext = getCalleeContextFor(DIL, CalleeName);
207   if (CalleeContext) {
208     FunctionSamples *FSamples = CalleeContext->getFunctionSamples();
209     LLVM_DEBUG(if (FSamples) {
210       dbgs() << "  Callee context found: " << FSamples->getContext() << "\n";
211     });
212     return FSamples;
213   }
214 
215   return nullptr;
216 }
217 
218 std::vector<const FunctionSamples *>
219 SampleContextTracker::getIndirectCalleeContextSamplesFor(
220     const DILocation *DIL) {
221   std::vector<const FunctionSamples *> R;
222   if (!DIL)
223     return R;
224 
225   ContextTrieNode *CallerNode = getContextFor(DIL);
226   LineLocation CallSite = FunctionSamples::getCallSiteIdentifier(DIL);
227   for (auto &It : CallerNode->getAllChildContext()) {
228     ContextTrieNode &ChildNode = It.second;
229     if (ChildNode.getCallSiteLoc() != CallSite)
230       continue;
231     if (FunctionSamples *CalleeSamples = ChildNode.getFunctionSamples())
232       R.push_back(CalleeSamples);
233   }
234 
235   return R;
236 }
237 
238 FunctionSamples *
239 SampleContextTracker::getContextSamplesFor(const DILocation *DIL) {
240   assert(DIL && "Expect non-null location");
241 
242   ContextTrieNode *ContextNode = getContextFor(DIL);
243   if (!ContextNode)
244     return nullptr;
245 
246   // We may have inlined callees during pre-LTO compilation, in which case
247   // we need to rely on the inline stack from !dbg to mark context profile
248   // as inlined, instead of `MarkContextSamplesInlined` during inlining.
249   // Sample profile loader walks through all instructions to get profile,
250   // which calls this function. So once that is done, all previously inlined
251   // context profile should be marked properly.
252   FunctionSamples *Samples = ContextNode->getFunctionSamples();
253   if (Samples && ContextNode->getParentContext() != &RootContext)
254     Samples->getContext().setState(InlinedContext);
255 
256   return Samples;
257 }
258 
259 FunctionSamples *
260 SampleContextTracker::getContextSamplesFor(const SampleContext &Context) {
261   ContextTrieNode *Node = getContextFor(Context);
262   if (!Node)
263     return nullptr;
264 
265   return Node->getFunctionSamples();
266 }
267 
268 SampleContextTracker::ContextSamplesTy &
269 SampleContextTracker::getAllContextSamplesFor(const Function &Func) {
270   StringRef CanonName = FunctionSamples::getCanonicalFnName(Func);
271   return FuncToCtxtProfiles[CanonName];
272 }
273 
274 SampleContextTracker::ContextSamplesTy &
275 SampleContextTracker::getAllContextSamplesFor(StringRef Name) {
276   return FuncToCtxtProfiles[Name];
277 }
278 
279 FunctionSamples *SampleContextTracker::getBaseSamplesFor(const Function &Func,
280                                                          bool MergeContext) {
281   StringRef CanonName = FunctionSamples::getCanonicalFnName(Func);
282   return getBaseSamplesFor(CanonName, MergeContext);
283 }
284 
285 FunctionSamples *SampleContextTracker::getBaseSamplesFor(StringRef Name,
286                                                          bool MergeContext) {
287   LLVM_DEBUG(dbgs() << "Getting base profile for function: " << Name << "\n");
288   // Base profile is top-level node (child of root node), so try to retrieve
289   // existing top-level node for given function first. If it exists, it could be
290   // that we've merged base profile before, or there's actually context-less
291   // profile from the input (e.g. due to unreliable stack walking).
292   ContextTrieNode *Node = getTopLevelContextNode(Name);
293   if (MergeContext) {
294     LLVM_DEBUG(dbgs() << "  Merging context profile into base profile: " << Name
295                       << "\n");
296 
297     // We have profile for function under different contexts,
298     // create synthetic base profile and merge context profiles
299     // into base profile.
300     for (auto *CSamples : FuncToCtxtProfiles[Name]) {
301       SampleContext &Context = CSamples->getContext();
302       ContextTrieNode *FromNode = getContextFor(Context);
303       if (FromNode == Node)
304         continue;
305 
306       // Skip inlined context profile and also don't re-merge any context
307       if (Context.hasState(InlinedContext) || Context.hasState(MergedContext))
308         continue;
309 
310       ContextTrieNode &ToNode = promoteMergeContextSamplesTree(*FromNode);
311       assert((!Node || Node == &ToNode) && "Expect only one base profile");
312       Node = &ToNode;
313     }
314   }
315 
316   // Still no profile even after merge/promotion (if allowed)
317   if (!Node)
318     return nullptr;
319 
320   return Node->getFunctionSamples();
321 }
322 
323 void SampleContextTracker::markContextSamplesInlined(
324     const FunctionSamples *InlinedSamples) {
325   assert(InlinedSamples && "Expect non-null inlined samples");
326   LLVM_DEBUG(dbgs() << "Marking context profile as inlined: "
327                     << InlinedSamples->getContext() << "\n");
328   InlinedSamples->getContext().setState(InlinedContext);
329 }
330 
331 ContextTrieNode &SampleContextTracker::getRootContext() { return RootContext; }
332 
333 void SampleContextTracker::promoteMergeContextSamplesTree(
334     const Instruction &Inst, StringRef CalleeName) {
335   LLVM_DEBUG(dbgs() << "Promoting and merging context tree for instr: \n"
336                     << Inst << "\n");
337   // Get the caller context for the call instruction, we don't use callee
338   // name from call because there can be context from indirect calls too.
339   DILocation *DIL = Inst.getDebugLoc();
340   ContextTrieNode *CallerNode = getContextFor(DIL);
341   if (!CallerNode)
342     return;
343 
344   // Get the context that needs to be promoted
345   LineLocation CallSite = FunctionSamples::getCallSiteIdentifier(DIL);
346   // For indirect call, CalleeName will be empty, in which case we need to
347   // promote all non-inlined child context profiles.
348   if (CalleeName.empty()) {
349     for (auto &It : CallerNode->getAllChildContext()) {
350       ContextTrieNode *NodeToPromo = &It.second;
351       if (CallSite != NodeToPromo->getCallSiteLoc())
352         continue;
353       FunctionSamples *FromSamples = NodeToPromo->getFunctionSamples();
354       if (FromSamples && FromSamples->getContext().hasState(InlinedContext))
355         continue;
356       promoteMergeContextSamplesTree(*NodeToPromo);
357     }
358     return;
359   }
360 
361   // Get the context for the given callee that needs to be promoted
362   ContextTrieNode *NodeToPromo =
363       CallerNode->getChildContext(CallSite, CalleeName);
364   if (!NodeToPromo)
365     return;
366 
367   promoteMergeContextSamplesTree(*NodeToPromo);
368 }
369 
370 ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
371     ContextTrieNode &NodeToPromo) {
372   // Promote the input node to be directly under root. This can happen
373   // when we decided to not inline a function under context represented
374   // by the input node. The promote and merge is then needed to reflect
375   // the context profile in the base (context-less) profile.
376   FunctionSamples *FromSamples = NodeToPromo.getFunctionSamples();
377   assert(FromSamples && "Shouldn't promote a context without profile");
378   LLVM_DEBUG(dbgs() << "  Found context tree root to promote: "
379                     << FromSamples->getContext() << "\n");
380 
381   assert(!FromSamples->getContext().hasState(InlinedContext) &&
382          "Shouldn't promote inlined context profile");
383   StringRef ContextStrToRemove = FromSamples->getContext().getCallingContext();
384   return promoteMergeContextSamplesTree(NodeToPromo, RootContext,
385                                         ContextStrToRemove);
386 }
387 
388 void SampleContextTracker::dump() {
389   dbgs() << "Context Profile Tree:\n";
390   std::queue<ContextTrieNode *> NodeQueue;
391   NodeQueue.push(&RootContext);
392 
393   while (!NodeQueue.empty()) {
394     ContextTrieNode *Node = NodeQueue.front();
395     NodeQueue.pop();
396     Node->dump();
397 
398     for (auto &It : Node->getAllChildContext()) {
399       ContextTrieNode *ChildNode = &It.second;
400       NodeQueue.push(ChildNode);
401     }
402   }
403 }
404 
405 ContextTrieNode *
406 SampleContextTracker::getContextFor(const SampleContext &Context) {
407   return getOrCreateContextPath(Context, false);
408 }
409 
410 ContextTrieNode *
411 SampleContextTracker::getCalleeContextFor(const DILocation *DIL,
412                                           StringRef CalleeName) {
413   assert(DIL && "Expect non-null location");
414 
415   ContextTrieNode *CallContext = getContextFor(DIL);
416   if (!CallContext)
417     return nullptr;
418 
419   // When CalleeName is empty, the child context profile with max
420   // total samples will be returned.
421   return CallContext->getChildContext(
422       FunctionSamples::getCallSiteIdentifier(DIL), CalleeName);
423 }
424 
425 ContextTrieNode *SampleContextTracker::getContextFor(const DILocation *DIL) {
426   assert(DIL && "Expect non-null location");
427   SmallVector<std::pair<LineLocation, StringRef>, 10> S;
428 
429   // Use C++ linkage name if possible.
430   const DILocation *PrevDIL = DIL;
431   for (DIL = DIL->getInlinedAt(); DIL; DIL = DIL->getInlinedAt()) {
432     StringRef Name = PrevDIL->getScope()->getSubprogram()->getLinkageName();
433     if (Name.empty())
434       Name = PrevDIL->getScope()->getSubprogram()->getName();
435     S.push_back(
436         std::make_pair(FunctionSamples::getCallSiteIdentifier(DIL), Name));
437     PrevDIL = DIL;
438   }
439 
440   // Push root node, note that root node like main may only
441   // a name, but not linkage name.
442   StringRef RootName = PrevDIL->getScope()->getSubprogram()->getLinkageName();
443   if (RootName.empty())
444     RootName = PrevDIL->getScope()->getSubprogram()->getName();
445   S.push_back(std::make_pair(LineLocation(0, 0), RootName));
446 
447   ContextTrieNode *ContextNode = &RootContext;
448   int I = S.size();
449   while (--I >= 0 && ContextNode) {
450     LineLocation &CallSite = S[I].first;
451     StringRef &CalleeName = S[I].second;
452     ContextNode = ContextNode->getChildContext(CallSite, CalleeName);
453   }
454 
455   if (I < 0)
456     return ContextNode;
457 
458   return nullptr;
459 }
460 
461 ContextTrieNode *
462 SampleContextTracker::getOrCreateContextPath(const SampleContext &Context,
463                                              bool AllowCreate) {
464   ContextTrieNode *ContextNode = &RootContext;
465   StringRef ContextRemain = Context;
466   StringRef ChildContext;
467   StringRef CalleeName;
468   LineLocation CallSiteLoc(0, 0);
469 
470   while (ContextNode && !ContextRemain.empty()) {
471     auto ContextSplit = SampleContext::splitContextString(ContextRemain);
472     ChildContext = ContextSplit.first;
473     ContextRemain = ContextSplit.second;
474     LineLocation NextCallSiteLoc(0, 0);
475     SampleContext::decodeContextString(ChildContext, CalleeName,
476                                        NextCallSiteLoc);
477 
478     // Create child node at parent line/disc location
479     if (AllowCreate) {
480       ContextNode =
481           ContextNode->getOrCreateChildContext(CallSiteLoc, CalleeName);
482     } else {
483       ContextNode = ContextNode->getChildContext(CallSiteLoc, CalleeName);
484     }
485     CallSiteLoc = NextCallSiteLoc;
486   }
487 
488   assert((!AllowCreate || ContextNode) &&
489          "Node must exist if creation is allowed");
490   return ContextNode;
491 }
492 
493 ContextTrieNode *SampleContextTracker::getTopLevelContextNode(StringRef FName) {
494   assert(!FName.empty() && "Top level node query must provide valid name");
495   return RootContext.getChildContext(LineLocation(0, 0), FName);
496 }
497 
498 ContextTrieNode &SampleContextTracker::addTopLevelContextNode(StringRef FName) {
499   assert(!getTopLevelContextNode(FName) && "Node to add must not exist");
500   return *RootContext.getOrCreateChildContext(LineLocation(0, 0), FName);
501 }
502 
503 void SampleContextTracker::mergeContextNode(ContextTrieNode &FromNode,
504                                             ContextTrieNode &ToNode,
505                                             StringRef ContextStrToRemove) {
506   FunctionSamples *FromSamples = FromNode.getFunctionSamples();
507   FunctionSamples *ToSamples = ToNode.getFunctionSamples();
508   if (FromSamples && ToSamples) {
509     // Merge/duplicate FromSamples into ToSamples
510     ToSamples->merge(*FromSamples);
511     ToSamples->getContext().setState(SyntheticContext);
512     FromSamples->getContext().setState(MergedContext);
513   } else if (FromSamples) {
514     // Transfer FromSamples from FromNode to ToNode
515     ToNode.setFunctionSamples(FromSamples);
516     FromSamples->getContext().setState(SyntheticContext);
517     FromSamples->getContext().promoteOnPath(ContextStrToRemove);
518     FromNode.setFunctionSamples(nullptr);
519   }
520 }
521 
522 ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
523     ContextTrieNode &FromNode, ContextTrieNode &ToNodeParent,
524     StringRef ContextStrToRemove) {
525   assert(!ContextStrToRemove.empty() && "Context to remove can't be empty");
526 
527   // Ignore call site location if destination is top level under root
528   LineLocation NewCallSiteLoc = LineLocation(0, 0);
529   LineLocation OldCallSiteLoc = FromNode.getCallSiteLoc();
530   ContextTrieNode &FromNodeParent = *FromNode.getParentContext();
531   ContextTrieNode *ToNode = nullptr;
532   bool MoveToRoot = (&ToNodeParent == &RootContext);
533   if (!MoveToRoot) {
534     NewCallSiteLoc = OldCallSiteLoc;
535   }
536 
537   // Locate destination node, create/move if not existing
538   ToNode = ToNodeParent.getChildContext(NewCallSiteLoc, FromNode.getFuncName());
539   if (!ToNode) {
540     // Do not delete node to move from its parent here because
541     // caller is iterating over children of that parent node.
542     ToNode = &ToNodeParent.moveToChildContext(
543         NewCallSiteLoc, std::move(FromNode), ContextStrToRemove, false);
544   } else {
545     // Destination node exists, merge samples for the context tree
546     mergeContextNode(FromNode, *ToNode, ContextStrToRemove);
547     LLVM_DEBUG({
548       if (ToNode->getFunctionSamples())
549         dbgs() << "  Context promoted and merged to: "
550                << ToNode->getFunctionSamples()->getContext() << "\n";
551     });
552 
553     // Recursively promote and merge children
554     for (auto &It : FromNode.getAllChildContext()) {
555       ContextTrieNode &FromChildNode = It.second;
556       promoteMergeContextSamplesTree(FromChildNode, *ToNode,
557                                      ContextStrToRemove);
558     }
559 
560     // Remove children once they're all merged
561     FromNode.getAllChildContext().clear();
562   }
563 
564   // For root of subtree, remove itself from old parent too
565   if (MoveToRoot)
566     FromNodeParent.removeChildContext(OldCallSiteLoc, ToNode->getFuncName());
567 
568   return *ToNode;
569 }
570 } // namespace llvm
571