[llvm-profgen] Use shared pointer to fix memory leak issue from asan

[wlei]

Summary

As titled.

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

299

I think you can use unique_ptr here. Everywhere else consuming the binary can use the raw pointer, until PerfReaderBase is deleted.

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

299

Yeah, just realized this, thanks!

llvm/tools/llvm-profgen/PerfReader.h

struct MMapEvent {
     uint64_t Size = 0;
     uint64_t Offset = 0;
     StringRef BinaryPath;
   };
 
   /// Load symbols and disassemble the code of a given binary.
   void loadBinary(const StringRef BinaryPath);
   void updateBinaryAddress(const MMapEvent &Event);
-  ProfiledBinary *getBinary() const { return Binary; }
+  std::shared_ptr<ProfiledBinary> getBinary() const { return Binary; }
   PerfScriptType getPerfScriptType() const { return PerfType; }
   // Entry of the reader to parse multiple perf traces
   void parsePerfTraces(cl::list<std::string> &PerfTraceFilenames);
   const ContextSampleCounterMap &getSampleCounters() const {
     return SampleCounters;
   }
 
 protected:
protected:
   void parseSample(TraceStream &TraceIt);
   // An aggregated count is given to indicate how many times the sample is
   // repeated.
   virtual void parseSample(TraceStream &TraceIt, uint64_t Count) = 0;
   // Post process the profile after trace aggregation, we will do simple range
   // overlap computation for AutoFDO, or unwind for CSSPGO(hybrid sample).
   virtual void generateRawProfile() = 0;
 
-  ProfiledBinary *Binary = nullptr;
+  std::shared_ptr<ProfiledBinary> Binary;
 
   ContextSampleCounterMap SampleCounters;
   // Samples with the repeating time generated by the perf reader
   AggregatedCounter AggregatedSamples;
   PerfScriptType PerfType = PERF_UNKNOWN;
 };
 
 /*

llvm/tools/llvm-profgen/PerfReader.cpp

 
 PerfReaderBase::PerfReaderBase(StringRef BinaryPath) {
   // Load the binary.
   loadBinary(BinaryPath);
 }
 
 void PerfReaderBase::loadBinary(const StringRef BinaryPath) {
   // Call to load the binary in the ctor of ProfiledBinary.
-  Binary = new ProfiledBinary(BinaryPath);
+  Binary = std::make_shared<ProfiledBinary>(BinaryPath);
   // Initialize the base address to preferred address.
   Binary->setBaseAddress(Binary->getPreferredBaseAddress());
 }
 
 void PerfReaderBase::updateBinaryAddress(const MMapEvent &Event) {
   // Drop the event which doesn't belong to user-provided binary
   StringRef BinaryName = llvm::sys::path::filename(Event.BinaryPath);
   if (Binary->getName() != BinaryName)
for (auto &CI : Counter) {
     OrderedCounter[getContextKeyStr(CI.first.getPtr(), Binary)] =
         CI.second.BranchCounter;
   }
   printSampleCounter(OrderedCounter);
 }
 
 void HybridPerfReader::printUnwinderOutput() {
     outs() << "Binary(" << Binary->getName().str() << ")'s Range Counter:\n";
-    printRangeCounter(SampleCounters, Binary);
+    printRangeCounter(SampleCounters, Binary.get());
     outs() << "\nBinary(" << Binary->getName().str() << ")'s Branch Counter:\n";
-    printBranchCounter(SampleCounters, Binary);
+    printBranchCounter(SampleCounters, Binary.get());
 }
 
 void HybridPerfReader::unwindSamples() {
   for (const auto &Item : AggregatedSamples) {
     const HybridSample *Sample = dyn_cast<HybridSample>(Item.first.getPtr());
-    VirtualUnwinder Unwinder(&SampleCounters, Binary);
+    VirtualUnwinder Unwinder(&SampleCounters, Binary.get());
     Unwinder.unwind(Sample, Item.second);
   }
 
   if (ShowUnwinderOutput)
     printUnwinderOutput();
 }
 
 bool PerfReaderBase::extractLBRStack(TraceStream &TraceIt,
void HybridPerfReader::parseSample(TraceStream &TraceIt, uint64_t Count) {
   // intermediately by LBR sample
   // e.g.
   //            4005dc    # call stack leaf
   //            400634
   //            400684    # call stack root
   // 0x4005c8/0x4005dc/P/-/-/0   0x40062f/0x4005b0/P/-/-/0 ...
   //          ... 0x4005c8/0x4005dc/P/-/-/0    # LBR Entries
   //
-  std::shared_ptr<HybridSample> Sample = std::make_shared<HybridSample>(Binary);
+  std::shared_ptr<HybridSample> Sample =
+      std::make_shared<HybridSample>(Binary.get());
 
   // Parsing call stack and populate into HybridSample.CallStack
   if (!extractCallstack(TraceIt, Sample->CallStack)) {
     // Skip the next LBR line matched current call stack
     if (!TraceIt.isAtEoF() && TraceIt.getCurrentLine().startswith(" 0x"))
       TraceIt.advance();
     return;
   }

llvm/tools/llvm-profgen/ProfileGenerator.h

 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 class ProfileGenerator {
 
 public:
-  ProfileGenerator(ProfiledBinary *B) : Binary(B){};
+  ProfileGenerator(std::shared_ptr<ProfiledBinary> B) : Binary(B){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
-  create(ProfiledBinary *Binary, const ContextSampleCounterMap &SampleCounters,
+  create(std::shared_ptr<ProfiledBinary> Binary,
+         const ContextSampleCounterMap &SampleCounters,
          enum PerfScriptType SampleType);
   virtual void generateProfile() = 0;
   // Use SampleProfileWriter to serialize profile map
   virtual void write(std::unique_ptr<SampleProfileWriter> Writer,
                      StringMap<FunctionSamples> &ProfileMap);
   void write();
 
 protected:
protected:
   sample count for disjoint region [A,B] would be 300.
   */
   void findDisjointRanges(RangeSample &DisjointRanges,
                           const RangeSample &Ranges);
 
   // Used by SampleProfileWriter
   StringMap<FunctionSamples> ProfileMap;
 
