[llvm-profgen] Move profiled binary loading out of PerfReader

Change to use unique pointer of profiled binary to unblock asan.

At same time, I realized we can decouple to move the profiled binary loading out of PerfReader, so I made some other related refactors.

Reviewed By: hoy

Differential Revision: https://reviews.llvm.org/D108254

llvm/tools/llvm-profgen/PerfReader.cpp

@@ -7,7 +7,6 @@
 //===----------------------------------------------------------------------===//
 #include "PerfReader.h"
 #include "ProfileGenerator.h"
-#include "llvm/Support/FileSystem.h"
 
 static cl::opt<bool> ShowMmapEvents("show-mmap-events", cl::ReallyHidden,
                                     cl::init(false), cl::ZeroOrMore,
@@ -244,41 +243,13 @@ bool VirtualUnwinder::unwind(const HybridSample *Sample, uint64_t Repeat) {
   return true;
 }
 
-void PerfReaderBase::validateCommandLine(
-    StringRef BinaryPath, cl::list<std::string> &PerfTraceFilenames) {
-  // Allow the invalid perfscript if we only use to show binary disassembly
-  if (!ShowDisassemblyOnly) {
-    for (auto &File : PerfTraceFilenames) {
-      if (!llvm::sys::fs::exists(File)) {
-        std::string Msg = "Input perf script(" + File + ") doesn't exist!";
-        exitWithError(Msg);
-      }
-    }
-  }
-
-  if (!llvm::sys::fs::exists(BinaryPath)) {
-    std::string Msg = "Input binary(" + BinaryPath.str() + ") doesn't exist!";
-    exitWithError(Msg);
-  }
-
-  if (CSProfileGenerator::MaxCompressionSize < -1) {
-    exitWithError("Value of --compress-recursion should >= -1");
-  }
-  if (ShowSourceLocations && !ShowDisassemblyOnly) {
-    exitWithError("--show-source-locations should work together with "
-                  "--show-disassembly-only!");
-  }
-}
-
 std::unique_ptr<PerfReaderBase>
-PerfReaderBase::create(StringRef BinaryPath,
+PerfReaderBase::create(ProfiledBinary *Binary,
                        cl::list<std::string> &PerfTraceFilenames) {
-  validateCommandLine(BinaryPath, PerfTraceFilenames);
-
   PerfScriptType PerfType = extractPerfType(PerfTraceFilenames);
   std::unique_ptr<PerfReaderBase> PerfReader;
   if (PerfType == PERF_LBR_STACK) {
-    PerfReader.reset(new HybridPerfReader(BinaryPath));
+    PerfReader.reset(new HybridPerfReader(Binary));
   } else if (PerfType == PERF_LBR) {
     // TODO:
     exitWithError("Unsupported perfscript!");
@@ -289,18 +260,6 @@ PerfReaderBase::create(StringRef BinaryPath,
   return PerfReader;
 }
 
-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);
-  // 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);

llvm/tools/llvm-profgen/PerfReader.h

@@ -295,7 +295,6 @@ struct UnwindState {
            "IP should align with context leaf");
   }
 
-  const ProfiledBinary *getBinary() const { return Binary; }
   bool hasNextLBR() const { return LBRIndex < LBRStack.size(); }
   uint64_t getCurrentLBRSource() const { return LBRStack[LBRIndex].Source; }
   uint64_t getCurrentLBRTarget() const { return LBRStack[LBRIndex].Target; }
@@ -531,13 +530,16 @@ private:
   const ProfiledBinary *Binary;
 };
 
-// Load binaries and read perf trace to parse the events and samples
+// Read perf trace to parse the events and samples.
 class PerfReaderBase {
 public:
-  PerfReaderBase(StringRef BinaryPath);
+  PerfReaderBase(ProfiledBinary *B) : Binary(B) {
+    // Initialize the base address to preferred address.
+    Binary->setBaseAddress(Binary->getPreferredBaseAddress());
+  };
   virtual ~PerfReaderBase() = default;
   static std::unique_ptr<PerfReaderBase>
-  create(StringRef BinaryPath, cl::list<std::string> &PerfTraceFilenames);
+  create(ProfiledBinary *Binary, cl::list<std::string> &PerfTraceFilenames);
 
   // A LBR sample is like:
   // 0x5c6313f/0x5c63170/P/-/-/0  0x5c630e7/0x5c63130/P/-/-/0 ...
@@ -593,10 +595,7 @@ public:
     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; }
   PerfScriptType getPerfScriptType() const { return PerfType; }
   // Entry of the reader to parse multiple perf traces
   void parsePerfTraces(cl::list<std::string> &PerfTraceFilenames);
@@ -605,9 +604,6 @@ public:
   }
 
 protected:
