[llvm-profdata] Condense string allocations for MD5 function names

[huangjd]

Summary

Allocate a big buffer enough to hold all MD5 strings, instead of allocating it on heap one by one.
Also materialize all MD5 strings if the full profile is being read since all these strings are generated anyways.

Improved around 3% speed on loading an ExtBinary profile with 1GB total size/10 million names.

This refactoring takes precendence over the refactoring of using MD5 as key

[davidxl]

llvm/include/llvm/ProfileData/SampleProfReader.h

662

Can we reduce the size using hex encoding?

llvm/include/llvm/ProfileData/SampleProfReader.h

public:
   /// Return whether any name in the profile contains ".__uniq." suffix.
   virtual bool hasUniqSuffix() { return false; }
 
   SampleProfileReaderItaniumRemapper *getRemapper() { return Remapper.get(); }
 
   void setModule(const Module *Mod) { M = Mod; }
 
 protected:
+  BumpPtrAllocator Allocator;
+
   /// Map every function to its associated profile.
   ///
   /// The profile of every function executed at runtime is collected
   /// in the structure FunctionSamples. This maps function objects
   /// to their corresponding profiles.
   SampleProfileMap Profiles;
 
   /// LLVM context used to emit diagnostics.
protected:
   std::error_code readSummary();
 
   /// Read the whole name table with regular strings.
   std::error_code readNameTable();
 
   /// Read the whole name table with ULEB128 encoded MD5 values.
   std::error_code readMD5NameTable();
 
+  /// Set MD5StringBuf size to be allocated for the whole MD5 name table.
+  /// uint64_t::max() == 18446744073709551616 has 20 bytes.
+  void SetMD5StringBufSize(size_t N) { MD5StringBufSize = 21 * N; }
+
+  /// Materialize a string from MD5 value into MD5StringBuf.
+  StringRef allocateMD5String(uint64_t FID);
+
   /// Read a string indirectly via the name table.
   ErrorOr<StringRef> readStringFromTable();
 
   /// Read a context indirectly via the CSNameTable.
   ErrorOr<SampleContextFrames> readContextFromTable();
 
   /// Read a context indirectly via the CSNameTable if the profile has context,
   /// otherwise same as readStringFromTable.
protected:
 
   /// Function name table.
   std::vector<StringRef> NameTable;
 
   /// If MD5 is used in NameTable section, the section saves uint64_t data.
   /// The uint64_t data has to be converted to a string and then the string
   /// will be used to initialize StringRef in NameTable.
   /// Note NameTable contains StringRef so it needs another buffer to own
-  /// the string data. MD5StringBuf serves as the string buffer that is
-  /// referenced by NameTable (vector of StringRef). We make sure
-  /// the lifetime of MD5StringBuf is not shorter than that of NameTable.
-  std::vector<std::string> MD5StringBuf;
+  /// the string data. MD5StringBuf is a pointer to an allocated slab from
+  /// Allocator that serves as the char buffer referenced by NameTable.
+  /// It is incremented to the next writable position every time a string is
+  /// allocated, or a new slab is allocated if it exceeds MD5StringBufLimit.
+  char *MD5StringBuf = reinterpret_cast<char *>(1);
+
+  /// End of the current MD5 string buffer.
+  const char *MD5StringBufLimit = nullptr;
+
+  /// Size of MD5 string buffer per allocation.
+  size_t MD5StringBufSize;
 
   /// The starting address of NameTable containing fixed length MD5.
   const uint8_t *MD5NameMemStart = nullptr;
 
   /// CSNameTable is used to save full context vectors. It is the backing buffer
   /// for SampleContextFrames.
   std::vector<SampleContextFrameVector> CSNameTable;
 

 /// possible to define other types of profile inherited from
 /// SampleProfileReaderExtBinaryBase/SampleProfileWriterExtBinaryBase.
 class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
 private:
   std::error_code decompressSection(const uint8_t *SecStart,
                                     const uint64_t SecSize,
                                     const uint8_t *&DecompressBuf,
                                     uint64_t &DecompressBufSize);
-
-  BumpPtrAllocator Allocator;
-
 protected:
   std::vector<SecHdrTableEntry> SecHdrTable;
   std::error_code readSecHdrTableEntry(uint64_t Idx);
   std::error_code readSecHdrTable();
 
   std::error_code readFuncMetadata(bool ProfileHasAttribute);
   std::error_code readFuncMetadata(bool ProfileHasAttribute,
                                    FunctionSamples *FProfile);

llvm/lib/ProfileData/SampleProfReader.cpp

inline ErrorOr<size_t> SampleProfileReaderBinary::readStringIndex(T &Table) {
   auto Idx = readNumber<size_t>();
   if (std::error_code EC = Idx.getError())
     return EC;
   if (*Idx >= Table.size())
     return sampleprof_error::truncated_name_table;
   return *Idx;
 }
 
