[CSSPGO][llvm-profgen] Refactor to unify hashable interface for trace sample and context-sensitive counter

[wlei]

Summary

As we plan to support both CSSPGO and AutoFDO for llvm-profgen, we will have different kinds of perf sample and different kinds of sample counter(cs/non-cs, with/without pseudo probe) which both need to do aggregation in hash map. This change implements the hashable interface(Hashable) and the unified base class for them to have better extensibility and reusability.

Currently perf trace sample and sample counter with context implemented this Hashable and the class hierarchy is like:

| Hashable  
           | PerfSample
                          | HybridSample
                          | LBRSample
           | ContextKey
                          | StringBasedCtxKey
                          | ProbeBasedCtxKey
                          | CallsiteBasedCtxKey
           | ...

For perf sample, HybridSample includes the call stack with LBR stack and LBRSample only includes LBR stack.
For context key used for virtual unwinding counter aggregation, we use string based context(StringBasedCtxKey) by default for good debug experience. For pseudo probe, we switch to use a stack of probe pointer(ProbeBasedCtxKey) to avoid redundant string handling. In the future, we can also speed up StringBasedCtxKey to use the original function frame stack, which is here named CallsiteBasedCtxKey

Implementation

• Class specifying Hashable should implement getHashCode and isEqual. Here we make getHashCode a non-virtual function to avoid vtable overhead, so derived class should calculate and assign the base class's HashCode manually. This also provides the flexibility for calculating the hash code incrementally(like rolling hash) during frame stack unwinding
• isEqual is a virtual function, which will have perf overhead. In the future, if we redesign a better hash function, then we can just skip this or switch to non-virtual function.
• Added PerfSample and ContextKey as base class for perf sample and counter context key, leveraging llvm-style RTTI for this.
• Added StringBasedCtxKey class extending ContextKey to use string as context id.
• Refactor AggregationCounter to take all kinds of PerfSample as key
• Refactor ContextSampleCounter to take all kinds of ContextKey as key
• Other refactoring work:
  • Create a wrapper class SampleCounter to wrap RangeCounter and BranchCounter
  • Hoist ContextId and FunctionProfile out of populateFunctionBodySamples and populateFunctionBoundarySamples to reuse them in ProfileGenerator

[wlei]

add more comments

[wlei]

fix a missing genHashCode, also add a assertion for this

[wlei]

rebase

[hoy]

llvm/tools/llvm-profgen/PerfReader.h

104

Nit: could this be defined as an overloaded == operator?

[wlei]

llvm/tools/llvm-profgen/PerfReader.h

104

== should work. Here using Equal is intentional, I want to diff from == to indicate that this is a 'Equal' function for hash not for the regular compare. Currently we exactly compare the elements like ==, but maybe in the future if we have a good hash function, we can use the hash code for the comparison. Or we can have different custom Equal function, like ContextKeyEqual, SampleEqual. Or maybe the 'Equal' is not a proper naming. Any thoughts on this?

[hoy]

llvm/tools/llvm-profgen/PerfReader.h

104

Sounds good. It makes sense to name the comparison more specifically.

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

76

Name it getOrCreateSampleCounter?

llvm/tools/llvm-profgen/PerfReader.h

110

I don't see a consumer of the two APIs. Maybe exclude them from this patch? I'm wondering if the shared pointer should be returned (for reference counting) when there is a need of exposing the underlying data.

[wmi]

llvm/tools/llvm-profgen/PerfReader.h

116

263

