[llvm-profgen] Clean up code dealing with multiple binaries

[wlei]

Summary

As we decided to support only one binary each time, this patch cleans up the related code dealing with multiple binaries. We can use llvm-profdata to merge profile from multiple binaries.

[wlei]

fix comments

[hoy]

Thanks for the cleanup. The new code looks cleaner and I believe it'll also speed up the tool.

llvm/tools/llvm-profgen/PerfReader.h

533–534

Do we still need this? Looks like this can be just a single field of PerfReaderBase.

[wlei]

remove unneeded code

llvm/tools/llvm-profgen/PerfReader.h

533–534

Good catch! We shouldn't need this, removed.

[wlei]

fixed typos

[hoy]

lgtm except for a minor issue.

llvm/tools/llvm-profgen/PerfReader.cpp

350

The comment could still be helpful?

[wlei]

Tweak some comments

[wenlei]

llvm/tools/llvm-profgen/PerfReader.cpp

296–297

is the removal of const intentional?

[wlei]

add "const" qualifier

llvm/tools/llvm-profgen/PerfReader.cpp

296–297

Good catch!

llvm/tools/llvm-profgen/PerfReader.h

 struct HybridSample : public PerfSample {
   // Profiled binary that current frame address belongs to
   ProfiledBinary *Binary;
   // Call stack recorded in FILO(leaf to root) order
   SmallVector<uint64_t, 16> CallStack;
   // LBR stack recorded in FIFO order
   SmallVector<LBREntry, 16> LBRStack;
 
-  HybridSample() : PerfSample(PK_HybridSample){};
+  HybridSample(ProfiledBinary *B) : PerfSample(PK_HybridSample), Binary(B){};
   static bool classof(const PerfSample *K) {
     return K->getKind() == PK_HybridSample;
   }
 
   // Used for sample aggregation
   bool isEqual(const PerfSample *K) const override {
     const HybridSample *Other = dyn_cast<HybridSample>(K);
     if (Other->Binary != Binary)
void recordRangeCount(uint64_t Start, uint64_t End, UnwindState &State,
                         uint64_t Repeat);
   void recordBranchCount(const LBREntry &Branch, UnwindState &State,
                          uint64_t Repeat);
 
   ContextSampleCounterMap *CtxCounterMap;
   // Profiled binary that current frame address belongs to
   const ProfiledBinary *Binary;
 };
 
-// Filename to binary map
-using BinaryMap = StringMap<ProfiledBinary>;
-// Address to binary map for fast look-up
-using AddressBinaryMap = std::map<uint64_t, ProfiledBinary *>;
-// Binary to ContextSampleCounters Map to support multiple binary, we may have
-// same binary loaded at different addresses, they should share the same sample
-// counter
-using BinarySampleCounterMap =
-    std::unordered_map<ProfiledBinary *, ContextSampleCounterMap>;
-
 // Load binaries and read perf trace to parse the events and samples
 class PerfReaderBase {
 public:
-  PerfReaderBase(cl::list<std::string> &BinaryFilenames);
+  PerfReaderBase(StringRef BinaryPath);
   virtual ~PerfReaderBase() = default;
   static std::unique_ptr<PerfReaderBase>
-  create(cl::list<std::string> &BinaryFilenames,
-         cl::list<std::string> &PerfTraceFilenames);
+  create(StringRef BinaryPath, 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 give binary.
-  /// Also register the binary in the binary table.
-  ///
-  ProfiledBinary &loadBinary(const StringRef BinaryPath,
-                             bool AllowNameConflict = true);
+  /// Load symbols and disassemble the code of a given binary.
+  void loadBinary(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 BinarySampleCounterMap &getBinarySampleCounters() const {
-    return BinarySampleCounters;
+  const ContextSampleCounterMap &getSampleCounters() const {
+    return SampleCounters;
   }
 
 protected:
   /// Validate the command line input
-  static void validateCommandLine(cl::list<std::string> &BinaryFilenames,
+  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);
   // Parse either an MMAP event or a perf sample
   void parseEventOrSample(TraceStream &TraceIt);
   // Extract call stack from the perf trace lines
   bool extractCallstack(TraceStream &TraceIt,
                         SmallVectorImpl<uint64_t> &CallStack);
   // Extract LBR stack from one perf trace line
   bool extractLBRStack(TraceStream &TraceIt,
-                       SmallVectorImpl<LBREntry> &LBRStack,
-                       ProfiledBinary *Binary);
+                       SmallVectorImpl<LBREntry> &LBRStack);
   uint64_t parseAggregatedCount(TraceStream &TraceIt);
   // Parse one sample from multiple perf lines, override this for different
   // sample type
   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;
-  // Helper function for looking up binary in AddressBinaryMap
-  ProfiledBinary *getBinary(uint64_t Address);
 
-  BinaryMap BinaryTable;
-  AddressBinaryMap AddrToBinaryMap; // Used by address-based lookup.
+  ProfiledBinary *Binary = nullptr;
 
-  BinarySampleCounterMap BinarySampleCounters;
+  ContextSampleCounterMap SampleCounters;
   // Samples with the repeating time generated by the perf reader
   AggregatedCounter AggregatedSamples;
   PerfScriptType PerfType = PERF_UNKNOWN;
 };
 
 /*
   Hybrid perf script includes a group of hybrid samples(LBRs + call stack),
   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(cl::list<std::string> &BinaryFilenames)
-      : PerfReaderBase(BinaryFilenames) {
+  HybridPerfReader(StringRef BinaryPath) : PerfReaderBase(BinaryPath) {
     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

bool VirtualUnwinder::unwind(const HybridSample *Sample, uint64_t Repeat) {
   }
   // As samples are aggregated on trie, record them into counter map
   collectSamplesFromFrameTrie(State.getDummyRootPtr());
 
   return true;
 }
 
 void PerfReaderBase::validateCommandLine(
-    cl::list<std::string> &BinaryFilenames,
-    cl::list<std::string> &PerfTraceFilenames) {
+    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 (BinaryFilenames.size() > 1) {
-    // TODO: remove this if everything is ready to support multiple binaries.
-    exitWithError(
-        "Currently only support one input binary, multiple binaries' "
-        "profile will be merged in one profile and make profile "
-        "summary info inaccurate. Please use `llvm-perfdata` to merge "
-        "profiles from multiple binaries.");
-  }
-  for (auto &Binary : BinaryFilenames) {
-    if (!llvm::sys::fs::exists(Binary)) {
-      std::string Msg = "Input binary(" + Binary + ") doesn't exist!";
-      exitWithError(Msg);
-    }
-  }
+
+  if (!llvm::sys::fs::exists(BinaryPath)) {
+    std::string Msg = "Input binary(" + BinaryPath.str() + ") doesn't exist!";
+
   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(cl::list<std::string> &BinaryFilenames,
+PerfReaderBase::create(StringRef BinaryPath,
                        cl::list<std::string> &PerfTraceFilenames) {
-  validateCommandLine(BinaryFilenames, PerfTraceFilenames);
+  validateCommandLine(BinaryPath, PerfTraceFilenames);
 
   PerfScriptType PerfType = extractPerfType(PerfTraceFilenames);
   std::unique_ptr<PerfReaderBase> PerfReader;
   if (PerfType == PERF_LBR_STACK) {
-    PerfReader.reset(new HybridPerfReader(BinaryFilenames));
+    PerfReader.reset(new HybridPerfReader(BinaryPath));
   } else if (PerfType == PERF_LBR) {
     // TODO:
     exitWithError("Unsupported perfscript!");
   } else {
     exitWithError("Unsupported perfscript!");
   }
 
   return PerfReader;
 }
 
-PerfReaderBase::PerfReaderBase(cl::list<std::string> &BinaryFilenames) {
-  // Load the binaries.
-  for (auto Filename : BinaryFilenames)
-    loadBinary(Filename, /*AllowNameConflict*/ false);
-}
-
-ProfiledBinary &PerfReaderBase::loadBinary(const StringRef BinaryPath,
-                                           bool AllowNameConflict) {
-  // The binary table is currently indexed by the binary name not the full
-  // binary path. This is because the user-given path may not match the one
-  // that was actually executed.
-  StringRef BinaryName = llvm::sys::path::filename(BinaryPath);
-
-  // Call to load the binary in the ctor of ProfiledBinary.
-  auto Ret = BinaryTable.insert({BinaryName, ProfiledBinary(BinaryPath)});
-
-  if (!Ret.second && !AllowNameConflict) {
-    std::string ErrorMsg = "Binary name conflict: " + BinaryPath.str() +
-                           " and " + Ret.first->second.getPath().str() + " \n";
-    exitWithError(ErrorMsg);
-  }
+PerfReaderBase::PerfReaderBase(StringRef BinaryPath) {
+  // Load the binary.
+  loadBinary(BinaryPath);
 
+void PerfReaderBase::loadBinary(StringRef BinaryPath) {
+  Binary = new ProfiledBinary(BinaryPath);
   // Initialize the base address to preferred address.
-  ProfiledBinary &B = Ret.first->second;
-  uint64_t PreferredAddr = B.getPreferredBaseAddress();
-  AddrToBinaryMap[PreferredAddr] = &B;
-  B.setBaseAddress(PreferredAddr);
-
-  return B;
+  Binary->setBaseAddress(Binary->getPreferredBaseAddress());
 }
 
 void PerfReaderBase::updateBinaryAddress(const MMapEvent &Event) {
-  // Load the binary.
-  StringRef BinaryPath = Event.BinaryPath;
-  StringRef BinaryName = llvm::sys::path::filename(BinaryPath);
-
-  auto I = BinaryTable.find(BinaryName);
-  // Drop the event which doesn't belong to user-provided binaries
-  if (I == BinaryTable.end())
+  // Drop the event which doesn't belong to user-provided binary
+  StringRef BinaryName = llvm::sys::path::filename(Event.BinaryPath);
+  if (Binary->getName() != BinaryName)
     return;
 
-  ProfiledBinary &Binary = I->second;
   // Drop the event if its image is loaded at the same address
-  if (Event.Address == Binary.getBaseAddress()) {
-    Binary.setIsLoadedByMMap(true);
+  if (Event.Address == Binary->getBaseAddress()) {
+    Binary->setIsLoadedByMMap(true);
     return;
   }
 
-  if (Event.Offset == Binary.getTextSegmentOffset()) {
-    // A binary image could be unloaded and then reloaded at different
-    // place, so update the address map here.
-    // Only update for the first executable segment and assume all other
-    // segments are loaded at consecutive memory addresses, which is the case on
-    // X64.
-    AddrToBinaryMap.erase(Binary.getBaseAddress());
-    AddrToBinaryMap[Event.Address] = &Binary;
-
+  if (Event.Offset == Binary->getTextSegmentOffset()) {
     // Update binary load address.
-    Binary.setBaseAddress(Event.Address);
-
-    Binary.setIsLoadedByMMap(true);
+    Binary->setBaseAddress(Event.Address);
+    Binary->setIsLoadedByMMap(true);
   } else {
     // Verify segments are loaded consecutively.
-    const auto &Offsets = Binary.getTextSegmentOffsets();
+    const auto &Offsets = Binary->getTextSegmentOffsets();
     auto It = std::lower_bound(Offsets.begin(), Offsets.end(), Event.Offset);
     if (It != Offsets.end() && *It == Event.Offset) {
       // The event is for loading a separate executable segment.
       auto I = std::distance(Offsets.begin(), It);
-      const auto &PreferredAddrs = Binary.getPreferredTextSegmentAddresses();
-      if (PreferredAddrs[I] - Binary.getPreferredBaseAddress() !=
-          Event.Address - Binary.getBaseAddress())
+      const auto &PreferredAddrs = Binary->getPreferredTextSegmentAddresses();
+      if (PreferredAddrs[I] - Binary->getPreferredBaseAddress() !=
+          Event.Address - Binary->getBaseAddress())
         exitWithError("Executable segments not loaded consecutively");
     } else {
       if (It == Offsets.begin())
         exitWithError("File offset not found");
       else {
         // Find the segment the event falls in. A large segment could be loaded
         // via multiple mmap calls with consecutive memory addresses.
         --It;
         assert(*It < Event.Offset);
-        if (Event.Offset - *It != Event.Address - Binary.getBaseAddress())
+        if (Event.Offset - *It != Event.Address - Binary->getBaseAddress())
           exitWithError("Segment not loaded by consecutive mmaps");
       }
     }
   }
 }
 
-ProfiledBinary *PerfReaderBase::getBinary(uint64_t Address) {
-  auto Iter = AddrToBinaryMap.lower_bound(Address);
-  if (Iter == AddrToBinaryMap.end() || Iter->first != Address) {
-    if (Iter == AddrToBinaryMap.begin())
-      return nullptr;
-    Iter--;
-  }
-  return Iter->second;
-}
-
 // Use ordered map to make the output deterministic
 using OrderedCounterForPrint = std::map<std::string, RangeSample>;
 
 static void printSampleCounter(OrderedCounterForPrint &OrderedCounter) {
   for (auto Range : OrderedCounter) {
     outs() << Range.first << "\n";
     for (auto I : Range.second) {
       outs() << "  (" << format("%" PRIx64, I.first.first) << ", "
static void printBranchCounter(ContextSampleCounterMap &Counter,
   for (auto &CI : Counter) {
     OrderedCounter[getContextKeyStr(CI.first.getPtr(), Binary)] =
         CI.second.BranchCounter;
   }
   printSampleCounter(OrderedCounter);
 }
 
 void HybridPerfReader::printUnwinderOutput() {
-  for (auto I : BinarySampleCounters) {
-    const ProfiledBinary *Binary = I.first;
     outs() << "Binary(" << Binary->getName().str() << ")'s Range Counter:\n";
-    printRangeCounter(I.second, Binary);
+    printRangeCounter(SampleCounters, Binary);
     outs() << "\nBinary(" << Binary->getName().str() << ")'s Branch Counter:\n";
-    printBranchCounter(I.second, Binary);
-  }
+    printBranchCounter(SampleCounters, Binary);
 }
 
 void HybridPerfReader::unwindSamples() {
   for (const auto &Item : AggregatedSamples) {
     const HybridSample *Sample = dyn_cast<HybridSample>(Item.first.getPtr());
-    VirtualUnwinder Unwinder(&BinarySampleCounters[Sample->Binary],
-                             Sample->Binary);
+    VirtualUnwinder Unwinder(&SampleCounters, Binary);
     Unwinder.unwind(Sample, Item.second);
   }
 
   if (ShowUnwinderOutput)
     printUnwinderOutput();
 }
 
 bool PerfReaderBase::extractLBRStack(TraceStream &TraceIt,
-                                     SmallVectorImpl<LBREntry> &LBRStack,
-                                     ProfiledBinary *Binary) {
+                                     SmallVectorImpl<LBREntry> &LBRStack) {
   // The raw format of LBR stack is like:
   // 0x4005c8/0x4005dc/P/-/-/0 0x40062f/0x4005b0/P/-/-/0 ...
   //                           ... 0x4005c8/0x4005dc/P/-/-/0
   // It's in FIFO order and seperated by whitespace.
   SmallVector<StringRef, 32> Records;
   TraceIt.getCurrentLine().split(Records, " ");
 
   // Extract leading instruction pointer if present, use single
bool PerfReaderBase::extractCallstack(TraceStream &TraceIt,
                                       SmallVectorImpl<uint64_t> &CallStack) {
   // The raw format of call stack is like:
   //            4005dc      # leaf frame
   //            400634
   //            400684      # root frame
   // It's in bottom-up order with each frame in one line.
 
   // Extract stack frames from sample
-  ProfiledBinary *Binary = nullptr;
   while (!TraceIt.isAtEoF() && !TraceIt.getCurrentLine().startswith(" 0x")) {
     StringRef FrameStr = TraceIt.getCurrentLine().ltrim();
     uint64_t FrameAddr = 0;
     if (FrameStr.getAsInteger(16, FrameAddr)) {
       // We might parse a non-perf sample line like empty line and comments,
       // skip it
       TraceIt.advance();
       return false;
     }
     TraceIt.advance();
-    if (!Binary) {
-      Binary = getBinary(FrameAddr);
-      // we might have addr not match the MMAP, skip it
-      if (!Binary) {
-        if (AddrToBinaryMap.size() == 0)
-          WithColor::warning() << "No MMAP event in the perfscript, create it "
-                                  "with '--show-mmap-events'\n";
-        break;
-      }
-    }
     // Currently intermixed frame from different binaries is not supported.
     // Ignore bottom frames not from binary of interest.
     if (!Binary->addressIsCode(FrameAddr))
       break;
 
     // We need to translate return address to call address
     // for non-leaf frames
     if (!CallStack.empty()) {
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>();
+  std::shared_ptr<HybridSample> Sample = std::make_shared<HybridSample>(Binary);
 
   // 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;
   }
-  // Set the binary current sample belongs to
-  ProfiledBinary *PB = getBinary(Sample->CallStack.front());
-  Sample->Binary = PB;
-  if (!PB->getMissingMMapWarned() && !PB->getIsLoadedByMMap()) {
+
+  if (!Binary->getMissingMMapWarned() && !Binary->getIsLoadedByMMap()) {
     WithColor::warning() << "No relevant mmap event is matched, will use "
                             "preferred address as the base loading address!\n";
     // Avoid redundant warning, only warn at the first unmatched sample.
-    PB->setMissingMMapWarned(true);
+    Binary->setMissingMMapWarned(true);
   }
 
   if (!TraceIt.isAtEoF() && TraceIt.getCurrentLine().startswith(" 0x")) {
     // Parsing LBR stack and populate into HybridSample.LBRStack
-    if (extractLBRStack(TraceIt, Sample->LBRStack, Sample->Binary)) {
+    if (extractLBRStack(TraceIt, Sample->LBRStack)) {
       // Canonicalize stack leaf to avoid 'random' IP from leaf frame skew LBR
       // ranges
       Sample->CallStack.front() = Sample->LBRStack[0].Target;
       // Record samples by aggregation
       Sample->genHashCode();
       AggregatedSamples[Hashable<PerfSample>(Sample)] += Count;
     }
   } else {

llvm/tools/llvm-profgen/ProfileGenerator.h

 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 class ProfileGenerator {
 
 public:
-  ProfileGenerator(){};
+  ProfileGenerator(ProfiledBinary *B) : Binary(B){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
-  create(const BinarySampleCounterMap &BinarySampleCounters,
+  create(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;
 };
 
 class CSProfileGenerator : public ProfileGenerator {
 protected:
-  const BinarySampleCounterMap &BinarySampleCounters;
+  const ContextSampleCounterMap &SampleCounters;
 
 public:
-  CSProfileGenerator(const BinarySampleCounterMap &Counters)
-      : BinarySampleCounters(Counters){};
+  CSProfileGenerator(ProfiledBinary *Binary,
+                     const ContextSampleCounterMap &Counters)
+      : ProfileGenerator(Binary), SampleCounters(Counters){};
 
 public:
   void generateProfile() override;
 
   // Trim the context stack at a given depth.
   template <typename T>
   static void trimContext(SmallVectorImpl<T> &S, int Depth = MaxContextDepth) {
     if (Depth < 0 || static_cast<size_t>(Depth) >= S.size())

 
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
-                                   const RangeSample &RangeCounters,
-                                   ProfiledBinary *Binary);
+                                   const RangeSample &RangeCounters);
   void populateFunctionBoundarySamples(StringRef ContextId,
                                        FunctionSamples &FunctionProfile,
-                                       const BranchSample &BranchCounters,
-                                       ProfiledBinary *Binary);
+                                       const BranchSample &BranchCounters);
   void populateInferredFunctionSamples();
 
 public:
   // Deduplicate adjacent repeated context sequences up to a given sequence
   // length. -1 means no size limit.
   static int32_t MaxCompressionSize;
   static int MaxContextDepth;
 };
 
 using ProbeCounterMap =
     std::unordered_map<const MCDecodedPseudoProbe *, uint64_t>;
 
 class PseudoProbeCSProfileGenerator : public CSProfileGenerator {
 
 public:
-  PseudoProbeCSProfileGenerator(const BinarySampleCounterMap &Counters)
-      : CSProfileGenerator(Counters) {}
+  PseudoProbeCSProfileGenerator(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,
-                              ProbeCounterMap &ProbeCounter,
-                              ProfiledBinary *Binary);
+                              ProbeCounterMap &ProbeCounter);
   // Fill in function body samples from probes
   void
   populateBodySamplesWithProbes(const RangeSample &RangeCounter,
-                                SmallVectorImpl<std::string> &ContextStrStack,
-                                ProfiledBinary *Binary);
+                                SmallVectorImpl<std::string> &ContextStrStack);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(
       const BranchSample &BranchCounter,
-      SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary);
+      SmallVectorImpl<std::string> &ContextStrStack);
   // Helper function to get FunctionSamples for the leaf inlined context
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
                                  const MCPseudoProbeFuncDesc *LeafFuncDesc,
                                  bool WasLeafInlined);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
-                                 const MCDecodedPseudoProbe *LeafProbe,
-                                 ProfiledBinary *Binary);
+                                 const MCDecodedPseudoProbe *LeafProbe);
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 namespace llvm {
 namespace sampleprof {
 
 // Initialize the MaxCompressionSize to -1 which means no size limit
 int32_t CSProfileGenerator::MaxCompressionSize = -1;
 
 int CSProfileGenerator::MaxContextDepth = -1;
 
-static bool
-usePseudoProbes(const BinarySampleCounterMap &BinarySampleCounters) {
-  return BinarySampleCounters.size() &&
-         BinarySampleCounters.begin()->first->usePseudoProbes();
-}
-
 std::unique_ptr<ProfileGenerator>
-ProfileGenerator::create(const BinarySampleCounterMap &BinarySampleCounters,
+ProfileGenerator::create(ProfiledBinary *Binary,
+                         const ContextSampleCounterMap &SampleCounters,
                          enum PerfScriptType SampleType) {
   std::unique_ptr<ProfileGenerator> ProfileGenerator;
   if (SampleType == PERF_LBR_STACK) {
-    if (usePseudoProbes(BinarySampleCounters)) {
+    if (Binary->usePseudoProbes()) {
       ProfileGenerator.reset(
-          new PseudoProbeCSProfileGenerator(BinarySampleCounters));
+          new PseudoProbeCSProfileGenerator(Binary, SampleCounters));
     } else {
-      ProfileGenerator.reset(new CSProfileGenerator(BinarySampleCounters));
+      ProfileGenerator.reset(new CSProfileGenerator(Binary, SampleCounters));
     }
   } else {
     // TODO:
     llvm_unreachable("Unsupported perfscript!");
   }
 
   return ProfileGenerator;
 }
if (Ret.second) {
     FProfile.setContext(FContext);
     FProfile.setName(FContext.getNameWithoutContext());
   }
   return Ret.first->second;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
-  for (const auto &BI : BinarySampleCounters) {
-    ProfiledBinary *Binary = BI.first;
-    for (const auto &CI : BI.second) {
-      const StringBasedCtxKey *CtxKey =
-          dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
-      StringRef ContextId(CtxKey->Context);
-      // Get or create function profile for the range
-      FunctionSamples &FunctionProfile =
-          getFunctionProfileForContext(ContextId, CtxKey->WasLeafInlined);
+  for (const auto &CI : SampleCounters) {
 
-      // Fill in function body samples
-      populateFunctionBodySamples(FunctionProfile, CI.second.RangeCounter,
-                                  Binary);
-      // Fill in boundary sample counts as well as call site samples for calls
-      populateFunctionBoundarySamples(ContextId, FunctionProfile,
-                                      CI.second.BranchCounter, Binary);
-    }
+    populateFunctionBodySamples(FunctionProfile, CI.second.RangeCounter);
+                                    CI.second.BranchCounter);
   }
   // Fill in call site value sample for inlined calls and also use context to
   // infer missing samples. Since we don't have call count for inlined
   // functions, we estimate it from inlinee's profile using the entry of the
   // body sample.
   populateInferredFunctionSamples();
 
   postProcessProfiles();
void CSProfileGenerator::updateBodySamplesforFunctionProfile(
   if (PreviousCount < Count) {
     FunctionProfile.addBodySamples(LeafLoc.second.LineOffset,
                                    LeafLoc.second.Discriminator,
                                    Count - PreviousCount);
   }
 }
 
 void CSProfileGenerator::populateFunctionBodySamples(
-    FunctionSamples &FunctionProfile, const RangeSample &RangeCounter,
-    ProfiledBinary *Binary) {
+    FunctionSamples &FunctionProfile, const RangeSample &RangeCounter) {
   // Compute disjoint ranges first, so we can use MAX
   // for calculating count for each location.
   RangeSample Ranges;
   findDisjointRanges(Ranges, RangeCounter);
   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;
while (IP.Address <= RangeEnd) {
       // Move to next IP within the range
       IP.advance();
     }
   }
 }
 
 void CSProfileGenerator::populateFunctionBoundarySamples(
     StringRef ContextId, FunctionSamples &FunctionProfile,
-    const BranchSample &BranchCounters, ProfiledBinary *Binary) {
+    const BranchSample &BranchCounters) {
 
   for (auto Entry : BranchCounters) {
     uint64_t SourceOffset = Entry.first.first;
     uint64_t TargetOffset = Entry.first.second;
     uint64_t Count = Entry.second;
     // Get the callee name by branch target if it's a call branch
     StringRef CalleeName = FunctionSamples::getCanonicalFnName(
         Binary->getFuncFromStartOffset(TargetOffset));
void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
                                StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = Writer->write(ProfileMap))
     exitWithError(std::move(EC));
 }
 
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
 // be added compressed while looking up function profile
-static void
-extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
+static void extractPrefixContextStack(
+    SmallVectorImpl<std::string> &ContextStrStack,
     const SmallVectorImpl<const MCDecodedPseudoProbe *> &Probes,
     ProfiledBinary *Binary) {
   for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStrStack, true);
   }
 }
 
 void PseudoProbeCSProfileGenerator::generateProfile() {
   // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;
   FunctionSamples::ProfileIsCS = true;
-  for (const auto &BI : BinarySampleCounters) {
-    ProfiledBinary *Binary = BI.first;
-    for (const auto &CI : BI.second) {
-      const ProbeBasedCtxKey *CtxKey =
-          dyn_cast<ProbeBasedCtxKey>(CI.first.getPtr());
-      SmallVector<std::string, 16> ContextStrStack;
-      extractPrefixContextStack(ContextStrStack, CtxKey->Probes, Binary);
-      // Fill in function body samples from probes, also infer caller's samples
-      // from callee's probe
-      populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStrStack,
-                                    Binary);
-      // Fill in boundary samples for a call probe
-      populateBoundarySamplesWithProbes(CI.second.BranchCounter,
-                                        ContextStrStack, Binary);
-    }
+  for (const auto &CI : SampleCounters) {
+    populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStrStack);
+    populateBoundarySamplesWithProbes(CI.second.BranchCounter, ContextStrStack);
   }
 
   postProcessProfiles();
 }
 
 void PseudoProbeCSProfileGenerator::extractProbesFromRange(
-    const RangeSample &RangeCounter, ProbeCounterMap &ProbeCounter,
-    ProfiledBinary *Binary) {
+    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.
while (IP.Address <= RangeEnd) {
 
       IP.advance();
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBodySamplesWithProbes(
     const RangeSample &RangeCounter,
-    SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {
+    SmallVectorImpl<std::string> &ContextStrStack) {
   ProbeCounterMap ProbeCounter;
   // Extract the top frame probes by looking up each address among the range in
   // the Address2ProbeMap
-  extractProbesFromRange(RangeCounter, ProbeCounter, Binary);
+  extractProbesFromRange(RangeCounter, ProbeCounter);
   std::unordered_map<MCDecodedPseudoProbeInlineTree *,
                      std::unordered_set<FunctionSamples *>>
       FrameSamples;
   for (auto PI : ProbeCounter) {
     const MCDecodedPseudoProbe *Probe = PI.first;
     uint64_t Count = PI.second;
     FunctionSamples &FunctionProfile =
-        getFunctionProfileForLeafProbe(ContextStrStack, Probe, Binary);
+        getFunctionProfileForLeafProbe(ContextStrStack, Probe);
     // Record the current frame and FunctionProfile whenever samples are
     // collected for non-danglie probes. This is for reporting all of the
     // zero count probes of the frame later.
     FrameSamples[Probe->getInlineTreeNode()].insert(&FunctionProfile);
     FunctionProfile.addBodySamplesForProbe(Probe->getIndex(), Count);
     FunctionProfile.addTotalSamples(Count);
     if (Probe->isEntry()) {
       FunctionProfile.addHeadSamples(Count);
for (auto *FunctionProfile : I.second) {
         FunctionProfile->addBodySamplesForProbe(Probe->getIndex(), 0);
       }
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
     const BranchSample &BranchCounter,
-    SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {
+    SmallVectorImpl<std::string> &ContextStrStack) {
   for (auto BI : BranchCounter) {
     uint64_t SourceOffset = BI.first.first;
     uint64_t TargetOffset = BI.first.second;
     uint64_t Count = BI.second;
     uint64_t SourceAddress = Binary->offsetToVirtualAddr(SourceOffset);
     const MCDecodedPseudoProbe *CallProbe =
         Binary->getCallProbeForAddr(SourceAddress);
     if (CallProbe == nullptr)
       continue;
     FunctionSamples &FunctionProfile =
-        getFunctionProfileForLeafProbe(ContextStrStack, CallProbe, Binary);
+        getFunctionProfileForLeafProbe(ContextStrStack, CallProbe);
     FunctionProfile.addBodySamples(CallProbe->getIndex(), 0, Count);
     FunctionProfile.addTotalSamples(Count);
     StringRef CalleeName = FunctionSamples::getCanonicalFnName(
         Binary->getFuncFromStartOffset(TargetOffset));
     if (CalleeName.size() == 0)
       continue;
     FunctionProfile.addCalledTargetSamples(CallProbe->getIndex(), 0, CalleeName,
                                            Count);
FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
   FunctionSamples &FunctionProile =
       getFunctionProfileForContext(OContextStr.str(), WasLeafInlined);
   FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
   return FunctionProile;
 }
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
     SmallVectorImpl<std::string> &ContextStrStack,
-    const MCDecodedPseudoProbe *LeafProbe, ProfiledBinary *Binary) {
+    const MCDecodedPseudoProbe *LeafProbe) {
 
   // Explicitly copy the context for appending the leaf context
   SmallVector<std::string, 16> ContextStrStackCopy(ContextStrStack.begin(),
                                                    ContextStrStack.end());
   Binary->getInlineContextForProbe(LeafProbe, ContextStrStackCopy, true);
   const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->getGuid());
   bool WasLeafInlined = LeafProbe->getInlineTreeNode()->hasInlineSite();
   return getFunctionProfileForLeafProbe(ContextStrStackCopy, FuncDesc,
                                         WasLeafInlined);
 }
 
 } // end namespace sampleprof
 } // end namespace llvm

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

 
 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::list<std::string>
-    BinaryFilenames("binary", cl::value_desc("binary"), cl::OneOrMore,
-                    llvm::cl::MiscFlags::CommaSeparated,
-                    cl::desc("Path of profiled binary files"),
-                    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."),
 
 extern cl::opt<bool> ShowDisassemblyOnly;
 
 using namespace llvm;
 using namespace sampleprof;
 
 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");
 
   if (ShowDisassemblyOnly) {
-    for (auto BinaryPath : BinaryFilenames) {
-      (void)ProfiledBinary(BinaryPath);
-    }
     return EXIT_SUCCESS;
   }
 
   // Load binaries and parse perf events and samples
   std::unique_ptr<PerfReaderBase> Reader =
-      PerfReaderBase::create(BinaryFilenames, PerfTraceFilenames);
+      PerfReaderBase::create(BinaryPath, PerfTraceFilenames);
   Reader->parsePerfTraces(PerfTraceFilenames);
 
-  std::unique_ptr<ProfileGenerator> Generator = ProfileGenerator::create(
-      Reader->getBinarySampleCounters(), Reader->getPerfScriptType());
+  std::unique_ptr<ProfileGenerator> Generator =
+      ProfileGenerator::create(Reader->getBinary(), Reader->getSampleCounters(),
+                               Reader->getPerfScriptType());
   Generator->generateProfile();
   Generator->write();
 
   return EXIT_SUCCESS;
 }