[llvm-profgen] Use shared string table for function name to save memory usage

[wlei]

Summary

Currently we have three maps all using function name std::string, a) FuncName in pseudo probe section(GUID2FuncDescMap) b) FuncName in debug info(BinaryFunctions) c) symbolizer (NameStrings). We can optimize to use one shared map and other two use StringRef to save memory usage.

[wenlei]

llvm/lib/MC/MCPseudoProbe.cpp

354

Just double checking, so the underlying data for probe section is freed after decoding? Where does that happen?

[hoy]

Thanks for the change. How much space do you see this helps save?

llvm/lib/MC/MCPseudoProbe.cpp

354

The real name is on the disk. When read, a temp object ErrorOrName is created which is then freed.

[wenlei]

llvm/lib/MC/MCPseudoProbe.cpp

354

ErrorOrName doesn't own the underlying data. I was asking about the underlying data. The data is MCPseudoProbeDecoder::Data. Is OwningBinary taking care of the freeing?

[hoy]

llvm/lib/MC/MCPseudoProbe.cpp

354

Yes, the binary object owns that underlying data buffer, i.e, pointed by const uint8_t *Start.

[wlei]

In D115950#3200422, @hoy wrote:

Thanks for the change. How much space do you see this helps save?

I just tested it on our search infra, surprisingly I didn't see any memory saved by this. I dived more and it showed that the FuncNameStrings is not a big consumer, only ~1GB compared to the total consumption (200GB+), that's petty small. Another guess is the library have some optimization to allocate the string memory, but I'm not quite familiar with underlying optimization of string allocation.

llvm/lib/MC/MCPseudoProbe.cpp

354

Yeah, the input binary info OwningBinary<Binary> OBinary is all freed after load() completed, otherwise we can directly use StringRef before.

[hoy]

It's still good to have this optimization, thanks.

[wenlei]

lgtm, thanks.

llvm/include/llvm/MC/MCPseudoProbe.h

 #include "llvm/Support/ErrorOr.h"
 #include <list>
 #include <map>
 #include <memory>
 #include <string>
 #include <tuple>
 #include <type_traits>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace llvm {
 
 class MCSection;
 class MCStreamer;
 class MCSymbol;
 class MCObjectStreamer;
 class raw_ostream;
 
 enum class MCPseudoProbeFlag {
   // If set, indicates that the probe is encoded as an address delta
   // instead of a real code address.
   AddressDelta = 0x1,
 };
 
 // Function descriptor decoded from .pseudo_probe_desc section
 struct MCPseudoProbeFuncDesc {
   uint64_t FuncGUID = 0;
   uint64_t FuncHash = 0;
-  std::string FuncName;
+  StringRef FuncName;
 
   MCPseudoProbeFuncDesc(uint64_t GUID, uint64_t Hash, StringRef Name)
       : FuncGUID(GUID), FuncHash(Hash), FuncName(Name){};
 
   void print(raw_ostream &OS);
 };
 
 class MCPseudoProbe;
class MCPseudoProbeDecoder {
   // Decoding helper function
   template <typename T> ErrorOr<T> readUnencodedNumber();
   template <typename T> ErrorOr<T> readUnsignedNumber();
   template <typename T> ErrorOr<T> readSignedNumber();
   ErrorOr<StringRef> readString(uint32_t Size);
 
 public:
   // Decode pseudo_probe_desc section to build GUID to PseudoProbeFuncDesc map.
-  bool buildGUID2FuncDescMap(const uint8_t *Start, std::size_t Size);
+  bool buildGUID2FuncDescMap(const uint8_t *Start, std::size_t Size,
+                             std::unordered_set<std::string> &FuncNameStrings);
 
   // Decode pseudo_probe section to build address to probes map.
   bool buildAddress2ProbeMap(const uint8_t *Start, std::size_t Size);
 
   // Print pseudo_probe_desc section info
   void printGUID2FuncDescMap(raw_ostream &OS);
 
