Transforms/IPO/SampleContextTracker.h

0001 //===- Transforms/IPO/SampleContextTracker.h --------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 /// \file
0010 /// This file provides the interface for context-sensitive profile tracker used
0011 /// by CSSPGO.
0012 //
0013 //===----------------------------------------------------------------------===//
0014
0015 #ifndef LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H
0016 #define LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H
0017
0018 #include "llvm/ADT/StringRef.h"
0019 #include "llvm/ADT/iterator.h"
0020 #include "llvm/ProfileData/SampleProf.h"
0021 #include <map>
0022 #include <queue>
0023 #include <vector>
0024
0025 namespace llvm {
0026 class CallBase;
0027 class DILocation;
0028 class Function;
0029 class Instruction;
0030
0031 // Internal trie tree representation used for tracking context tree and sample
0032 // profiles. The path from root node to a given node represents the context of
0033 // that nodes' profile.
0034 class ContextTrieNode {
0035 public:
0036   ContextTrieNode(ContextTrieNode *Parent = nullptr,
0037                   FunctionId FName = FunctionId(),
0038                   FunctionSamples *FSamples = nullptr,
0039                   LineLocation CallLoc = {0, 0})
0040       : ParentContext(Parent), FuncName(FName), FuncSamples(FSamples),
0041         CallSiteLoc(CallLoc){};
0042   ContextTrieNode *getChildContext(const LineLocation &CallSite,
0043                                    FunctionId ChildName);
0044   ContextTrieNode *getHottestChildContext(const LineLocation &CallSite);
0045   ContextTrieNode *getOrCreateChildContext(const LineLocation &CallSite,
0046                                            FunctionId ChildName,
0047                                            bool AllowCreate = true);
0048   void removeChildContext(const LineLocation &CallSite, FunctionId ChildName);
0049   std::map<uint64_t, ContextTrieNode> &getAllChildContext();
0050   FunctionId getFuncName() const;
0051   FunctionSamples *getFunctionSamples() const;
0052   void setFunctionSamples(FunctionSamples *FSamples);
0053   std::optional<uint32_t> getFunctionSize() const;
0054   void addFunctionSize(uint32_t FSize);
0055   LineLocation getCallSiteLoc() const;
0056   ContextTrieNode *getParentContext() const;
0057   void setParentContext(ContextTrieNode *Parent);
0058   void setCallSiteLoc(const LineLocation &Loc);
0059   void dumpNode();
0060   void dumpTree();
0061
0062 private:
0063   // Map line+discriminator location to child context
0064   std::map<uint64_t, ContextTrieNode> AllChildContext;
0065
0066   // Link to parent context node
0067   ContextTrieNode *ParentContext;
0068
0069   // Function name for current context
0070   FunctionId FuncName;
0071
0072   // Function Samples for current context
0073   FunctionSamples *FuncSamples;
0074
0075   // Function size for current context
0076   std::optional<uint32_t> FuncSize;
0077
0078   // Callsite location in parent context
0079   LineLocation CallSiteLoc;
0080 };
0081
0082 // Profile tracker that manages profiles and its associated context. It
0083 // provides interfaces used by sample profile loader to query context profile or
0084 // base profile for given function or location; it also manages context tree
0085 // manipulation that is needed to accommodate inline decisions so we have
0086 // accurate post-inline profile for functions. Internally context profiles
0087 // are organized in a trie, with each node representing profile for specific
0088 // calling context and the context is identified by path from root to the node.
0089 class SampleContextTracker {
0090 public:
0091   using ContextSamplesTy = std::vector<FunctionSamples *>;
0092
0093   SampleContextTracker() = default;
0094   SampleContextTracker(SampleProfileMap &Profiles,
0095                        const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap);
0096   // Populate the FuncToCtxtProfiles map after the trie is built.
0097   void populateFuncToCtxtMap();
0098   // Query context profile for a specific callee with given name at a given
0099   // call-site. The full context is identified by location of call instruction.
0100   FunctionSamples *getCalleeContextSamplesFor(const CallBase &Inst,
0101                                               StringRef CalleeName);
0102   // Get samples for indirect call targets for call site at given location.
0103   std::vector<const FunctionSamples *>
0104   getIndirectCalleeContextSamplesFor(const DILocation *DIL);
0105   // Query context profile for a given location. The full context
0106   // is identified by input DILocation.
0107   FunctionSamples *getContextSamplesFor(const DILocation *DIL);
0108   // Query context profile for a given sample contxt of a function.
0109   FunctionSamples *getContextSamplesFor(const SampleContext &Context);
0110   // Get all context profile for given function.
0111   ContextSamplesTy &getAllContextSamplesFor(const Function &Func);
0112   ContextSamplesTy &getAllContextSamplesFor(StringRef Name);
0113   ContextTrieNode *getOrCreateContextPath(const SampleContext &Context,
0114                                           bool AllowCreate);
0115   // Query base profile for a given function. A base profile is a merged view
0116   // of all context profiles for contexts that are not inlined.
0117   FunctionSamples *getBaseSamplesFor(const Function &Func,
0118                                      bool MergeContext = true);
0119   // Query base profile for a given function by name.
0120   FunctionSamples *getBaseSamplesFor(FunctionId Name,
0121                                      bool MergeContext = true);
0122   // Retrieve the context trie node for given profile context
0123   ContextTrieNode *getContextFor(const SampleContext &Context);
0124   // Get real function name for a given trie node.
0125   StringRef getFuncNameFor(ContextTrieNode *Node) const;
0126   // Mark a context profile as inlined when function is inlined.
0127   // This makes sure that inlined context profile will be excluded in
0128   // function's base profile.
0129   void markContextSamplesInlined(const FunctionSamples *InlinedSamples);
0130   ContextTrieNode &getRootContext();
0131   void promoteMergeContextSamplesTree(const Instruction &Inst,
0132                                       FunctionId CalleeName);
0133
0134   // Create a merged conext-less profile map.
0135   void createContextLessProfileMap(SampleProfileMap &ContextLessProfiles);
0136   ContextTrieNode *
0137   getContextNodeForProfile(const FunctionSamples *FSamples) const {
0138     auto I = ProfileToNodeMap.find(FSamples);
0139     if (I == ProfileToNodeMap.end())
0140       return nullptr;
0141     return I->second;
0142   }
0143   HashKeyMap<std::unordered_map, FunctionId, ContextSamplesTy>
0144       &getFuncToCtxtProfiles() {
0145     return FuncToCtxtProfiles;
0146   }
0147
0148   class Iterator : public llvm::iterator_facade_base<
0149                        Iterator, std::forward_iterator_tag, ContextTrieNode *,
0150                        std::ptrdiff_t, ContextTrieNode **, ContextTrieNode *> {
0151     std::queue<ContextTrieNode *> NodeQueue;
0152
0153   public:
0154     explicit Iterator() = default;
0155     explicit Iterator(ContextTrieNode *Node) { NodeQueue.push(Node); }
0156     Iterator &operator++() {
0157       assert(!NodeQueue.empty() && "Iterator already at the end");
0158       ContextTrieNode *Node = NodeQueue.front();
0159       NodeQueue.pop();
0160       for (auto &It : Node->getAllChildContext())
0161         NodeQueue.push(&It.second);
0162       return *this;
0163     }
0164
0165     bool operator==(const Iterator &Other) const {
0166       if (NodeQueue.empty() && Other.NodeQueue.empty())
0167         return true;
0168       if (NodeQueue.empty() || Other.NodeQueue.empty())
0169         return false;
0170       return NodeQueue.front() == Other.NodeQueue.front();
0171     }
0172
0173     ContextTrieNode *operator*() const {
0174       assert(!NodeQueue.empty() && "Invalid access to end iterator");
0175       return NodeQueue.front();
0176     }
0177   };
0178
0179   Iterator begin() { return Iterator(&RootContext); }
0180   Iterator end() { return Iterator(); }
0181
0182 #ifndef NDEBUG
0183   // Get a context string from root to current node.
0184   std::string getContextString(const FunctionSamples &FSamples) const;
0185   std::string getContextString(ContextTrieNode *Node) const;
0186 #endif
0187   // Dump the internal context profile trie.
0188   void dump();
0189
0190 private:
0191   ContextTrieNode *getContextFor(const DILocation *DIL);
0192   ContextTrieNode *getCalleeContextFor(const DILocation *DIL,
0193                                        FunctionId CalleeName);
0194   ContextTrieNode *getTopLevelContextNode(FunctionId FName);
0195   ContextTrieNode &addTopLevelContextNode(FunctionId FName);
0196   ContextTrieNode &promoteMergeContextSamplesTree(ContextTrieNode &NodeToPromo);
0197   void mergeContextNode(ContextTrieNode &FromNode, ContextTrieNode &ToNode);
0198   ContextTrieNode &
0199   promoteMergeContextSamplesTree(ContextTrieNode &FromNode,
0200                                  ContextTrieNode &ToNodeParent);
0201   ContextTrieNode &moveContextSamples(ContextTrieNode &ToNodeParent,
0202                                       const LineLocation &CallSite,
0203                                       ContextTrieNode &&NodeToMove);
0204   void setContextNode(const FunctionSamples *FSample, ContextTrieNode *Node) {
0205     ProfileToNodeMap[FSample] = Node;
0206   }
0207   // Map from function name to context profiles (excluding base profile)
0208   HashKeyMap<std::unordered_map, FunctionId, ContextSamplesTy>
0209       FuncToCtxtProfiles;
0210
0211   // Map from current FunctionSample to the belonged context trie.
0212   std::unordered_map<const FunctionSamples *, ContextTrieNode *>
0213       ProfileToNodeMap;
0214
0215   // Map from function guid to real function names. Only used in md5 mode.
0216   const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap;
0217
0218   // Root node for context trie tree
0219   ContextTrieNode RootContext;
0220 };
0221
0222 } // end namespace llvm
0223 #endif // LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H