[llvm-profdata] Use StringRef for CallTargetMap

[huangjd]

Summary

Use StringRef as key to CallTargetMap. Previously CallTargetMap uses a StringMap, which actually stores a copy of the string. All strings put in the map is found to be backed by the profile data so a StringRef can be safely used.

This refactoring is a prerequisite to implement phase 2 of MD5 refactoring, where StringRef can represent an MD5 value directly, instead of casting it to a string first.

[hoy]

Interestingly we've seen regression on memory usage with similar change. The reason is that an empty DenseMap preallocates a chunk of memory that turns out more expansive than a single string for our workload. cc @marksantaniello @wenlei

[marksantaniello]

Yeah I tested pretty much the same thing, and it was a big memory regression. The problem is that, upon first insert, the DenseMap grows to 64 entries. Even if I use a const char* key, the KV-pair is 16B, so 64 entries is 1K.

[huangjd]

Holding on to this patch. Currently I am also exploring writing a specialized container for all the map types used by SampleProfile, given the fact that in production profiles maps typically have 0 or 1 elements (~90%), and map allocation is a huge bottleneck (~25% full profile load time), having a specialized map structure is beneficial

[marksantaniello]

FWIW I also tried SmallDenseMap with various inline sizes. It still seems to be the case that, once you exceed the in-situ space, you immediately grow to 64 entries.

[huangjd]

(temporary) use std::map in place of StringMap

[marksantaniello]

Another thing worth considering is just to use const char* as the map key. I believe it's safe in this case, and smaller than StringRef which is a pointer and a length.

llvm/include/llvm/ProfileData/SampleProf.h

bool operator()(const CallTarget &LHS, const CallTarget &RHS) const {
       if (LHS.second != RHS.second)
         return LHS.second > RHS.second;
 
       return LHS.first < RHS.first;
     }
   };
 
   using SortedCallTargetSet = std::set<CallTarget, CallTargetComparator>;
