[CSSPGO][llvm-profgen] Fix external address issues of perf reader (return to external addr part)

[wlei]

Summary

Before we have an issue with artificial LBR whose source is a return, recalling that "an internal code(A) can return to external address, then from the external address call a new internal code(B), making an artificial branch that looks like a return from A to B can confuse the unwinder". We just ignore the LBRs after this artificial LBR which can miss some samples. This change aims at fixing this by correctly unwinding them instead of ignoring them.

List some typical scenarios covered by this change.

1. multiple sequential call back happen in external address, e.g.

[ext, call, foo] [foo, return, ext] [ext, call, bar]

Unwinder should avoid having foo return from bar. Wrong call stack is like [foo, bar]

2. the call stack before and after external call should be correctly unwinded.

{call stack1}                                            {call stack2}  
[foo, call, ext]  [ext, call, bar]  [bar, return, ext]  [ext, return, foo ]

call stack 1 should be the same to call stack2. Both shouldn't be truncated

3. call stack should be truncated after call into external code since we can't do inlining with external code.

[foo, call, ext]  [ext, call, bar]  [bar, call, baz] [baz, return, bar ] [bar, return, ext]

the call stack of code in baz should not include foo.

Implementation:

We leverage artificial frame to fix #2 and #3: when we got a return artificial LBR, push an extra artificial frame to the stack. when we pop frame, check if the parent is an artificial frame to pop(fix #2). Therefore, call/ return artificial LBR is just the same as regular LBR which can keep the call stack.

