[llvm-profgen] Add process filter for perf reader

For profile generation, we need to filter raw perf samples for binary of interest. Sometimes binary name along isn't enough as we can have binary of the same name running in the system. This change adds a process id filter to allow users to further disambiguiate the input raw samples.

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

Author: WenleiHe

llvm/tools/llvm-profgen/PerfReader.cpp

@@ -341,7 +341,8 @@ bool VirtualUnwinder::unwind(const PerfSample *Sample, uint64_t Repeat) {
 }
 
 std::unique_ptr<PerfReaderBase>
-PerfReaderBase::create(ProfiledBinary *Binary, PerfInputFile &PerfInput) {
+PerfReaderBase::create(ProfiledBinary *Binary, PerfInputFile &PerfInput,
+                       Optional<uint32_t> PIDFilter) {
   std::unique_ptr<PerfReaderBase> PerfReader;
 
   if (PerfInput.Format == PerfFormat::UnsymbolizedProfile) {
@@ -352,26 +353,28 @@ PerfReaderBase::create(ProfiledBinary *Binary, PerfInputFile &PerfInput) {
 
   // For perf data input, we need to convert them into perf script first.
   if (PerfInput.Format == PerfFormat::PerfData)
-    PerfInput = PerfScriptReader::convertPerfDataToTrace(Binary, PerfInput);
+    PerfInput =
+        PerfScriptReader::convertPerfDataToTrace(Binary, PerfInput, PIDFilter);
 
   assert((PerfInput.Format == PerfFormat::PerfScript) &&
          "Should be a perfscript!");
 
   PerfInput.Content =
       PerfScriptReader::checkPerfScriptType(PerfInput.InputFile);
   if (PerfInput.Content == PerfContent::LBRStack) {
-    PerfReader.reset(new HybridPerfReader(Binary, PerfInput.InputFile));
+    PerfReader.reset(
+        new HybridPerfReader(Binary, PerfInput.InputFile, PIDFilter));
   } else if (PerfInput.Content == PerfContent::LBR) {
-    PerfReader.reset(new LBRPerfReader(Binary, PerfInput.InputFile));
+    PerfReader.reset(new LBRPerfReader(Binary, PerfInput.InputFile, PIDFilter));
   } else {
     exitWithError("Unsupported perfscript!");
   }
 
   return PerfReader;
 }
 
-PerfInputFile PerfScriptReader::convertPerfDataToTrace(ProfiledBinary *Binary,
-                                                       PerfInputFile &File) {
+PerfInputFile PerfScriptReader::convertPerfDataToTrace(
+    ProfiledBinary *Binary, PerfInputFile &File, Optional<uint32_t> PIDFilter) {
   StringRef PerfData = File.InputFile;
   // Run perf script to retrieve PIDs matching binary we're interested in.
   auto PerfExecutable = sys::Process::FindInEnvPath("PATH", "perf");
@@ -397,7 +400,7 @@ PerfInputFile PerfScriptReader::convertPerfDataToTrace(ProfiledBinary *Binary,
     if (isMMap2Event(TraceIt.getCurrentLine()) &&
         extractMMap2EventForBinary(Binary, TraceIt.getCurrentLine(), MMap)) {
       auto It = PIDSet.emplace(MMap.PID);
-      if (It.second) {
+      if (It.second && (!PIDFilter || MMap.PID == *PIDFilter)) {
         if (!PIDs.empty()) {
           PIDs.append(",");
         }
@@ -426,6 +429,10 @@ void PerfScriptReader::updateBinaryAddress(const MMapEvent &Event) {
   if (Binary->getName() != BinaryName)
     return;
 
+  // Drop the event if process does not match pid filter
+  if (PIDFilter && Event.PID != *PIDFilter)
+    return;
+
   // Drop the event if its image is loaded at the same address
   if (Event.Address == Binary->getBaseAddress()) {
     Binary->setIsLoadedByMMap(true);

llvm/tools/llvm-profgen/PerfReader.h

@@ -13,6 +13,7 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Regex.h"
+#include <bits/stdint-uintn.h>
 #include <cstdint>
 #include <fstream>
 #include <list>
@@ -544,7 +545,8 @@ class PerfReaderBase {
   };
   virtual ~PerfReaderBase() = default;
   static std::unique_ptr<PerfReaderBase> create(ProfiledBinary *Binary,
-                                                PerfInputFile &PerfInput);
+                                                PerfInputFile &PerfInput,
+                                                Optional<uint32_t> PIDFilter);
 
   // Entry of the reader to parse multiple perf traces
   virtual void parsePerfTraces() = 0;
@@ -568,14 +570,16 @@ class PerfReaderBase {
 // Read perf script to parse the events and samples.
 class PerfScriptReader : public PerfReaderBase {
 public:
-  PerfScriptReader(ProfiledBinary *B, StringRef PerfTrace)
-      : PerfReaderBase(B, PerfTrace){};
+  PerfScriptReader(ProfiledBinary *B, StringRef PerfTrace,
+                   Optional<uint32_t> PID)
+      : PerfReaderBase(B, PerfTrace), PIDFilter(PID){};
 
   // Entry of the reader to parse multiple perf traces
   virtual void parsePerfTraces() override;
   // Generate perf script from perf data
   static PerfInputFile convertPerfDataToTrace(ProfiledBinary *Binary,
-                                              PerfInputFile &File);
+                                              PerfInputFile &File,
+                                              Optional<uint32_t> PIDFilter);
   // Extract perf script type by peaking at the input
   static PerfContent checkPerfScriptType(StringRef FileName);
 
@@ -635,6 +639,8 @@ class PerfScriptReader : public PerfReaderBase {
   AggregatedCounter AggregatedSamples;
   // Keep track of all invalid return addresses
   std::set<uint64_t> InvalidReturnAddresses;
+  // PID for the process of interest
+  Optional<uint32_t> PIDFilter;
 };
 
 /*
@@ -645,8 +651,9 @@ class PerfScriptReader : public PerfReaderBase {
 */
 class LBRPerfReader : public PerfScriptReader {
 public:
-  LBRPerfReader(ProfiledBinary *Binary, StringRef PerfTrace)
-      : PerfScriptReader(Binary, PerfTrace){};
+  LBRPerfReader(ProfiledBinary *Binary, StringRef PerfTrace,
+                Optional<uint32_t> PID)
+      : PerfScriptReader(Binary, PerfTrace, PID){};
   // Parse the LBR only sample.
   virtual void parseSample(TraceStream &TraceIt, uint64_t Count) override;
 };
@@ -662,8 +669,9 @@ class LBRPerfReader : public PerfScriptReader {
 */
 class HybridPerfReader : public PerfScriptReader {
 public:
-  HybridPerfReader(ProfiledBinary *Binary, StringRef PerfTrace)
-      : PerfScriptReader(Binary, PerfTrace){};
+  HybridPerfReader(ProfiledBinary *Binary, StringRef PerfTrace,
+                   Optional<uint32_t> PID)
+      : PerfScriptReader(Binary, PerfTrace, PID){};
   // Parse the hybrid sample including the call and LBR line
   void parseSample(TraceStream &TraceIt, uint64_t Count) override;
   void generateUnsymbolizedProfile() override;

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

@@ -60,6 +60,11 @@ static cl::opt<std::string>
                cl::desc("Path of profiled executable binary."),
                cl::cat(ProfGenCategory));
 
+static cl::opt<uint32_t>
+    ProcessId("pid", cl::value_desc("process Id"), cl::ZeroOrMore, cl::init(0),
+              cl::desc("Process Id for the profiled executable binary."),
+              cl::cat(ProfGenCategory));
+
 static cl::opt<std::string> DebugBinPath(
     "debug-binary", cl::value_desc("debug-binary"), cl::ZeroOrMore,
     cl::desc("Path of debug info binary, llvm-profgen will load the DWARF info "
@@ -168,9 +173,12 @@ int main(int argc, const char *argv[]) {
     Generator->generateProfile();
     Generator->write();
   } else {
+    Optional<uint32_t> PIDFilter;
+    if (ProcessId.getNumOccurrences())
+      PIDFilter = ProcessId;
     PerfInputFile PerfFile = getPerfInputFile();
     std::unique_ptr<PerfReaderBase> Reader =
-        PerfReaderBase::create(Binary.get(), PerfFile);
+        PerfReaderBase::create(Binary.get(), PerfFile, PIDFilter);
     // Parse perf events and samples
     Reader->parsePerfTraces();