[BOLT] Set EntryDiscriminator in YAML profile for indirect calls

Indirect call handling missed setting an `EntryDiscriminator` while it's
set for direct calls and tail calls.

Improve YAML profile accuracy by unifying the destination setting
between direct and indirect calls into `setCSIDestination` method.

Depends on: https://github.com/llvm/llvm-project/pull/86848

Test Plan: Updated bolt/test/X86/yaml-secondary-entry-discriminator.s

Reviewers: ayermolo, maksfb, rafaelauler

Reviewed By: maksfb

Pull Request: https://github.com/llvm/llvm-project/pull/82128

2 files changed, 30 insertions, 19 deletions

diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index 6fcc4a9..0f08208 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -25,6 +25,25 @@ extern llvm::cl::opt<bool> ProfileUseDFS;
 namespace llvm {
 namespace bolt {
 
+/// Set CallSiteInfo destination fields from \p Symbol and return a target
+/// BinaryFunction for that symbol.
+static const BinaryFunction *setCSIDestination(const BinaryContext &BC,
+                                               yaml::bolt::CallSiteInfo &CSI,
+                                               const MCSymbol *Symbol) {
+  CSI.DestId = 0; // designated for unknown functions
+  CSI.EntryDiscriminator = 0;
+  if (Symbol) {
+    uint64_t EntryID = 0;
+    if (const BinaryFunction *const Callee =
+            BC.getFunctionForSymbol(Symbol, &EntryID)) {
+      CSI.DestId = Callee->getFunctionNumber();
+      CSI.EntryDiscriminator = EntryID;
+      return Callee;
+    }
+  }
+  return nullptr;
+}
+
 yaml::bolt::BinaryFunctionProfile
 YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS) {
   yaml::bolt::BinaryFunctionProfile YamlBF;
@@ -79,31 +98,20 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS) {
           continue;
         for (const IndirectCallProfile &CSP : ICSP.get()) {
           StringRef TargetName = "";
-          CSI.DestId = 0; // designated for unknown functions
-          CSI.EntryDiscriminator = 0;
-          if (CSP.Symbol) {
-            const BinaryFunction *Callee = BC.getFunctionForSymbol(CSP.Symbol);
-            if (Callee) {
-              CSI.DestId = Callee->getFunctionNumber();
-              TargetName = Callee->getOneName();
-            }
-          }
+          const BinaryFunction *Callee = setCSIDestination(BC, CSI, CSP.Symbol);
+          if (Callee)
+            TargetName = Callee->getOneName();
           CSI.Count = CSP.Count;
           CSI.Mispreds = CSP.Mispreds;
           CSTargets.emplace_back(TargetName, CSI);
         }
       } else { // direct call or a tail call
-        uint64_t EntryID = 0;
-        CSI.DestId = 0;
         StringRef TargetName = "";
         const MCSymbol *CalleeSymbol = BC.MIB->getTargetSymbol(Instr);
         const BinaryFunction *const Callee =
-            BC.getFunctionForSymbol(CalleeSymbol, &EntryID);
-        if (Callee) {
-          CSI.DestId = Callee->getFunctionNumber();
-          CSI.EntryDiscriminator = EntryID;
+            setCSIDestination(BC, CSI, CalleeSymbol);
+        if (Callee)
           TargetName = Callee->getOneName();
-        }
 
         auto getAnnotationWithDefault = [&](const MCInst &Inst, StringRef Ann) {
           return BC.MIB->getAnnotationWithDefault(Instr, Ann, 0ull);
diff --git a/bolt/test/X86/yaml-secondary-entry-discriminator.s b/bolt/test/X86/yaml-secondary-entry-discriminator.s
index 55dc024..43c2e2a 100644
--- a/bolt/test/X86/yaml-secondary-entry-discriminator.s
+++ b/bolt/test/X86/yaml-secondary-entry-discriminator.s
@@ -19,6 +19,10 @@
 # CHECK-NEXT:       insns: 1
 # CHECK-NEXT:       hash:  0x36A303CBA4360014
 # CHECK-NEXT:       calls: [ { off: 0x0, fid: 1, disc: 1, cnt: 1 } ]
+# CHECK:          - bid:   2
+# CHECK-NEXT:       insns: 5
+# CHECK-NEXT:       hash:  0x8B2F5747CD0019
+# CHECK-NEXT:       calls: [ { off: 0x0, fid: 1, disc: 1, cnt: 1, mis: 1 } ]
 
 # Make sure that the profile is attached correctly
 # RUN: llvm-bolt %t.exe -o %t.out --data %t.yaml --print-profile \
@@ -27,13 +31,12 @@
 # CHECK-CFG: Binary Function "main" after attaching profile {
 # CHECK-CFG:      callq secondary_entry # Offset: [[#]] # Count: 1
 # CHECK-CFG:      callq *%rax # Offset: [[#]] # CallProfile: 1 (1 misses) :
-# XXX: uncomment after discriminator is set correctly for indirect calls
-# COM: CHECK-CFG-NEXT:     { secondary_entry: 1 (1 misses) }
-# CHECK-CFG: }
+# CHECK-CFG-NEXT:     { secondary_entry: 1 (1 misses) }
 
 .globl func
 .type	func, @function
 func:
+# FDATA: 0 [unknown] 0 1 func 0 1 0
   .cfi_startproc
   pushq   %rbp
   movq    %rsp, %rbp