While recording context on the trie, artificial frame is used as a tag indicating that we should truncate the call stack(fix #3).

To differentiate #1 and #2, we leverage getCallAddrFromFrameAddr. Normally the target of the return should be the next inst of a call inst and getCallAddrFromFrameAddr will return the address of call inst. Otherwise, getCallAddrFromFrameAddr will return to 0 which is the case of #1.

[wlei]

Updating D115550: [CSSPGO][llvm-profgen] Fix external address issues of perf reader (return to external addr part)

[wenlei]

llvm/tools/llvm-profgen/PerfReader.cpp

147–159

Can we move this check into isReturnState? This should not be considered as return state, then we would naturally fall back to unwindBranchWithinFrame which does what we want. We may need to rename unwindBranchWithinFrame to simply unwindBranch because this now can be inter-function too.

llvm/tools/llvm-profgen/PerfReader.h

296

Comment for the definition and purpose for artificial frame.

423

nit: if we don't expect pushing artificial frame, the message can be explicit, something like "Artificial frame not expected for context stack."

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

222–224

nit: ... since this isn't a real call context the compiler will see

llvm/tools/llvm-profgen/PerfReader.h

309

nit: switch the two statements for readability? A comment would also be helpful.

423

Or we just return false for artificial frames so that it doesn't need to be handled on the caller side?

[wenlei]

llvm/tools/llvm-profgen/PerfReader.h

309

iiuc, switching changes the semantic though - we will be checking IsArtificialFrame the grand parent instead of parent ..

[hoy]

llvm/tools/llvm-profgen/PerfReader.h

309

Right, the code needs to be slightly tweaked. I was thinking of reading this routine like we unwind once first. When an artificial frame is found, we unwind once more to skip that frame. Not a strong preference though.

[wlei]

addressing reviewer's feedback

[wlei]

refine some comments

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

147–159

That's a good point! moved to isReturnState

222–224

Moved to pushFrame

llvm/tools/llvm-profgen/PerfReader.h

296

changed to reuse external frame instead of artificail frame. Added some comments near it's definition.

309

ok, I realized that moving this check`if (CurrentLeafFrame->Parent->IsArtificialFrame())` to its caller seems better, then we can have better handling the unmatched cases.

423

Sounds good, changed to return false for external frame.

[wenlei]

llvm/tools/llvm-profgen/PerfReader.cpp

147

replace "artificial frame" with "external frame" now that we're converging on the term?

llvm/tools/llvm-profgen/PerfReader.h

423

I actually think that it's cleaner to handle this on caller side and assert we don't have external frame when it comes here.

525–527

nit: return (CallAddr != 0); to make the boolean return explicit.

[hoy]

llvm/tools/llvm-profgen/PerfReader.h

423

Sounds good to me.

[wlei]

addressing feedbacks.

[wenlei]

lgtm, thanks.

[hoy]

lgtm, thx.

llvm/test/tools/llvm-profgen/Inputs/callback-external-addr.perfscript

+      7fe8d7620597
+            4007b0
+      7fe8d727e493
+  5541f689495641d7
+ 0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1  0x40072b/0x40074c/P/-/-/5  0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/7  0x7fe8d76205a3/0x400715/P/-/-/1  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/3  0x7fe8d7620589/0x400690/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/10  0x4006be/0x4006ec/P/-/-/2  0x4006e7/0x4006b0/P/-/-/3  0x400747/0x4006d0/P/-/-/2  0x7fe8d7620589/0x400730/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x4007ab/0x400590/P/-/-/2  0x4007bf/0x40077d/P/-/-/1  0x7fe8d76205a3/0x4007b0/P/-/-/2  0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1  0x40072b/0x40074c/P/-/-/4  0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/3  0x7fe8d76205a3/0x400715/P/-/-/4  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/3  0x7fe8d7620589/0x400690/P/-/-/4
+
+            4006ec
+            40074c
+      7fe8d762058b
+            4007b0
+      7fe8d727e493
+  5541f689495641d7
+ 0x4006be/0x4006ec/P/-/-/2  0x4006e7/0x4006b0/P/-/-/3  0x400747/0x4006d0/P/-/-/2  0x7fe8d7620589/0x400730/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x4007ab/0x400590/P/-/-/2  0x4007bf/0x40077d/P/-/-/3  0x7fe8d76205a3/0x4007b0/P/-/-/2  0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1  0x40072b/0x40074c/P/-/-/7  0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/2  0x7fe8d76205a3/0x400715/P/-/-/1  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/4  0x7fe8d7620589/0x400690/P/-/-/6  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/5  0x4006be/0x4006ec/P/-/-/3  0x4006e7/0x4006b0/P/-/-/3  0x400747/0x4006d0/P/-/-/2  0x7fe8d7620589/0x400730/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x4007ab/0x400590/P/-/-/2  0x4007bf/0x40077d/P/-/-/3  0x7fe8d76205a3/0x4007b0/P/-/-/2  0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1
+
+            40074c
+      7fe8d762058b
+            4007b0
+      7fe8d727e493
+  5541f689495641d7
+ 0x40072b/0x40074c/P/-/-/6  0x4006ce/0x400720/P/-/-/8  0x40071b/0x4006c0/P/-/-/1  0x7fe8d76205a3/0x400715/P/-/-/2  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/1  0x40069e/0x7fe8d762058b/P/-/-/2  0x7fe8d7620589/0x400690/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/2  0x4006be/0x4006ec/P/-/-/2  0x4006e7/0x4006b0/P/-/-/3  0x400747/0x4006d0/P/-/-/2  0x7fe8d7620589/0x400730/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x4007ab/0x400590/P/-/-/2  0x4007bf/0x40077d/P/-/-/1  0x7fe8d76205a3/0x4007b0/P/-/-/2  0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1  0x40072b/0x40074c/P/-/-/4  0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/10  0x7fe8d76205a3/0x400715/P/-/-/1  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/2  0x7fe8d7620589/0x400690/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/6  0x4006be/0x4006ec/P/-/-/4
+
+            400720
+            40074c
+      7fe8d762058b
+            4007b0
+      7fe8d727e493
+  5541f689495641d7
+ 0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/3  0x7fe8d76205a3/0x400715/P/-/-/1  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/7  0x7fe8d7620589/0x400690/P/-/-/5  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/5  0x4006be/0x4006ec/P/-/-/3  0x4006e7/0x4006b0/P/-/-/3  0x400747/0x4006d0/P/-/-/2  0x7fe8d7620589/0x400730/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x4007ab/0x400590/P/-/-/2  0x4007bf/0x40077d/P/-/-/2  0x7fe8d76205a3/0x4007b0/P/-/-/2  0x40069e/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x400690/P/-/-/1  0x400751/0x7fe8d762058b/P/-/-/1  0x40072b/0x40074c/P/-/-/4  0x4006ce/0x400720/P/-/-/4  0x40071b/0x4006c0/P/-/-/2  0x7fe8d76205a3/0x400715/P/-/-/2  0x4006ae/0x7fe8d7620597/P/-/-/2  0x7fe8d7620595/0x4006a0/P/-/-/3  0x40069e/0x7fe8d762058b/P/-/-/2  0x7fe8d7620589/0x400690/P/-/-/4  0x400590/0x7fe8d7620560/P/-/-/1  0x400710/0x400590/P/-/-/5  0x4006be/0x4006ec/P/-/-/2  0x4006e7/0x4006b0/P/-/-/3

llvm/test/tools/llvm-profgen/callback-external-addr.test

+; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/callback-external-addr.perfscript --binary=%S/Inputs/callback-external-addr.perfbin --output=%t --skip-symbolization
+; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-UNWINDER
+
+; Test if call stack is correctly truncated.
+; CHECK-UNWINDER-NOT: main:3 @ bar
+; CHECK-UNWINDER-NOT: main:3 @ foo
+; CHECK-UNWINDER-NOT: qux:3 @ baz
+; CHECK-UNWINDER-NOT: qux:3 @ bar
+
+; Test if return to wrong internal target
+; CHECK-UNWINDER-NOT: baz:0 @ bar
+; CHECK-UNWINDER-NOT: bar:0 @ baz
+; CHECK-UNWINDER-NOT: baz:0 @ main
+; CHECK-UNWINDER-NOT: bar:0 @ foo
+; CHECK-UNWINDER-NOT: baz:0 @ qux
+
+; Test for callback return from internal address to external address.
+; [foo:2 @ qux:2 @ callBeforeReturn] and [foo:2 @ qux:4 @ callAfterReturn] should exist
+; which means the callback return won't interrupt the previous call stack
+
+; CHECK-UNWINDER: [bar]
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   690-69e:12
+; CHECK-UNWINDER:   0
+; CHECK-UNWINDER: [baz]
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   6a0-6ae:7
+; CHECK-UNWINDER:   0
+; CHECK-UNWINDER: [foo]
+; CHECK-UNWINDER:   2
+; CHECK-UNWINDER:   730-747:5
+; CHECK-UNWINDER:   74c-751:5
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   747->6d0:5
+; CHECK-UNWINDER: [foo:2 @ qux]
+; CHECK-UNWINDER:   4
+; CHECK-UNWINDER:   6d0-6e7:5
+; CHECK-UNWINDER:   6ec-710:6
+; CHECK-UNWINDER:   715-71b:7
+; CHECK-UNWINDER:   720-72b:6
+; CHECK-UNWINDER:   3
+; CHECK-UNWINDER:   6e7->6b0:6
+; CHECK-UNWINDER:   71b->6c0:7
+; CHECK-UNWINDER:   72b->74c:6
+; CHECK-UNWINDER: [foo:2 @ qux:2 @ callBeforeReturn]
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   6b0-6be:6
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   6be->6ec:7
+; CHECK-UNWINDER: [foo:2 @ qux:4 @ callAfterReturn]
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   6c0-6ce:7
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   6ce->720:7
+; CHECK-UNWINDER: [main]
+; CHECK-UNWINDER:   2
+; CHECK-UNWINDER:   77d-7ab:5
+; CHECK-UNWINDER:   7b0-7bf:5
+; CHECK-UNWINDER:   1
+; CHECK-UNWINDER:   7bf->77d:5
+
+; libcallback.c
+; clang -shared -fPIC -o libcallback.so libcallback.c
+
+int callback(int *cnt, int (*func1)(int), int (*func2)(int), int p) {
+  (*cnt)++;
+  return func1(p) + func2(p);
+}
+
+; test.c
+; clang test.c  -O0 -g -fno-optimize-sibling-calls -fdebug-info-for-profiling -L $PWD -lcallback -fno-inline
+
+#include <stdio.h>
+
+int callbackCnt = 0;
+
+int callback(int *cnt, int (*func1)(int), int (*func2)(int), int p);
+
+int bar(int p) {
+  return p + 1;
+}
+
+int baz(int p) {
+  return p - 1;
+}
+
+int callBeforeReturn(int p) {
+  return p + 10;
+}
+
+int callAfterReturn(int p) {
+  return p - 10;
+}
+
+int qux(int p) {
+  p += 10;
+  int ret = callBeforeReturn(p);
+  ret = callback(&callbackCnt, bar, baz, ret);
+  ret = callAfterReturn(ret);
+  return ret;
+}
+
+int foo (int p) {
+  p -= 10;
+  return qux(p);
+}
+
+int  main(void) {
+  int sum = 0;
+  for (int i = 0; i < 1000 * 1000; i++) {
+    sum += callback(&callbackCnt, foo, bar, i);
+  }
+  printf("callback count=%d, sum=%d\n", callbackCnt, sum);
+}

llvm/test/tools/llvm-profgen/inline-noprobe2.test

 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/artificial-branch.perfscript --binary=%S/Inputs/inline-noprobe2.perfbin --output=%t  --skip-symbolization --use-offset=0
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-ARTIFICIAL-BRANCH
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/inline-noprobe2.perfscript --binary=%S/Inputs/inline-noprobe2.perfbin --output=%t  --skip-symbolization --use-offset=0
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-RAW-PROFILE
 ; RUN: llvm-profgen --format=text --unsymbolized-profile=%t --binary=%S/Inputs/inline-noprobe2.perfbin --output=%t1 --use-offset=0
 ; RUN: FileCheck %s --input-file %t1 --check-prefix=CHECK
 
 ; RUN: llvm-profgen --format=extbinary --perfscript=%S/Inputs/inline-noprobe2.perfscript --binary=%S/Inputs/inline-noprobe2.perfbin --output=%t --populate-profile-symbol-list=1
 ; RUN: llvm-profdata show -show-prof-sym-list -sample %t | FileCheck %s --check-prefix=CHECK-SYM-LIST
 
-; CHECK-ARTIFICIAL-BRANCH: 0
-; CHECK-ARTIFICIAL-BRANCH: 0
+; CHECK-ARTIFICIAL-BRANCH: 2
+; CHECK-ARTIFICIAL-BRANCH: 400870-400870:2
+; CHECK-ARTIFICIAL-BRANCH: 400875-4008bf:1
+; CHECK-ARTIFICIAL-BRANCH: 1
+; CHECK-ARTIFICIAL-BRANCH: 4008bf->400870:2
 
 ; CHECK-SYM-LIST: Dump profile symbol list
 ; CHECK-SYM-LIST: main
 ; CHECK-SYM-LIST: partition_pivot_first
 ; CHECK-SYM-LIST: partition_pivot_last
 ; CHECK-SYM-LIST: quick_sort
 ; CHECK-SYM-LIST: swap
 

llvm/tools/llvm-profgen/PerfReader.h

 // 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>;
 
 using SampleVector = SmallVector<std::tuple<uint64_t, uint64_t, uint64_t>, 16>;
 
-// The special frame addresses.
-enum SpecialFrameAddr {
-  // Dummy root of frame trie.
-  DummyRoot = 0,
-  // Represent all the addresses outside of current binary.
-  ExternalAddr = 1,
-};
-
 // 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;
   // Call stack trie node
   struct ProfiledFrame {
struct UnwindState {
   uint64_t getCurrentLBRSource() const { return LBRStack[LBRIndex].Source; }
   uint64_t getCurrentLBRTarget() const { return LBRStack[LBRIndex].Target; }
   const LBREntry &getCurrentLBR() const { return LBRStack[LBRIndex]; }
   bool IsLastLBR() const { return LBRIndex == 0; }
   bool getLBRStackSize() const { return LBRStack.size(); }
   void advanceLBR() { LBRIndex++; }
   ProfiledFrame *getParentFrame() { return CurrentLeafFrame->Parent; }
 
   void pushFrame(uint64_t Address) {
     CurrentLeafFrame = CurrentLeafFrame->getOrCreateChildFrame(Address);
   }
 
   void switchToFrame(uint64_t Address) {
     if (CurrentLeafFrame->Address == Address)
       return;
     CurrentLeafFrame = CurrentLeafFrame->Parent->getOrCreateChildFrame(Address);
   }
 
   void popFrame() { CurrentLeafFrame = CurrentLeafFrame->Parent; }
 
+  void clearCallStack() { CurrentLeafFrame = &DummyTrieRoot; }
+
   void initFrameTrie(const SmallVectorImpl<uint64_t> &CallStack) {
     ProfiledFrame *Cur = &DummyTrieRoot;
     for (auto Address : reverse(CallStack)) {
       Cur = Cur->getOrCreateChildFrame(Address);
     }
     CurrentLeafFrame = Cur;
   }
 
using ContextSampleCounterMap =
     std::unordered_map<Hashable<ContextKey>, SampleCounter,
                        Hashable<ContextKey>::Hash, Hashable<ContextKey>::Equal>;
 
 struct FrameStack {
   SmallVector<uint64_t, 16> Stack;
   ProfiledBinary *Binary;
   FrameStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
-    // Truncate the context for external frame since this isn't a real call
-    // context the compiler will see
-    if (Cur->isExternalFrame())
-      return false;
+    assert(!Cur->isExternalFrame() &&
+           "External frame's not expected for context stack.");
     Stack.push_back(Cur->Address);
     return true;
   }
 
   void popFrame() {
     if (!Stack.empty())
       Stack.pop_back();
   }
   std::shared_ptr<StringBasedCtxKey> getContextKey();
 };
 
 struct ProbeStack {
   SmallVector<const MCDecodedPseudoProbe *, 16> Stack;
   ProfiledBinary *Binary;
   ProbeStack(ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
-    // Truncate the context for external frame since this isn't a real call
-    // context the compiler will see
-    if (Cur->isExternalFrame())
-      return false;
+    assert(!Cur->isExternalFrame() &&
+           "External frame's not expected for context stack.");
     const MCDecodedPseudoProbe *CallProbe =
         Binary->getCallProbeForAddr(Cur->Address);
     // We may not find a probe for a merged or external callsite.
     // Callsite merging may cause the loss of original probe IDs.
     // Cutting off the context from here since the inliner will
     // not know how to consume a context with unknown callsites.
     if (!CallProbe)
       return false;

 */
 class VirtualUnwinder {
 public:
   VirtualUnwinder(ContextSampleCounterMap *Counter, ProfiledBinary *B)
       : CtxCounterMap(Counter), Binary(B) {}
   bool unwind(const PerfSample *Sample, uint64_t Repeat);
   std::set<uint64_t> &getUntrackedCallsites() { return UntrackedCallsites; }
 
+  uint64_t NumTotalBranches = 0;
+  uint64_t NumExtCallBranch = 0;
+  uint64_t NumMissingExternalFrame = 0;
+  uint64_t NumMismatchedProEpiBranch = 0;
+  uint64_t NumMismatchedExtCallBranch = 0;
+
 private:
   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 Binary->addressIsCall(State.getCurrentLBRSource());
   }
 
   bool isReturnState(UnwindState &State) const {
     // Simply check addressIsReturn, as ret is always reliable, both for
     // regular call and tail call.
-    return Binary->addressIsReturn(State.getCurrentLBRSource());
+    if (!Binary->addressIsReturn(State.getCurrentLBRSource()))
+      return false;
+
+    // In a callback case, a return from internal code, say A, to external
+    // runtime can happen. The external runtime can then call back to
+    // another internal routine, say B. Making an artificial branch that
+    // looks like a return from A to B can confuse the unwinder to treat
+    // the instruction before B as the call instruction. Here we detect this
+    // case if the return target is not the next inst of call inst, then we just
+    // do not treat it as a return.
+    uint64_t CallAddr =
+        Binary->getCallAddrFromFrameAddr(State.getCurrentLBRTarget());
+    return (CallAddr != 0);
   }
 
   void unwindCall(UnwindState &State);
   void unwindLinear(UnwindState &State, uint64_t Repeat);
   void unwindReturn(UnwindState &State);
   void unwindBranch(UnwindState &State);
 
   template <typename T>

llvm/tools/llvm-profgen/PerfReader.cpp

 extern cl::opt<bool> ShowDisassemblyOnly;
 extern cl::opt<bool> ShowSourceLocations;
 extern cl::opt<std::string> OutputFilename;
 
 namespace llvm {
 namespace sampleprof {
 
 void VirtualUnwinder::unwindCall(UnwindState &State) {
+  uint64_t Source = State.getCurrentLBRSource();
+  // An artificial return should push an external frame and an artificial call
+  // will match it and pop the external frame so that the context before and
+  // after the external call will be the same.
+  if (State.getCurrentLBR().IsArtificial) {
+    NumExtCallBranch++;
+    // A return is matched and pop the external frame.
+    if (State.getParentFrame()->isExternalFrame()) {
+      State.popFrame();
+    } else {
+      // An artificial return is missing, it happens that the sample is just hit
+      // in the middle of the external code. In this case, the leading branch is
+      // a call to external, we just keep unwinding use a context-less stack.
+      if (State.getParentFrame() != State.getDummyRootPtr())
+        NumMissingExternalFrame++;
+      State.clearCallStack();
+      State.pushFrame(Source);
+      State.InstPtr.update(Source);
+      return;
+    }
+  }
+
+  auto *ParentFrame = State.getParentFrame();
   // The 2nd frame after leaf could be missing if stack sample is
   // taken when IP is within prolog/epilog, as frame chain isn't
   // setup yet. Fill in the missing frame in that case.
   // TODO: Currently we just assume all the addr that can't match the
   // 2nd frame is in prolog/epilog. In the future, we will switch to
   // pro/epi tracker(Dwarf CFI) for the precise check.
-  uint64_t Source = State.getCurrentLBRSource();
-  auto *ParentFrame = State.getParentFrame();
-
   if (ParentFrame == State.getDummyRootPtr() ||
       ParentFrame->Address != Source) {
     State.switchToFrame(Source);
+    if (ParentFrame != State.getDummyRootPtr()) {
+      if (State.getCurrentLBR().IsArtificial)
+        NumMismatchedExtCallBranch++;
+      else
+        NumMismatchedProEpiBranch++;
+    }
   } else {
     State.popFrame();
   }
   State.InstPtr.update(Source);
 }
 
 void VirtualUnwinder::unwindLinear(UnwindState &State, uint64_t Repeat) {
   InstructionPointer &IP = State.InstPtr;
void VirtualUnwinder::unwindLinear(UnwindState &State, uint64_t Repeat) {
   }
 }
 
 void VirtualUnwinder::unwindReturn(UnwindState &State) {
   // Add extra frame as we unwind through the return
   const LBREntry &LBR = State.getCurrentLBR();
   uint64_t CallAddr = Binary->getCallAddrFromFrameAddr(LBR.Target);
   State.switchToFrame(CallAddr);
+  // Push an external frame for the case of returning to external
+  // address(callback), later if an aitificial call is matched and it will be
+  // popped up. This is to 1)avoid context being interrupted by callback,
+  // context before or after the callback should be the same. 2) the call stack
+  // of function called by callback should be truncated which is done during
+  // recording the context on trie. For example:
+  //  main (call)--> foo (call)--> callback (call)--> bar (return)--> callback
+  //  (return)--> foo (return)--> main
+  // Context for bar should not include main and foo.
+  // For the code of foo, the context of before and after callback should both
+  // be [foo, main].
+  if (LBR.IsArtificial)
+    State.pushFrame(ExternalAddr);
   State.pushFrame(LBR.Source);
   State.InstPtr.update(LBR.Source);
 }
 
 void VirtualUnwinder::unwindBranch(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();
for (auto &Item : Cur->BranchSamples) {
     SCounter.recordBranchCount(SourceOffset, TargetOffset, std::get<2>(Item));
   }
 }
 
 template <typename T>
 void VirtualUnwinder::collectSamplesFromFrameTrie(
     UnwindState::ProfiledFrame *Cur, T &Stack) {
   if (!Cur->isDummyRoot()) {
-    if (!Stack.pushFrame(Cur)) {
+    // Truncate the context for external frame since this isn't a real call
+    // context the compiler will see.
+    if (Cur->isExternalFrame() || !Stack.pushFrame(Cur)) {
       // Process truncated context
       // Start a new traversal ignoring its bottom context
       T EmptyStack(Binary);
       collectSamplesFromFrame(Cur, EmptyStack);
       for (const auto &Item : Cur->Children) {
         collectSamplesFromFrameTrie(Item.second.get(), EmptyStack);
       }
 
for (auto Address : Unwinder.getUntrackedCallsites())
                            << "for call instruction at "
                            << format("0x%" PRIx64, Address) << "\n";
   }
 
   emitWarningSummary(Unwinder.getUntrackedCallsites().size(),
                      SampleCounters.size(),
                      "of profiled contexts are truncated due to missing probe "
                      "for call instruction.");
+
+  emitWarningSummary(
+      Unwinder.NumMismatchedExtCallBranch, Unwinder.NumTotalBranches,
+      "of branches'source is a call instruction but doesn't match call frame "
+      "stack, likely due to unwinding error of external frame.");
+
+  emitWarningSummary(
+      Unwinder.NumMismatchedProEpiBranch, Unwinder.NumTotalBranches,
+      "of branches'source is a call instruction but doesn't match call frame "
+      "stack, likely due to frame in prolog/epilog.");
+
+  emitWarningSummary(Unwinder.NumMissingExternalFrame,
+                     Unwinder.NumExtCallBranch,
+                     "of artificial call branches but doesn't have an external "
+                     "frame to match.");
 }
 
 bool PerfScriptReader::extractLBRStack(TraceStream &TraceIt,
                                        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.
if (IsOutgoing) {
           continue;
         }
         // This is middle unpaired outgoing jump which is likely due to
         // interrupt or incomplete LBR trace. Ignore current and subsequent
         // entries since they are likely in different contexts.
         break;
       }
 
-      if (Binary->addressIsReturn(Src)) {
-        // In a callback case, a return from internal code, say A, to external
-        // runtime can happen. The external runtime can then call back to
-        // another internal routine, say B. Making an artificial branch that
-        // looks like a return from A to B can confuse the unwinder to treat
-        // the instruction before B as the call instruction.
-        break;
-      }
-
       // For transition to external code, group the Source with the next
       // availabe transition target.
       Dst = PrevTrDst;
       PrevTrDst = 0;
       IsArtificial = true;
     } else {
       if (PrevTrDst) {
         // If we have seen an incoming transition from external code to internal
void UnsymbolizedProfileReader::parsePerfTraces() {
   readUnsymbolizedProfile(PerfTraceFile);
 }
 
 void PerfScriptReader::computeCounterFromLBR(const PerfSample *Sample,
                                              uint64_t Repeat) {
   SampleCounter &Counter = SampleCounters.begin()->second;
   uint64_t EndOffeset = 0;
   for (const LBREntry &LBR : Sample->LBRStack) {
+    assert(LBR.Source != ExternalAddr &&
+           "Branch' source should not be an external address, it should be "
+           "converted to aritificial branch.");
     uint64_t SourceOffset = Binary->virtualAddrToOffset(LBR.Source);
-    uint64_t TargetOffset = Binary->virtualAddrToOffset(LBR.Target);
+    uint64_t TargetOffset = LBR.Target == ExternalAddr
+                                ? ExternalAddr
+                                : Binary->virtualAddrToOffset(LBR.Target);
 
-    if (!LBR.IsArtificial) {
+    if (!LBR.IsArtificial && TargetOffset != ExternalAddr) {
       Counter.recordBranchCount(SourceOffset, TargetOffset, Repeat);
     }
 
     // If this not the first LBR, update the range count between TO of current
     // LBR and FROM of next LBR.
     uint64_t StartOffset = TargetOffset;
     if (EndOffeset != 0)
       Counter.recordRangeCount(StartOffset, EndOffeset, Repeat);

llvm/tools/llvm-profgen/ProfiledBinary.h

struct InstructionPointer {
   uint64_t Index = 0;
   InstructionPointer(const ProfiledBinary *Binary, uint64_t Address,
                      bool RoundToNext = false);
   bool advance();
   bool backward();
   void update(uint64_t Addr);
 };
 
+// The special frame addresses.
+enum SpecialFrameAddr {
+  // Dummy root of frame trie.
+  DummyRoot = 0,
+  // Represent all the addresses outside of current binary.
+  // This's also used to indicate the call stack should be truncated since this
+  // isn't a real call context the compiler will see.
+  ExternalAddr = 1,
+};
+
 using RangesTy = std::vector<std::pair<uint64_t, uint64_t>>;
 
 struct BinaryFunction {
   StringRef FuncName;
   // End of range is an exclusive bound.
   RangesTy Ranges;
 
   uint64_t getFuncSize() {
uint32_t getIndexForOffset(uint64_t Offset) const {
     return Low - CodeAddrOffsets.begin();
   }
   uint32_t getIndexForAddr(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     return getIndexForOffset(Offset);
   }
 
   uint64_t getCallAddrFromFrameAddr(uint64_t FrameAddr) const {
+    if (FrameAddr == ExternalAddr)
+      return ExternalAddr;
     auto I = getIndexForAddr(FrameAddr);
     FrameAddr = I ? getAddressforIndex(I - 1) : 0;
     if (FrameAddr && addressIsCall(FrameAddr))
       return FrameAddr;
     return 0;
   }
 
   FuncRange *findFuncRangeForStartOffset(uint64_t Offset) {