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

[wlei]

Summary

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.

[wlei]

fix one wrong code copy by mistake

[wlei]

[NFC] remove two "getBinary"

[hoy]

The code structure looks cleaner now, thanks.

llvm/tools/llvm-profgen/PerfReader.cpp

259

Is this moved to anywhere?

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

259

Good catch, I was thinking the code below should catch the error, but I got it's better to remain the original logic.

// Attempt to open the binary.
OwningBinary<Binary> OBinary = unwrapOrError(createBinary(Path), Path);
Binary &Binary = *OBinary.getBinary();

auto *Obj = dyn_cast<ELFObjectFileBase>(&Binary);
if (!Obj)
  exitWithError("not a valid Elf image", Path);

[wlei]

address Hongtao's feedback

[wlei]

[NFC] tweak comments

llvm/tools/llvm-profgen/PerfReader.h

private:
   void recordBranchCount(const LBREntry &Branch, UnwindState &State,
                          uint64_t Repeat);
 
   ContextSampleCounterMap *CtxCounterMap;
   // Profiled binary that current frame address belongs to
   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 ...
   // A heuristic for fast detection by checking whether a
   // leading "  0x" and the '/' exist.
   static bool isLBRSample(StringRef Line) {
     if (!Line.startswith(" 0x"))
       return false;
public:
   struct MMapEvent {
     uint64_t PID = 0;
     uint64_t Address = 0;
     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; }
   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:
-  /// 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
   /// mapping between the binary name and its memory layout.
   ///
   void parseMMap2Event(TraceStream &TraceIt);
   // Parse perf events/samples and do aggregation
   void parseAndAggregateTrace(StringRef Filename);
which is used to generate CS profile. An example of hybrid sample:
     4005dc    # call stack leaf
     400634
     400684    # call stack root
     0x4005c8/0x4005dc/P/-/-/0   0x40062f/0x4005b0/P/-/-/0 ...
           ... 0x4005c8/0x4005dc/P/-/-/0    # LBR Entries
 */
 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
   void parseSample(TraceStream &TraceIt, uint64_t Count) override;
   void generateRawProfile() override;
 
 private:
   // Unwind the hybrid samples after aggregration
   void unwindSamples();
   void printUnwinderOutput();
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PerfReader.cpp

 //===-- PerfReader.cpp - perfscript reader  ---------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #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,
                                     cl::desc("Print binary load events."));
 
 static cl::opt<bool> ShowUnwinderOutput("show-unwinder-output",
                                         cl::ReallyHidden, cl::init(false),
                                         cl::ZeroOrMore,
while (State.hasNextLBR()) {
     recordBranchCount(Branch, State, Repeat);
   }
   // As samples are aggregated on trie, record them into counter map
   collectSamplesFromFrameTrie(State.getDummyRootPtr());
 
   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!");
   } else {
     exitWithError("Unsupported perfscript!");
   }
 
   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);
   if (Binary->getName() != BinaryName)
     return;
 
   // Drop the event if its image is loaded at the same address
   if (Event.Address == Binary->getBaseAddress()) {

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "ErrorHandling.h"
 #include "PerfReader.h"
 #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"
 
 static cl::OptionCategory ProfGenCategory("ProfGen Options");
 
 static cl::list<std::string> PerfTraceFilenames(
     "perfscript", cl::value_desc("perfscript"), cl::OneOrMore,
     llvm::cl::MiscFlags::CommaSeparated,
     cl::desc("Path of perf-script trace created by Linux perf tool with "
              "`script` command(the raw perf.data should be profiled with -b)"),
     cl::cat(ProfGenCategory));
 
 static cl::opt<std::string> BinaryPath(
     "binary", cl::value_desc("binary"), cl::Required,
     cl::desc("Path of profiled binary, only one binary is supported."),
     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(cl::list<std::string> &PerfTraceFilenames) {
+  for (auto &File : PerfTraceFilenames) {
+    if (!llvm::sys::fs::exists(File)) {
+      std::string Msg = "Input perf script(" + File + ") 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);
 
   // Initialize targets and assembly printers/parsers.
   InitializeAllTargetInfos();
   InitializeAllTargetMCs();
   InitializeAllDisassemblers();
 
   cl::HideUnrelatedOptions({&ProfGenCategory, &getColorCategory()});
   cl::ParseCommandLineOptions(argc, argv, "llvm SPGO profile generator\n");
+  validateCommandLine(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();
 
   return EXIT_SUCCESS;
 }