-  using CallTargetMap = StringMap<uint64_t>;
+  using CallTargetMap = std::map<StringRef, uint64_t>;
   SampleRecord() = default;
 
   /// Increment the number of samples for this record by \p S.
   /// Optionally scale sample count \p S by \p Weight.
   ///
   /// Sample counts accumulate using saturating arithmetic, to avoid wrapping
   /// around unsigned integers.
   sampleprof_error addSamples(uint64_t S, uint64_t Weight = 1) {
void findInlinedFunctions(DenseSet<GlobalValue::GUID> &S,
     if (isDeclaration(SymbolMap.lookup(getFuncName()))) {
       // Add to the import list only when it's defined out of module.
       S.insert(getGUID(getName()));
     }
     // Import hot CallTargets, which may not be available in IR because full
     // profile annotation cannot be done until backend compilation in ThinLTO.
     for (const auto &BS : BodySamples)
       for (const auto &TS : BS.second.getCallTargets())
-        if (TS.getValue() > Threshold) {
-          const Function *Callee = SymbolMap.lookup(getFuncName(TS.getKey()));
+        if (TS.second > Threshold) {
+          const Function *Callee = SymbolMap.lookup(getFuncName(TS.first));
           if (isDeclaration(Callee))
-            S.insert(getGUID(TS.getKey()));
+            S.insert(getGUID(TS.first));
         }
     for (const auto &CS : CallsiteSamples)
       for (const auto &NameFS : CS.second)
         NameFS.second.findInlinedFunctions(S, SymbolMap, Threshold);
   }
 
   /// Set the name of the function.
   void setName(StringRef FunctionName) { Context.setName(FunctionName); }

llvm/lib/ProfileData/SampleProf.cpp

 
 /// Merge the samples in \p Other into this record.
 /// Optionally scale sample counts by \p Weight.
 sampleprof_error SampleRecord::merge(const SampleRecord &Other,
                                      uint64_t Weight) {
   sampleprof_error Result;
   Result = addSamples(Other.getSamples(), Weight);
   for (const auto &I : Other.getCallTargets()) {
-    MergeResult(Result, addCalledTarget(I.first(), I.second, Weight));
+    MergeResult(Result, addCalledTarget(I.first, I.second, Weight));
   }
   return Result;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void LineLocation::dump() const { print(dbgs()); }
 #endif
 
const FunctionSamples *FunctionSamples::findFunctionSamples(
   }
   return FS;
 }
 
 void FunctionSamples::findAllNames(DenseSet<StringRef> &NameSet) const {
   NameSet.insert(getName());
   for (const auto &BS : BodySamples)
     for (const auto &TS : BS.second.getCallTargets())
-      NameSet.insert(TS.getKey());
+      NameSet.insert(TS.first);
 
   for (const auto &CS : CallsiteSamples) {
     for (const auto &NameFS : CS.second) {
       NameSet.insert(NameFS.first);
       NameFS.second.findAllNames(NameSet);
     }
   }
 }

llvm/lib/ProfileData/SampleProfWriter.cpp

void SampleProfileWriterBinary::addContext(const SampleContext &Context) {
   addName(Context.getName());
 }
 
 void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
   // Add all the names in indirect call targets.
   for (const auto &I : S.getBodySamples()) {
     const SampleRecord &Sample = I.second;
     for (const auto &J : Sample.getCallTargets())
-      addName(J.first());
+      addName(J.first);
   }
 
   // Recursively add all the names for inlined callsites.
   for (const auto &J : S.getCallsiteSamples())
     for (const auto &FS : J.second) {
       const FunctionSamples &CalleeSamples = FS.second;
       addName(CalleeSamples.getName());
       addNames(CalleeSamples);

llvm/lib/Target/X86/X86InsertPrefetch.cpp

bool X86InsertPrefetch::findPrefetchInfo(const FunctionSamples *TopSamples,
 
   const ErrorOr<PrefetchHints> T = getPrefetchHints(TopSamples, MI);
   if (!T)
     return false;
   int16_t max_index = -1;
   // Convert serialized prefetch hints into PrefetchInfo objects, and populate
   // the Prefetches vector.
   for (const auto &S_V : *T) {
-    StringRef Name = S_V.getKey();
+    StringRef Name = S_V.first;
     if (Name.consume_front(SerializedPrefetchPrefix)) {
       int64_t D = static_cast<int64_t>(S_V.second);
       unsigned IID = 0;
       for (const auto &HintType : HintTypes) {
         if (Name.startswith(HintType.first)) {
           Name = Name.drop_front(HintType.first.size());
           IID = HintType.second;
           break;

llvm/lib/Transforms/IPO/SampleProfile.cpp

while (!CalleeList.empty()) {
     // Add to the import list only when it's defined out of module.
     if (!Func || Func->isDeclaration())
       InlinedGUIDs.insert(FunctionSamples::getGUID(CalleeSample->getName()));
 
     // Import hot CallTargets, which may not be available in IR because full
     // profile annotation cannot be done until backend compilation in ThinLTO.
     for (const auto &BS : CalleeSample->getBodySamples())
       for (const auto &TS : BS.second.getCallTargets())
-        if (TS.getValue() > Threshold) {
-          StringRef CalleeName = CalleeSample->getFuncName(TS.getKey());
+        if (TS.second > Threshold) {
+          StringRef CalleeName = CalleeSample->getFuncName(TS.first);
           const Function *Callee = SymbolMap.lookup(CalleeName);
           if (!Callee || Callee->isDeclaration())
-            InlinedGUIDs.insert(FunctionSamples::getGUID(TS.getKey()));
+            InlinedGUIDs.insert(FunctionSamples::getGUID(TS.first));
         }
 
     // Import hot child context profile associted with callees. Note that this
     // may have some overlap with the call target loop above, but doing this
     // based child context profile again effectively allow us to use the max of
     // entry count and call target count to determine importing.
     for (auto &Child : Node->getAllChildContext()) {
       ContextTrieNode *CalleeNode = &Child.second;

 void SampleProfileMatcher::populateProfileCallsites(
     const FunctionSamples &FS,
     StringMap<std::set<LineLocation>> &CalleeToCallsitesMap) {
   for (const auto &I : FS.getBodySamples()) {
     const auto &Loc = I.first;
     const auto &CTM = I.second.getCallTargets();
     // Filter out possible indirect calls, use direct callee name as anchor.
     if (CTM.size() == 1) {
-      StringRef CalleeName = CTM.begin()->first();
+      StringRef CalleeName = CTM.begin()->first;
       const auto &Candidates = CalleeToCallsitesMap.try_emplace(
           CalleeName, std::set<LineLocation>());
       Candidates.first->second.insert(Loc);
     }
   }
 
   for (const auto &I : FS.getCallsiteSamples()) {
     const LineLocation &Loc = I.first;

llvm/tools/llvm-profdata/llvm-profdata.cpp

remapSamples(const sampleprof::FunctionSamples &Samples,
   for (const auto &BodySample : Samples.getBodySamples()) {
     uint32_t MaskedDiscriminator =
         BodySample.first.Discriminator & getDiscriminatorMask();
     Result.addBodySamples(BodySample.first.LineOffset, MaskedDiscriminator,
                           BodySample.second.getSamples());
     for (const auto &Target : BodySample.second.getCallTargets()) {
       Result.addCalledTargetSamples(BodySample.first.LineOffset,
                                     MaskedDiscriminator,
-                                    Remapper(Target.first()), Target.second);
+                                    Remapper(Target.first), Target.second);
     }
   }
   for (const auto &CallsiteSamples : Samples.getCallsiteSamples()) {
     sampleprof::FunctionSamplesMap &Target =
         Result.functionSamplesAt(CallsiteSamples.first);
     for (const auto &Callsite : CallsiteSamples.second) {
       sampleprof::FunctionSamples Remapped =
           remapSamples(Callsite.second, Remapper, Error);