   // Print pseudo_probe section info, used along with show-disassembly

llvm/lib/MC/MCPseudoProbe.cpp

ErrorOr<StringRef> MCPseudoProbeDecoder::readString(uint32_t Size) {
   StringRef Str(reinterpret_cast<const char *>(Data), Size);
   if (Data + Size > End) {
     return std::error_code();
   }
   Data += Size;
   return ErrorOr<StringRef>(Str);
 }
 
-bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start,
-                                                 std::size_t Size) {
+bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(
+    const uint8_t *Start, std::size_t Size,
+    std::unordered_set<std::string> &FuncNameStrings) {
   // The pseudo_probe_desc section has a format like:
   // .section .pseudo_probe_desc,"",@progbits
   // .quad -5182264717993193164   // GUID
   // .quad 4294967295             // Hash
   // .uleb 3                      // Name size
   // .ascii "foo"                 // Name
   // .quad -2624081020897602054
   // .quad 174696971957
while (Data < End) {
     auto ErrorOrName = readString(NameSize);
     if (!ErrorOrName)
       return false;
 
     uint64_t GUID = std::move(*ErrorOrGUID);
     uint64_t Hash = std::move(*ErrorOrHash);
     StringRef Name = std::move(*ErrorOrName);
 
+    auto It = FuncNameStrings.insert(Name.str());
     // Initialize PseudoProbeFuncDesc and populate it into GUID2FuncDescMap
-    GUID2FuncDescMap.emplace(GUID, MCPseudoProbeFuncDesc(GUID, Hash, Name));
+    GUID2FuncDescMap.emplace(GUID,
+                             MCPseudoProbeFuncDesc(GUID, Hash, *It.first));
   }
   assert(Data == End && "Have unprocessed data in pseudo_probe_desc section");
   return true;
 }
 
 bool MCPseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
                                                  std::size_t Size) {
   // The pseudo_probe section encodes an inline forest and each tree has a

llvm/tools/llvm-profgen/ProfiledBinary.h

class ProfiledBinary {
   std::unique_ptr<const MCInstrInfo> MII;
   std::unique_ptr<MCDisassembler> DisAsm;
   std::unique_ptr<const MCInstrAnalysis> MIA;
   std::unique_ptr<MCInstPrinter> IPrinter;
   // A list of text sections sorted by start RVA and size. Used to check
   // if a given RVA is a valid code address.
   std::set<std::pair<uint64_t, uint64_t>> TextSections;
 
+  // Shared string table owning function name strings, keep only one table to
+  // save the memory.
+  std::unordered_set<std::string> FuncNameStrings;
+
   // A map of mapping function name to BinaryFunction info.
-  std::unordered_map<std::string, BinaryFunction> BinaryFunctions;
+  StringMap<BinaryFunction> BinaryFunctions;
 
   // An ordered map of mapping function's start offset to function range
   // relevant info. Currently to determine if the offset of ELF is the start of
   // a real function, we leverage the function range info from DWARF.
   std::map<uint64_t, FuncRange> StartOffset2FuncRangeMap;
 
   // Offset to context location map. Used to expand the context.
   std::unordered_map<uint64_t, SampleContextFrameVector> Offset2LocStackMap;
class ProfiledBinary {
   PrologEpilogTracker ProEpilogTracker;
 
   // Track function sizes under different context
   BinarySizeContextTracker FuncSizeTracker;
 
   // The symbolizer used to get inline context for an instruction.
   std::unique_ptr<symbolize::LLVMSymbolizer> Symbolizer;
 
-  // String table owning function name strings created from the symbolizer.
-  std::unordered_set<std::string> NameStrings;
-
   // A collection of functions to print disassembly for.
   StringSet<> DisassembleFunctionSet;
 
   // Pseudo probe decoder
   MCPseudoProbeDecoder ProbeDecoder;
 
   bool UsePseudoProbes = false;
 
RangesTy getRangesForOffset(uint64_t Offset) {
     auto *FRange = findFuncRangeForOffset(Offset);
     // Ignore the range which falls into plt section or system lib.
     if (!FRange)
       return RangesTy();
 
     return FRange->Func->Ranges;
   }
 
-  const std::unordered_map<std::string, BinaryFunction> &
-  getAllBinaryFunctions() {
+  const StringMap<BinaryFunction> &getAllBinaryFunctions() {
     return BinaryFunctions;
   }
 
   BinaryFunction *getBinaryFunction(StringRef FName) {
     auto I = BinaryFunctions.find(FName.str());
     if (I == BinaryFunctions.end())
       return nullptr;
     return &I->second;

llvm/tools/llvm-profgen/ProfiledBinary.cpp

for (auto &R : F.second.Ranges) {
         }
       }
     }
 
     if (!hasFuncEntry) {
       NoFuncEntryNum++;
       if (ShowDetailedWarning)
         WithColor::warning()
-            << "Failed to determine function entry for " << F.first
+            << "Failed to determine function entry for " << F.first()
             << " due to inconsistent name from symbol table and dwarf info.\n";
     }
   }
   emitWarningSummary(NoFuncEntryNum, BinaryFunctions.size(),
                      "of functions failed to determine function entry due to "
                      "inconsistent name from symbol table and dwarf info.");
 }
 
for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
        SI != SE; ++SI) {
     const SectionRef &Section = *SI;
     StringRef SectionName = unwrapOrError(Section.getName(), FileName);
 
     if (SectionName == ".pseudo_probe_desc") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
       if (!ProbeDecoder.buildGUID2FuncDescMap(
               reinterpret_cast<const uint8_t *>(Contents.data()),
-              Contents.size()))
+              Contents.size(), FuncNameStrings))
         exitWithError("Pseudo Probe decoder fail in .pseudo_probe_desc section");
     } else if (SectionName == ".pseudo_probe") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
       if (!ProbeDecoder.buildAddress2ProbeMap(
               reinterpret_cast<const uint8_t *>(Contents.data()),
               Contents.size()))
         exitWithError("Pseudo Probe decoder fail in .pseudo_probe section");
       // set UsePseudoProbes flag, used for PerfReader