-  /// Validate the command line input
-  static void validateCommandLine(StringRef BinaryPath,
-                                  cl::list<std::string> &PerfTraceFilenames);
   static PerfScriptType
   extractPerfType(cl::list<std::string> &PerfTraceFilenames);
   /// Parse a single line of a PERF_RECORD_MMAP2 event looking for a
@@ -654,7 +650,7 @@ protected:
 */
 class HybridPerfReader : public PerfReaderBase {
 public:
-  HybridPerfReader(StringRef BinaryPath) : PerfReaderBase(BinaryPath) {
+  HybridPerfReader(ProfiledBinary *Binary) : PerfReaderBase(Binary) {
     PerfType = PERF_LBR_STACK;
   };
   // Parse the hybrid sample including the call and LBR line

llvm/tools/llvm-profgen/llvm-profgen.cpp

@@ -15,6 +15,7 @@
 #include "ProfileGenerator.h"
 #include "ProfiledBinary.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
@@ -33,10 +34,38 @@ static cl::opt<std::string> BinaryPath(
     cl::cat(ProfGenCategory));
 
 extern cl::opt<bool> ShowDisassemblyOnly;
+extern cl::opt<bool> ShowSourceLocations;
 
 using namespace llvm;
 using namespace sampleprof;
 
+// Validate the command line input.
+static void validateCommandLine(StringRef BinaryPath,
+                                cl::list<std::string> &PerfTraceFilenames) {
+  // Allow the invalid perfscript if we only use to show binary disassembly.
+  if (!ShowDisassemblyOnly) {
+    for (auto &File : PerfTraceFilenames) {
+      if (!llvm::sys::fs::exists(File)) {
+        std::string Msg = "Input perf script(" + File + ") doesn't exist!";
+        exitWithError(Msg);
+      }
+    }
+  }
+
+  if (!llvm::sys::fs::exists(BinaryPath)) {
+    std::string Msg = "Input binary(" + BinaryPath.str() + ") doesn't exist!";
+    exitWithError(Msg);
+  }
+
+  if (CSProfileGenerator::MaxCompressionSize < -1) {
+    exitWithError("Value of --compress-recursion should >= -1");
+  }
+  if (ShowSourceLocations && !ShowDisassemblyOnly) {
+    exitWithError("--show-source-locations should work together with "
+                  "--show-disassembly-only!");
+  }
+}
+
 int main(int argc, const char *argv[]) {
   InitLLVM X(argc, argv);
 
@@ -47,20 +76,21 @@ int main(int argc, const char *argv[]) {
 
   cl::HideUnrelatedOptions({&ProfGenCategory, &getColorCategory()});
   cl::ParseCommandLineOptions(argc, argv, "llvm SPGO profile generator\n");
+  validateCommandLine(BinaryPath, PerfTraceFilenames);
 
-  if (ShowDisassemblyOnly) {
-    (void)ProfiledBinary(BinaryPath);
+  // Load symbols and disassemble the code of a given binary.
+  std::unique_ptr<ProfiledBinary> Binary =
+      std::make_unique<ProfiledBinary>(BinaryPath);
+  if (ShowDisassemblyOnly)
     return EXIT_SUCCESS;
-  }
 
-  // Load binaries and parse perf events and samples
+  // Parse perf events and samples
   std::unique_ptr<PerfReaderBase> Reader =
-      PerfReaderBase::create(BinaryPath, PerfTraceFilenames);
+      PerfReaderBase::create(Binary.get(), PerfTraceFilenames);
   Reader->parsePerfTraces(PerfTraceFilenames);
 
-  std::unique_ptr<ProfileGenerator> Generator =
-      ProfileGenerator::create(Reader->getBinary(), Reader->getSampleCounters(),
-                               Reader->getPerfScriptType());
+  std::unique_ptr<ProfileGenerator> Generator = ProfileGenerator::create(
+      Binary.get(), Reader->getSampleCounters(), Reader->getPerfScriptType());
   Generator->generateProfile();
   Generator->write();