[SamplePGO] Fixing a memory issue when creating profiles on-demand

[hoy]

Summary

There is a on-dmeand creation of function profile during top-down processing in the sample loader when merging uninlined callees. During the profile creation, a stack string object is used to store a newly-created MD5 name, which is then used by reference as hash key in the profile map. This makes the hash key a dangling reference when later on the stack string object is deallocated.

The issue only happens with md5 profile use and was exposed by context split work for CS profile. I'm making a fix by storing newly created names in the reader.

[wlei]

LGTM, thanks for the fix!

[wenlei]

Thanks for the fix. I guess since we didn't really use the key of the map, it was just a hidden bug not causing trouble in the past (no ICE)?

llvm/include/llvm/ProfileData/SampleProfReader.h

419

Can we assert this only happens for MD5?

Additionally, NameBuffer -> MD5NameBuffer to make that explicit?

[hoy]

In D108142#2947508, @wenlei wrote:

Thanks for the fix. I guess since we didn't really use the key of the map, it was just a hidden bug not causing trouble in the past (no ICE)?

We have a test for this currently, but since it is small it doesn't capture the use-after-free case.

llvm/include/llvm/ProfileData/SampleProfReader.h

419

Sounds good.

[hoy]

Addressing Wenlei's comment.

[hoy]

Updating D108142: [SamplePGO] Fixing a memory issue when creating profiles on-demand

[wmi]

llvm/include/llvm/ProfileData/SampleProf.h

107–110

Thanks for the improvement.

llvm/include/llvm/ProfileData/SampleProfReader.h

419

Nit:

if (It != Profiles.end())
  return &It->second;
...

516

std::unordered_set is enough?

[hoy]

llvm/include/llvm/ProfileData/SampleProfReader.h

419

This looks cleaner.

516

Oops, that's what I meant to use. Good catch. Thanks.

[hoy]

Addresing Wei's comments.

[wenlei]

lgtm, thanks.

[wmi]

LGTM.

llvm/include/llvm/ProfileData/SampleProf.h

return uint64_t('S') << (64 - 8) | uint64_t('P') << (64 - 16) |
          uint64_t('F') << (64 - 40) | uint64_t('4') << (64 - 48) |
          uint64_t('2') << (64 - 56) | uint64_t(Format);
 }
 
 /// Get the proper representation of a string according to whether the
 /// current Format uses MD5 to represent the string.
 static inline StringRef getRepInFormat(StringRef Name, bool UseMD5,
                                        std::string &GUIDBuf) {
-  if (Name.empty())
+  if (Name.empty() || !UseMD5)
     return Name;
   GUIDBuf = std::to_string(Function::getGUID(Name));
-  return UseMD5 ? StringRef(GUIDBuf) : Name;
+  return GUIDBuf;
 }
 
 static inline uint64_t SPVersion() { return 103; }
 
 // Section Type used by SampleProfileExtBinaryBaseReader and
 // SampleProfileExtBinaryBaseWriter. Never change the existing
 // value of enum. Only append new ones.
 enum SecType {

llvm/include/llvm/ProfileData/SampleProfReader.h

 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SymbolRemappingReader.h"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <string>
 #include <system_error>
+#include <unordered_set>
 #include <vector>
 
 namespace llvm {
 
 class raw_ostream;
 class Twine;
 
 namespace sampleprof {
public:
   }
 
   /// Return the samples collected for function \p F, create empty
   /// FunctionSamples if it doesn't exist.
   FunctionSamples *getOrCreateSamplesFor(const Function &F) {
     std::string FGUID;
     StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
     CanonName = getRepInFormat(CanonName, useMD5(), FGUID);
+    auto It = Profiles.find(CanonName);
+    if (It != Profiles.end())
+      return &It->second;
+    if (!FGUID.empty()) {
+      assert(useMD5() && "New name should only be generated for md5 profile");
+      CanonName = *MD5NameBuffer.insert(FGUID).first;
+    }
     return &Profiles[CanonName];
   }
 
   /// Return the samples collected for function \p F.
   virtual FunctionSamples *getSamplesFor(StringRef Fname) {
     std::string FGUID;
     Fname = getRepInFormat(Fname, useMD5(), FGUID);
     auto It = Profiles.find(Fname);
     if (It != Profiles.end())
       return &It->second;
protected:
   StringMap<FunctionSamples> Profiles;
 
   /// LLVM context used to emit diagnostics.
   LLVMContext &Ctx;
 
   /// Memory buffer holding the profile file.
   std::unique_ptr<MemoryBuffer> Buffer;
 
+  /// Extra name buffer holding names created on demand.
+  /// This should only be needed for md5 profiles.
+  std::unordered_set<std::string> MD5NameBuffer;
+
   /// Profile summary information.
   std::unique_ptr<ProfileSummary> Summary;
 
   /// Take ownership of the summary of this reader.
   static std::unique_ptr<ProfileSummary>
   takeSummary(SampleProfileReader &Reader) {
     return std::move(Reader.Summary);
   }