for (const auto &DieInfo : CompilationUnit->dies()) {
         continue;
       const DWARFAddressRangesVector &Ranges = RangesOrError.get();
 
       if (Ranges.empty())
         continue;
 
       // Different DWARF symbols can have same function name, search or create
       // BinaryFunction indexed by the name.
-      auto Ret = BinaryFunctions.emplace(Name, BinaryFunction());
+      // Stablize in a shared function name string table.
+      auto It = FuncNameStrings.insert(Name);
+      auto Ret = BinaryFunctions.try_emplace(*It.first, BinaryFunction());
       auto &Func = Ret.first->second;
       if (Ret.second)
-        Func.FuncName = Ret.first->first;
+        Func.FuncName = Ret.first->first();
 
       for (const auto &Range : Ranges) {
         uint64_t FuncStart = Range.LowPC;
         uint64_t FuncSize = Range.HighPC - FuncStart;
 
         if (FuncSize == 0 || FuncStart < getPreferredBaseAddress())
           continue;
 
for (int32_t I = InlineStack.getNumberOfFrames() - 1; I >= 0; I--) {
     uint32_t LineOffset = (CallerFrame.Line - CallerFrame.StartLine) & 0xffff;
     if (UseProbeDiscriminator) {
       LineOffset =
           PseudoProbeDwarfDiscriminator::extractProbeIndex(Discriminator);
       Discriminator = 0;
     }
 
     LineLocation Line(LineOffset, Discriminator);
-    auto It = NameStrings.insert(FunctionName.str());
+    auto It = FuncNameStrings.insert(FunctionName.str());
     CallStack.emplace_back(*It.first, Line);
   }
 
   return CallStack;
 }
 
 void ProfiledBinary::computeInlinedContextSizeForRange(uint64_t StartOffset,
                                                        uint64_t EndOffset) {