+StringRef SampleProfileReaderBinary::allocateMD5String(uint64_t FID) {
+  if (MD5StringBuf > MD5StringBufLimit) {
+    // uint64_t::max() == 18446744073709551616 has 20 bytes. The limit is
+    // reduced by 21 so that the next write always have enough space.
+    size_t AllocSize = std::max(MD5StringBufSize, (size_t) 4096);
+    MD5StringBuf = Allocator.Allocate<char>(AllocSize);
+    MD5StringBufLimit = MD5StringBuf + AllocSize - 21;
+  }
+  char *Start = MD5StringBuf;
+  size_t N = sprintf(MD5StringBuf, "%llu", (unsigned long long) FID);
+  MD5StringBuf += N + 1;
+  return StringRef(Start, N);
+}
+
 ErrorOr<StringRef> SampleProfileReaderBinary::readStringFromTable() {
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
 
   // Check whether the name to be accessed has been accessed before,
   // if not, read it from memory directly.
   StringRef &SR = NameTable[*Idx];
   if (!SR.data()) {
     assert(MD5NameMemStart);
     using namespace support;
     uint64_t FID = endian::read<uint64_t, little, unaligned>(
         MD5NameMemStart + (*Idx) * sizeof(uint64_t));
-    SR = MD5StringBuf.emplace_back(std::to_string(FID));
+    SR = allocateMD5String(FID);
   }
   return SR;
 }
 
 std::error_code
 SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
   auto NumSamples = readNumber<uint64_t>();
   if (std::error_code EC = NumSamples.getError())
if (std::error_code EC = Size.getError())
     return EC;
 
   // Normally if useMD5 is true, the name table should have MD5 values, not
   // strings, however in the case that ExtBinary profile has multiple name
   // tables mixing string and MD5, all of them have to be normalized to use MD5,
   // because optimization passes can only handle either type.
   bool UseMD5 = useMD5();
   if (UseMD5)
-    MD5StringBuf.reserve(MD5StringBuf.size() + *Size);
-
+    SetMD5StringBufSize(*Size);
   NameTable.clear();
   NameTable.reserve(*Size);
   for (size_t I = 0; I < *Size; ++I) {
     auto Name(readString());
     if (std::error_code EC = Name.getError())
       return EC;
     if (UseMD5) {
       uint64_t MD5 = MD5Hash(*Name);
-      NameTable.emplace_back(MD5StringBuf.emplace_back(std::to_string(MD5)));
+      NameTable.emplace_back(allocateMD5String(MD5));
     } else
       NameTable.push_back(*Name);
   }
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderBinary::readMD5NameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
 
-  MD5StringBuf.reserve(MD5StringBuf.size() + *Size);
+  SetMD5StringBufSize(*Size);
   NameTable.clear();
   NameTable.reserve(*Size);
   for (size_t I = 0; I < *Size; ++I) {
     auto FID = readNumber<uint64_t>();
     if (std::error_code EC = FID.getError())
       return EC;
-    NameTable.emplace_back(MD5StringBuf.emplace_back(std::to_string(*FID)));
+    NameTable.emplace_back(allocateMD5String(*FID));
   }
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5,
                                                    bool FixedLengthMD5) {
   if (!FixedLengthMD5) {
assert(Data + (*Size) * sizeof(uint64_t) == End &&
          "entries");
   if (Data + (*Size) * sizeof(uint64_t) > End)
     return sampleprof_error::truncated;
 
   // Preallocate and initialize NameTable so we can check whether a name
   // index has been read before by checking whether the element in the
   // NameTable is empty, meanwhile readStringIndex can do the boundary
   // check using the size of NameTable.
+  SetMD5StringBufSize(*Size);
   NameTable.clear();
+  // If in tools mode, every profile is read so every name is read, loading them
+  // all now.
+  if (!M) {
+    NameTable.reserve(*Size);
+    for (size_t I = 0; I < *Size; I++) {
+      using namespace support;
+      uint64_t FID = endian::read<uint64_t, little, unaligned>(Data);
+      NameTable.emplace_back(allocateMD5String(FID));
+      Data += sizeof(uint64_t);
+    }
+  } else {
     NameTable.resize(*Size);
     MD5NameMemStart = Data;
     Data = Data + (*Size) * sizeof(uint64_t);
+  }
   return sampleprof_error::success;
 }
 
 // Read in the CS name table section, which basically contains a list of context
 // vectors. Each element of a context vector, aka a frame, refers to the
 // underlying raw function names that are stored in the name table, as well as
 // a callsite identifier that only makes sense for non-leaf frames.
 std::error_code SampleProfileReaderExtBinaryBase::readCSNameTableSec() {