Can you use Hashable as a base class instead so you can remove the isEqual virtual function (The CRTP pattern: https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern#Static_polymorphism)? It will also simplify the interface using Hashable.

struct StringBasedCtxKey : public Hashable<StringBasedCtxKey> {
...
}

And other type of Key later on:
struct ProbeBasedCtxKey : public Hashable<ProbeBasedCtxKey> {
...
}

[wmi]

llvm/tools/llvm-profgen/PerfReader.h

298

Can we rename it to ContextSampleCounterMap and its object to something like CtxCounterMap?

[wlei]

address reviewers' feedback

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

76

renamed

llvm/tools/llvm-profgen/PerfReader.h

110

Removed getData(), getPtr() is used when getting the key data, like PerfRead.cpp:232

if the shared pointer should be returned

Not sure const shared_ptr<Key> *K = dyn_cast<shared_ptr<Key>>(I.first.getPtr()); can work for this, let me try.
Also since it's the hashkey, it's only used like:

for(auto I : Map) {
    const Key *K = dyn_cast<Key>(I.first.getPtr());
    // use K
}

So K is always used inside of Map's scope.

263

Thanks for your suggestion and detailed example.
Here I want to put the different types derived from the same base class in the same container.
like :
struct StringBasedCtxKey : public CtxKey {
...
}
struct ProbeBasedCtxKey : public CtxKey {
...
}
Then I can use only one container unordered_map<CtxKey, ...> ContextSampleCounterMap . If use CRTP, I guess it need two containers. any ideas about this?

298

Good point, the value is actually a counter, renamed.

[wmi]

llvm/tools/llvm-profgen/PerfReader.h

263

Ok, I don't have better idea. A possible solution is described in https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern#Pitfalls, but that prevents Hashable class from being reused by PerfSample without duplication.

I am fine with your existing solution If no one else has better idea.

[wlei]

llvm/tools/llvm-profgen/PerfReader.h

263

Yeah, we want both PerfSample and ContextKey could implements a unified Hashable. Thank you.

[hoy]

llvm/tools/llvm-profgen/PerfReader.h

110

Current getPtr() implementation and its use looks good. I was thinking getPtr() might leak the underlying data out of the shared pointer scope when destroyed.

263

Both patterns have their specialties. The current solution looks a bit cleaner though it falls back to use virtual methods to solve hash collision. Hopefully that could be mitigated by using a more efficient hash code.

[wenlei]

Thanks for the refactoring and clean up, looks great!

LBRSample will be added when AutoFDO support is moved into llvm-profgen, right? For AutoFDO, we could use the same infrastructure except that context will always be empty.

What is CallsiteBasedCtxKey in the description? A stack of call site addresses? Brief comment for each of the leaf class that was mentioned in the description (even if not implemented) would be useful.

llvm/tools/llvm-profgen/PerfReader.cpp

216

Does this map also needs to own a copy of the context string?

nit: perhaps remove this comment below? determinism mentioned above is good enough.

This is due to a build failure on sanitizer build(asan/msan/ubsan)

llvm/tools/llvm-profgen/PerfReader.h

171–174

nit: make hashCombine a lambda function inside genHashCode if there's no other use?

192

nit: rename it AggregatedCounter? we have AggregatedSamples as names.

[wlei]

Address Wenlei's feedback

[wlei]

LBRSample will be added when AutoFDO support is moved into llvm-profgen, right? For AutoFDO, we could use the same infrastructure except that context will always be empty.

Good question, Yes, my initial thought is to decouple this by using separated LBRSample when it comes to AutoFDO, perhaps this's good for readability. I guess your suggestion is for the performance since we won't have virtual call. That's good. If so and we don't have other kinds of perf sample, we might don't need the base class PerfSample and can just name HybridSample to more general name(like PerfSample).

What is CallsiteBasedCtxKey in the description? A stack of call site addresses? Brief comment for each of the leaf class that was mentioned in the description (even if not implemented) would be useful.

Yes, it's a vector of original call site address(Callstack). I was thinking whether we could use this to replace the string based context key in unwinder stage, later in ProfileGenerator convert to string. Just an idea. Will add more comments to the summary part.

llvm/tools/llvm-profgen/PerfReader.cpp

216

Good point. Changed to StringRef and remove this comment.

[wenlei]

In D92584#2458810, @wlei wrote:

LBRSample will be added when AutoFDO support is moved into llvm-profgen, right? For AutoFDO, we could use the same infrastructure except that context will always be empty.

Good question, Yes, my initial thought is to decouple this by using separated LBRSample when it comes to AutoFDO, perhaps this's good for readability. I guess your suggestion is for the performance since we won't have virtual call. That's good. If so and we don't have other kinds of perf sample, we might don't need the base class PerfSample and can just name HybridSample to more general name(like PerfSample).

Right.. LBR sample can be a Hybrid sample with empty stack, in which case we don't need hierarchy. But we can refine and deal with that later.

What is CallsiteBasedCtxKey in the description? A stack of call site addresses? Brief comment for each of the leaf class that was mentioned in the description (even if not implemented) would be useful.

Yes, it's a vector of original call site address(Callstack). I was thinking whether we could use this to replace the string based context key in unwinder stage, later in ProfileGenerator convert to string. Just an idea. Will add more comments to the summary part.

[wlei]

Kindly ping @wmi ,
I have addressed feedbacks from Hongtao and Wenlei, want to know whether this's also good for you or any other questions.
Also for other pseudo probe patches(https://reviews.llvm.org/D92334, https://reviews.llvm.org/D92896, https://reviews.llvm.org/D92998). Thanks!

[wmi]

Sorry for the delay. LGTM.

llvm/tools/llvm-profgen/PerfReader.h

 //===-- PerfReader.h - 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TOOLS_LLVM_PROFGEN_PERFREADER_H
 #define LLVM_TOOLS_LLVM_PROFGEN_PERFREADER_H
 #include "ErrorHandling.h"
 #include "ProfiledBinary.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Regex.h"
 #include <fstream>
 #include <list>
 #include <map>
 #include <vector>
 
 using namespace llvm;
struct LBREntry {
   // An artificial branch stands for a series of consecutive branches starting
   // from the current binary with a transition through external code and
   // eventually landing back in the current binary.
   bool IsArtificial = false;
   LBREntry(uint64_t S, uint64_t T, bool I)
       : Source(S), Target(T), IsArtificial(I) {}
 };
 
+// Hash interface for generic data of type T
+// Data should implement a \fn getHashCode and a \fn isEqual
+// Currently getHashCode is non-virtual to avoid the overhead of calling vtable,
+// i.e we explicitly calculate hash of derived class, assign to base class's
+// HashCode. This also provides the flexibility for calculating the hash code
+// incrementally(like rolling hash) during frame stack unwinding since unwinding
+// only changes the leaf of frame stack. \fn isEqual is a virtual function,
+// which will have perf overhead. In the future, if we redesign a better hash
+// function, then we can just skip this or switch to non-virtual function(like
+// just ignore comparision if hash conflicts probabilities is low)
+template <class T> class Hashable {
+public:
+  std::shared_ptr<T> Data;
+  Hashable(const std::shared_ptr<T> &D) : Data(D) {}
+
+  // Hash code generation
+  struct Hash {
+    uint64_t operator()(const Hashable<T> &Key) const {
+      // Don't make it virtual for getHashCode
+      assert(Key.Data->getHashCode() && "Should generate HashCode for it!");
+      return Key.Data->getHashCode();
+    }
+  };
+
+  // Hash equal
+  struct Equal {
+    bool operator()(const Hashable<T> &LHS, const Hashable<T> &RHS) const {
+      // Precisely compare the data, vtable will have overhead.
+      return LHS.Data->isEqual(RHS.Data.get());
+    }
+  };
+
+  T *getPtr() const { return Data.get(); }
+};
+
+// Base class to extend for all types of perf sample
+struct PerfSample {
+  uint64_t HashCode = 0;
+
+  virtual ~PerfSample() = default;
+  uint64_t getHashCode() const { return HashCode; }
+  virtual bool isEqual(const PerfSample *K) const {
+    return HashCode == K->HashCode;
+  };
+
+  // Utilities for LLVM-style RTTI
+  enum PerfKind { PK_HybridSample };
+  const PerfKind Kind;
+  PerfKind getKind() const { return Kind; }
+  PerfSample(PerfKind K) : Kind(K){};
+};
+
 // The parsed hybrid sample including call stack and LBR stack.
-struct HybridSample {
+struct HybridSample : public PerfSample {
   // Profiled binary that current frame address belongs to
   ProfiledBinary *Binary;
   // Call stack recorded in FILO(leaf to root) order
   std::list<uint64_t> CallStack;
   // LBR stack recorded in FIFO order
   SmallVector<LBREntry, 16> LBRStack;
 
+  HybridSample() : PerfSample(PK_HybridSample){};
+  static bool classof(const PerfSample *K) {
+    return K->getKind() == PK_HybridSample;
+  }
+
   // Used for sample aggregation
-  bool operator==(const HybridSample &Other) const {
-    if (Other.Binary != Binary)
+  bool isEqual(const PerfSample *K) const override {
+    const HybridSample *Other = dyn_cast<HybridSample>(K);
+    if (Other->Binary != Binary)
       return false;
-    const std::list<uint64_t> &OtherCallStack = Other.CallStack;
-    const SmallVector<LBREntry, 16> &OtherLBRStack = Other.LBRStack;
+    const std::list<uint64_t> &OtherCallStack = Other->CallStack;
+    const SmallVector<LBREntry, 16> &OtherLBRStack = Other->LBRStack;
 
     if (CallStack.size() != OtherCallStack.size() ||
         LBRStack.size() != OtherLBRStack.size())
       return false;
 
     auto Iter = CallStack.begin();
     for (auto Address : OtherCallStack) {
       if (Address != *Iter++)
         return false;
     }
 
     for (size_t I = 0; I < OtherLBRStack.size(); I++) {
       if (LBRStack[I].Source != OtherLBRStack[I].Source ||
           LBRStack[I].Target != OtherLBRStack[I].Target)
         return false;
     }
     return true;
   }
+
+  void genHashCode() {
+    // Use simple DJB2 hash
+    auto HashCombine = [](uint64_t H, uint64_t V) {
+      return ((H << 5) + H) + V;
+    };
+    uint64_t Hash = 5381;
+    Hash = HashCombine(Hash, reinterpret_cast<uint64_t>(Binary));
+    for (const auto &Value : CallStack) {
+      Hash = HashCombine(Hash, Value);
+    }
+    for (const auto &Entry : LBRStack) {
+      Hash = HashCombine(Hash, Entry.Source);
+      Hash = HashCombine(Hash, Entry.Target);
+    }
+    HashCode = Hash;
+  }
 };
 
+// After parsing the sample, we record the samples by aggregating them
+// into this counter. The key stores the sample data and the value is
+// the sample repeat times.
+using AggregatedCounter =
+    std::unordered_map<Hashable<PerfSample>, uint64_t,
+                       Hashable<PerfSample>::Hash, Hashable<PerfSample>::Equal>;
+
 // The state for the unwinder, it doesn't hold the data but only keep the
 // pointer/index of the data, While unwinding, the CallStack is changed
 // dynamicially and will be recorded as the context of the sample
 struct UnwindState {
   // Profiled binary that current frame address belongs to
   const ProfiledBinary *Binary;
   // TODO: switch to use trie for call stack
   std::list<uint64_t> CallStack;
   // Used to fall through the LBR stack
   uint32_t LBRIndex = 0;
   // Reference to HybridSample.LBRStack
   const SmallVector<LBREntry, 16> &LBRStack;
   // Used to iterate the address range
   InstructionPointer InstPtr;
-  UnwindState(const HybridSample &Sample)
-      : Binary(Sample.Binary), CallStack(Sample.CallStack),
-        LBRStack(Sample.LBRStack),
-        InstPtr(Sample.Binary, Sample.CallStack.front()) {}
+  UnwindState(const HybridSample *Sample)
+      : Binary(Sample->Binary), CallStack(Sample->CallStack),
+        LBRStack(Sample->LBRStack),
+        InstPtr(Sample->Binary, Sample->CallStack.front()) {}
 
   bool validateInitialState() {
     uint64_t LBRLeaf = LBRStack[LBRIndex].Target;
     uint64_t StackLeaf = CallStack.front();
     // When we take a stack sample, ideally the sampling distance between the
     // leaf IP of stack and the last LBR target shouldn't be very large.
     // Use a heuristic size (0x100) to filter out broken records.
     if (StackLeaf < LBRLeaf || StackLeaf >= LBRLeaf + 0x100) {
struct UnwindState {
   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; }
   const LBREntry &getCurrentLBR() const { return LBRStack[LBRIndex]; }
   void advanceLBR() { LBRIndex++; }
 };
 
-// The counter of branch samples for one function indexed by the branch,
-// which is represented as the source and target offset pair.
-using BranchSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
-// The counter of range samples for one function indexed by the range,
-// which is represented as the start and end offset pair.
-using RangeSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
-// Range sample counters indexed by the context string
-using ContextRangeCounter = std::unordered_map<std::string, RangeSample>;
-// Branch sample counters indexed by the context string
-using ContextBranchCounter = std::unordered_map<std::string, BranchSample>;
+// Base class for sample counter key with context
+struct ContextKey {
+  uint64_t HashCode = 0;
+  virtual ~ContextKey() = default;
+  uint64_t getHashCode() const { return HashCode; }
+  virtual bool isEqual(const ContextKey *K) const {
+    return HashCode == K->HashCode;
+  };
 
-// For Hybrid sample counters
-struct ContextSampleCounters {
-  ContextRangeCounter RangeCounter;
-  ContextBranchCounter BranchCounter;
+  // Utilities for LLVM-style RTTI
+  enum ContextKind { CK_StringBased };
+  const ContextKind Kind;
+  ContextKind getKind() const { return Kind; }
+  ContextKey(ContextKind K) : Kind(K){};
+};
 
-  void recordRangeCount(std::string &ContextId, uint64_t Start, uint64_t End,
-                        uint64_t Repeat) {
-    RangeCounter[ContextId][{Start, End}] += Repeat;
+// String based context id
+struct StringBasedCtxKey : public ContextKey {
+  std::string Context;
+  StringBasedCtxKey() : ContextKey(CK_StringBased){};
+  static bool classof(const ContextKey *K) {
+    return K->getKind() == CK_StringBased;
   }
-  void recordBranchCount(std::string &ContextId, uint64_t Source,
-                         uint64_t Target, uint64_t Repeat) {
-    BranchCounter[ContextId][{Source, Target}] += Repeat;
+
+  bool isEqual(const ContextKey *K) const override {
+    const StringBasedCtxKey *Other = dyn_cast<StringBasedCtxKey>(K);
+    return Context == Other->Context;
   }
+
+  void genHashCode() { HashCode = hash_value(Context); }
 };
 
-struct HybridSampleHash {
-  uint64_t hashCombine(uint64_t Hash, uint64_t Value) const {
-    // Simple DJB2 hash
-    return ((Hash << 5) + Hash) + Value;
-  }
+// The counter of branch samples for one function indexed by the branch,
+// which is represented as the source and target offset pair.
+using BranchSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
+// The counter of range samples for one function indexed by the range,
+// which is represented as the start and end offset pair.
+using RangeSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
+// Wrapper for sample counters including range counter and branch counter
+struct SampleCounter {
+  RangeSample RangeCounter;
+  BranchSample BranchCounter;
 
-  uint64_t operator()(const HybridSample &Sample) const {
-    uint64_t Hash = 5381;
-    Hash = hashCombine(Hash, reinterpret_cast<uint64_t>(Sample.Binary));
-    for (const auto &Value : Sample.CallStack) {
-      Hash = hashCombine(Hash, Value);
-    }
-    for (const auto &Entry : Sample.LBRStack) {
-      Hash = hashCombine(Hash, Entry.Source);
-      Hash = hashCombine(Hash, Entry.Target);
-    }
-    return Hash;
+  void recordRangeCount(uint64_t Start, uint64_t End, uint64_t Repeat) {
+    RangeCounter[{Start, End}] += Repeat;
+  void recordBranchCount(uint64_t Source, uint64_t Target, uint64_t Repeat) {
+    BranchCounter[{Source, Target}] += Repeat;
   }
 };
 
-// After parsing the sample, we record the samples by aggregating them
-// into this structure and the value is the sample counter.
-using AggregationCounter =
-    std::unordered_map<HybridSample, uint64_t, HybridSampleHash>;
+// Sample counter with context to support context-sensitive profile
+using ContextSampleCounterMap =
+    std::unordered_map<Hashable<ContextKey>, SampleCounter,
+                       Hashable<ContextKey>::Hash, Hashable<ContextKey>::Equal>;
 
 /*
 As in hybrid sample we have a group of LBRs and the most recent sampling call
 stack, we can walk through those LBRs to infer more call stacks which would be
 used as context for profile. VirtualUnwinder is the class to do the call stack
 unwinding based on LBR state. Two types of unwinding are processd here:
 1) LBR unwinding and 2) linear range unwinding.
 Specifically, for each LBR entry(can be classified into call, return, regular
 branch), LBR unwinding will replay the operation by pushing, popping or
 switching leaf frame towards the call stack and since the initial call stack
 is most recently sampled, the replay should be in anti-execution order, i.e. for
 the regular case, pop the call stack when LBR is call, push frame on call stack
 when LBR is return. After each LBR processed, it also needs to align with the
 next LBR by going through instructions from previous LBR's target to current
 LBR's source, which is the linear unwinding. As instruction from linear range
 can come from different function by inlining, linear unwinding will do the range
 splitting and record counters by the range with same inline context. Over those
 unwinding process we will record each call stack as context id and LBR/linear
 range as sample counter for further CS profile generation.
 */
 class VirtualUnwinder {
 public:
-  VirtualUnwinder(ContextSampleCounters *Counters) : SampleCounters(Counters) {}
+  VirtualUnwinder(ContextSampleCounterMap *Counter) : CtxCounterMap(Counter) {}
 
   bool isCallState(UnwindState &State) const {
     // The tail call frame is always missing here in stack sample, we will
     // use a specific tail call tracker to infer it.
     return State.getBinary()->addressIsCall(State.getCurrentLBRSource());
   }
 
   bool isReturnState(UnwindState &State) const {
     // Simply check addressIsReturn, as ret is always reliable, both for
     // regular call and tail call.
     return State.getBinary()->addressIsReturn(State.getCurrentLBRSource());
   }
 
   void unwindCall(UnwindState &State);
   void unwindLinear(UnwindState &State, uint64_t Repeat);
   void unwindReturn(UnwindState &State);
   void unwindBranchWithinFrame(UnwindState &State);
-  bool unwind(const HybridSample &Sample, uint64_t Repeat);
+  bool unwind(const HybridSample *Sample, uint64_t Repeat);
   void recordRangeCount(uint64_t Start, uint64_t End, UnwindState &State,
                         uint64_t Repeat);
   void recordBranchCount(const LBREntry &Branch, UnwindState &State,
                          uint64_t Repeat);
+  SampleCounter &getOrCreateSampleCounter(const ProfiledBinary *Binary,
+                                          std::list<uint64_t> &CallStack);
 
 private:
-  ContextSampleCounters *SampleCounters;
+  ContextSampleCounterMap *CtxCounterMap;
 };
 
 // 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 *, ContextSampleCounters>;
+    std::unordered_map<ProfiledBinary *, ContextSampleCounterMap>;
 
 // Load binaries and read perf trace to parse the events and samples
 class PerfReader {
 
 public:
   PerfReader(cl::list<std::string> &BinaryFilenames);
 
   // Hybrid sample(call stack + LBRs) profile traces are seprated by double line
private:
   ProfiledBinary *getBinary(uint64_t Address);
 
   BinaryMap BinaryTable;
   AddressBinaryMap AddrToBinaryMap; // Used by address-based lookup.
 
 private:
   BinarySampleCounterMap BinarySampleCounters;
   // Samples with the repeating time generated by the perf reader
-  AggregationCounter AggregatedSamples;
+  AggregatedCounter AggregatedSamples;
   PerfScriptType PerfType;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PerfReader.cpp

 void VirtualUnwinder::unwindBranchWithinFrame(UnwindState &State) {
   // TODO: Tolerate tail call for now, as we may see tail call from libraries.
   // This is only for intra function branches, excluding tail calls.
   uint64_t Source = State.getCurrentLBRSource();
   State.CallStack.front() = Source;
   State.InstPtr.update(Source);
 }
 
+SampleCounter &
+VirtualUnwinder::getOrCreateSampleCounter(const ProfiledBinary *Binary,
+                                          std::list<uint64_t> &CallStack) {
+  std::shared_ptr<StringBasedCtxKey> KeyStr =
+      std::make_shared<StringBasedCtxKey>();
+  KeyStr->Context = Binary->getExpandedContextStr(CallStack);
+  KeyStr->genHashCode();
+  auto Ret =
+      CtxCounterMap->emplace(Hashable<ContextKey>(KeyStr), SampleCounter());
+  return Ret.first->second;
+}
+
 void VirtualUnwinder::recordRangeCount(uint64_t Start, uint64_t End,
                                        UnwindState &State, uint64_t Repeat) {
-  std::string &&ContextId = State.getExpandedContextStr();
   uint64_t StartOffset = State.getBinary()->virtualAddrToOffset(Start);
   uint64_t EndOffset = State.getBinary()->virtualAddrToOffset(End);
-  SampleCounters->recordRangeCount(ContextId, StartOffset, EndOffset, Repeat);
+  SampleCounter &SCounter =
+      getOrCreateSampleCounter(State.getBinary(), State.CallStack);
+  SCounter.recordRangeCount(StartOffset, EndOffset, Repeat);
 }
 
 void VirtualUnwinder::recordBranchCount(const LBREntry &Branch,
                                         UnwindState &State, uint64_t Repeat) {
   if (Branch.IsArtificial)
     return;
-  std::string &&ContextId = State.getExpandedContextStr();
   uint64_t SourceOffset = State.getBinary()->virtualAddrToOffset(Branch.Source);
   uint64_t TargetOffset = State.getBinary()->virtualAddrToOffset(Branch.Target);
-  SampleCounters->recordBranchCount(ContextId, SourceOffset, TargetOffset,
-                                    Repeat);
+  SampleCounter &SCounter =
+      getOrCreateSampleCounter(State.getBinary(), State.CallStack);
+  SCounter.recordBranchCount(SourceOffset, TargetOffset, Repeat);
 }
 
-bool VirtualUnwinder::unwind(const HybridSample &Sample, uint64_t Repeat) {
+bool VirtualUnwinder::unwind(const HybridSample *Sample, uint64_t Repeat) {
   // Capture initial state as starting point for unwinding.
   UnwindState State(Sample);
 
   // Sanity check - making sure leaf of LBR aligns with leaf of stack sample
   // Stack sample sometimes can be unreliable, so filter out bogus ones.
   if (!State.validateInitialState())
     return false;
 
ProfiledBinary *PerfReader::getBinary(uint64_t Address) {
   if (Iter == AddrToBinaryMap.end() || Iter->first != Address) {
     if (Iter == AddrToBinaryMap.begin())
       return nullptr;
     Iter--;
   }
   return Iter->second;
 }
 
-static void printSampleCounter(ContextRangeCounter &Counter) {
-  // Use ordered map to make the output deterministic
-  std::map<std::string, RangeSample> OrderedCounter(Counter.begin(),
-                                                    Counter.end());
+using OrderedCounterForPrint = std::map<StringRef, 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) << ", "
              << format("%" PRIx64, I.first.second) << "): " << I.second << "\n";
     }
   }
 }
 
+static void printRangeCounter(ContextSampleCounterMap &Counter) {
+  OrderedCounterForPrint OrderedCounter;
+  for (auto &CI : Counter) {
+    const StringBasedCtxKey *CtxKey =
+        dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
+    OrderedCounter[CtxKey->Context] = CI.second.RangeCounter;
+  }
+  printSampleCounter(OrderedCounter);
+}
+
+static void printBranchCounter(ContextSampleCounterMap &Counter) {
+  OrderedCounterForPrint OrderedCounter;
+  for (auto &CI : Counter) {
+    const StringBasedCtxKey *CtxKey =
+        dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
+    OrderedCounter[CtxKey->Context] = CI.second.BranchCounter;
+  }
+  printSampleCounter(OrderedCounter);
+}
+
 void PerfReader::printUnwinderOutput() {
   for (auto I : BinarySampleCounters) {
     const ProfiledBinary *Binary = I.first;
     outs() << "Binary(" << Binary->getName().str() << ")'s Range Counter:\n";
-    printSampleCounter(I.second.RangeCounter);
+    printRangeCounter(I.second);
     outs() << "\nBinary(" << Binary->getName().str() << ")'s Branch Counter:\n";
-    printSampleCounter(I.second.BranchCounter);
+    printBranchCounter(I.second);
   }
 }
 
 void PerfReader::unwindSamples() {
   for (const auto &Item : AggregatedSamples) {
-    const HybridSample &Sample = Item.first;
-    VirtualUnwinder Unwinder(&BinarySampleCounters[Sample.Binary]);
+    const HybridSample *Sample = dyn_cast<HybridSample>(Item.first.getPtr());
+    VirtualUnwinder Unwinder(&BinarySampleCounters[Sample->Binary]);
     Unwinder.unwind(Sample, Item.second);
   }
 
   if (ShowUnwinderOutput)
     printUnwinderOutput();
 }
 
 bool PerfReader::extractLBRStack(TraceStream &TraceIt,
void PerfReader::parseHybridSample(TraceStream &TraceIt) {
   // 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
   //
-  HybridSample Sample;
+  std::shared_ptr<HybridSample> Sample = std::make_shared<HybridSample>();
 
   // Parsing call stack and populate into HybridSample.CallStack
-  if (!extractCallstack(TraceIt, Sample.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
-  Sample.Binary = getBinary(Sample.CallStack.front());
+  Sample->Binary = getBinary(Sample->CallStack.front());
 
   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, Sample->Binary)) {
       // Canonicalize stack leaf to avoid 'random' IP from leaf frame skew LBR
       // ranges
-      Sample.CallStack.front() = Sample.LBRStack[0].Target;
+      Sample->CallStack.front() = Sample->LBRStack[0].Target;
       // Record samples by aggregation
-      AggregatedSamples[Sample]++;
+      Sample->genHashCode();
+      AggregatedSamples[Hashable<PerfSample>(Sample)]++;
     }
   } else {
     // LBR sample is encoded in single line after stack sample
     exitWithError("'Hybrid perf sample is corrupted, No LBR sample line");
   }
 }
 
 void PerfReader::parseMMap2Event(TraceStream &TraceIt) {

llvm/tools/llvm-profgen/ProfileGenerator.h

class CSProfileGenerator : public ProfileGenerator {
   const BinarySampleCounterMap &BinarySampleCounters;
 
 public:
   CSProfileGenerator(const BinarySampleCounterMap &Counters)
       : BinarySampleCounters(Counters){};
 
 public:
   void generateProfile() override {
-    // Fill in function body samples
-    populateFunctionBodySamples();
+    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);
 
+        // 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();
-
+        populateFunctionBoundarySamples(ContextId, FunctionProfile,
+                                        CI.second.BranchCounter, Binary);
+      }
+    }
     // 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();
   }
 
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
   // Lookup or create FunctionSamples for the context
   FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
-  void populateFunctionBodySamples();
-  void populateFunctionBoundarySamples();
+  void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
+                                   const RangeSample &RangeCounters,
+                                   ProfiledBinary *Binary);
+  void populateFunctionBoundarySamples(StringRef ContextId,
+                                       FunctionSamples &FunctionProfile,
+                                       const BranchSample &BranchCounters,
+                                       ProfiledBinary *Binary);
   void populateInferredFunctionSamples();
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

void CSProfileGenerator::updateBodySamplesforFunctionProfile(
   if (PreviousCount < Count) {
     FunctionProfile.addBodySamples(LeafLoc.second.LineOffset,
                                    LeafLoc.second.Discriminator,
                                    Count - PreviousCount);
     FunctionProfile.addTotalSamples(Count - PreviousCount);
   }
 }
 
-void CSProfileGenerator::populateFunctionBodySamples() {
-  for (const auto &BI : BinarySampleCounters) {
-    ProfiledBinary *Binary = BI.first;
-    for (const auto &CI : BI.second.RangeCounter) {
-      StringRef ContextId(CI.first);
-      // Get or create function profile for the range
-      FunctionSamples &FunctionProfile =
-          getFunctionProfileForContext(ContextId);
-      // Compute disjoint ranges first, so we can use MAX
-      // for calculating count for each location.
-      RangeSample Ranges;
-      findDisjointRanges(Ranges, CI.second);
-
-      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;
+void CSProfileGenerator::populateFunctionBodySamples(
+    FunctionSamples &FunctionProfile, const RangeSample &RangeCounter,
+    ProfiledBinary *Binary) {
+  findDisjointRanges(Ranges, RangeCounter);
 
-        InstructionPointer IP(Binary, 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) {
-          uint64_t Offset = Binary->virtualAddrToOffset(IP.Address);
-          const FrameLocation &LeafLoc = Binary->getInlineLeafFrameLoc(Offset);
-          // Recording body sample for this specific context
-          updateBodySamplesforFunctionProfile(FunctionProfile, LeafLoc, Count);
-          // Move to next IP within the range
-          IP.advance();
-        }
-      }
-    }
-  }
-}
 
-void CSProfileGenerator::populateFunctionBoundarySamples() {
-  for (const auto &BI : BinarySampleCounters) {
-    ProfiledBinary *Binary = BI.first;
-    for (const auto &CI : BI.second.BranchCounter) {
-      StringRef ContextId(CI.first);
-      // Get or create function profile for branch Source
-      FunctionSamples &FunctionProfile =
-          getFunctionProfileForContext(ContextId);
+void CSProfileGenerator::populateFunctionBoundarySamples(
+    StringRef ContextId, FunctionSamples &FunctionProfile,
+    const BranchSample &BranchCounters, ProfiledBinary *Binary) {
 
-      for (auto Entry : CI.second) {
-        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));
-        if (CalleeName.size() == 0)
-          continue;
+  for (auto Entry : BranchCounters) {
 
-        // Record called target sample and its count
-        const FrameLocation &LeafLoc =
-            Binary->getInlineLeafFrameLoc(SourceOffset);
+    const FrameLocation &LeafLoc = Binary->getInlineLeafFrameLoc(SourceOffset);
 
-        FunctionProfile.addCalledTargetSamples(LeafLoc.second.LineOffset,
-                                               LeafLoc.second.Discriminator,
-                                               CalleeName, Count);
-        FunctionProfile.addTotalSamples(Count);
 
-        // Record head sample for called target(callee)
-        // TODO: Cleanup ' @ '
-        std::string CalleeContextId =
-            getCallSite(LeafLoc) + " @ " + CalleeName.str();
-        if (ContextId.find(" @ ") != StringRef::npos) {
-          CalleeContextId =
-              ContextId.rsplit(" @ ").first.str() + " @ " + CalleeContextId;
-        }
 
-        if (ProfileMap.find(CalleeContextId) != ProfileMap.end()) {
-          FunctionSamples &CalleeProfile = ProfileMap[CalleeContextId];
-          assert(Count != 0 && "Unexpected zero weight branch");
-          if (CalleeProfile.getName().size()) {
-            CalleeProfile.addHeadSamples(Count);
-          }
-        }
-      }
-    }
-  }
-}
+    FunctionSamples &CalleeProfile =
+        getFunctionProfileForContext(CalleeContextId);
 
 static FrameLocation getCallerContext(StringRef CalleeContext,
                                       StringRef &CallerNameWithContext) {
   StringRef CallerContext = CalleeContext.rsplit(" @ ").first;
   CallerNameWithContext = CallerContext.rsplit(':').first;
   auto ContextSplit = CallerContext.rsplit(" @ ");
   FrameLocation LeafFrameLoc = {"", {0, 0}};
   StringRef Funcname;