-  ProfiledBinary *Binary = nullptr;
+  std::shared_ptr<ProfiledBinary> Binary;
 };
 
 class CSProfileGenerator : public ProfileGenerator {
 protected:
   const ContextSampleCounterMap &SampleCounters;
 
 public:
-  CSProfileGenerator(ProfiledBinary *Binary,
+  CSProfileGenerator(std::shared_ptr<ProfiledBinary> Binary,
                      const ContextSampleCounterMap &Counters)
       : ProfileGenerator(Binary), SampleCounters(Counters){};
 
 public:
   void generateProfile() override;
 
   // Trim the context stack at a given depth.
   template <typename T>

 };
 
 using ProbeCounterMap =
     std::unordered_map<const MCDecodedPseudoProbe *, uint64_t>;
 
 class PseudoProbeCSProfileGenerator : public CSProfileGenerator {
 
 public:
-  PseudoProbeCSProfileGenerator(ProfiledBinary *Binary,
+  PseudoProbeCSProfileGenerator(std::shared_ptr<ProfiledBinary> Binary,
                                 const ContextSampleCounterMap &Counters)
       : CSProfileGenerator(Binary, Counters) {}
   void generateProfile() override;
 
 private:
   // Go through each address from range to extract the top frame probe by
   // looking up in the Address2ProbeMap
   void extractProbesFromRange(const RangeSample &RangeCounter,

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 namespace sampleprof {
 
 // Initialize the MaxCompressionSize to -1 which means no size limit
 int32_t CSProfileGenerator::MaxCompressionSize = -1;
 
 int CSProfileGenerator::MaxContextDepth = -1;
 
 std::unique_ptr<ProfileGenerator>
-ProfileGenerator::create(ProfiledBinary *Binary,
+ProfileGenerator::create(std::shared_ptr<ProfiledBinary> Binary,
                          const ContextSampleCounterMap &SampleCounters,
                          enum PerfScriptType SampleType) {
   std::unique_ptr<ProfileGenerator> ProfileGenerator;
   if (SampleType == PERF_LBR_STACK) {
     if (Binary->usePseudoProbes()) {
       ProfileGenerator.reset(
           new PseudoProbeCSProfileGenerator(Binary, SampleCounters));
     } else {
for (auto Range : Ranges) {
     uint64_t RangeBegin = Binary->offsetToVirtualAddr(Range.first.first);
     uint64_t RangeEnd = Binary->offsetToVirtualAddr(Range.first.second);
     uint64_t Count = Range.second;
     // Disjoint ranges have introduce zero-filled gap that
     // doesn't belong to current context, filter them out.
     if (Count == 0)
       continue;
 
-    InstructionPointer IP(Binary, RangeBegin, true);
+    InstructionPointer IP(Binary.get(), RangeBegin, true);
 
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
     if (IP.Address > RangeEnd)
       continue;
 
     while (IP.Address <= RangeEnd) {

 void PseudoProbeCSProfileGenerator::generateProfile() {
   // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;
   FunctionSamples::ProfileIsCS = true;
   for (const auto &CI : SampleCounters) {
     const ProbeBasedCtxKey *CtxKey =
         dyn_cast<ProbeBasedCtxKey>(CI.first.getPtr());
     SmallVector<std::string, 16> ContextStrStack;
-    extractPrefixContextStack(ContextStrStack, CtxKey->Probes, Binary);
+    extractPrefixContextStack(ContextStrStack, CtxKey->Probes, Binary.get());
     // Fill in function body samples from probes, also infer caller's samples
     // from callee's probe
     populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStrStack);
     // Fill in boundary samples for a call probe
     populateBoundarySamplesWithProbes(CI.second.BranchCounter, ContextStrStack);
   }
 
   postProcessProfiles();
 }
 
 void PseudoProbeCSProfileGenerator::extractProbesFromRange(
     const RangeSample &RangeCounter, ProbeCounterMap &ProbeCounter) {
   RangeSample Ranges;
   findDisjointRanges(Ranges, RangeCounter);
   for (const auto &Range : Ranges) {
     uint64_t RangeBegin = Binary->offsetToVirtualAddr(Range.first.first);
     uint64_t RangeEnd = Binary->offsetToVirtualAddr(Range.first.second);
     uint64_t Count = Range.second;
     // Disjoint ranges have introduce zero-filled gap that
     // doesn't belong to current context, filter them out.
     if (Count == 0)
       continue;
 
-    InstructionPointer IP(Binary, RangeBegin, true);
+    InstructionPointer IP(Binary.get(), RangeBegin, true);
 
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
     if (IP.Address > RangeEnd)
       continue;
 
     while (IP.Address <= RangeEnd) {