[llvm-branch-commits] [llvm] [BOLT] Use BAT interfaces in YAMLProfileWriter::convert (PR #86219)

[Amir Ayupov via llvm-branch-commits]

================
@@ -27,25 +28,55 @@ 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) {
+static const BinaryFunction *
+setCSIDestination(const BinaryContext &BC, yaml::bolt::CallSiteInfo &CSI,
+                  const MCSymbol *Symbol, const BoltAddressTranslation *BAT) {
   CSI.DestId = 0; // designated for unknown functions
   CSI.EntryDiscriminator = 0;
+  auto setBATSecondaryEntry = [&](const BinaryFunction *const Callee) {
+    // The symbol could be a secondary entry in a cold fragment.
+    ErrorOr<uint64_t> SymbolValue = BC.getSymbolValue(*Symbol);
+    if (SymbolValue.getError())
+      return;
+
+    // Containing function, not necessarily the same as symbol value.
+    const uint64_t CalleeAddress = Callee->getAddress();
+    const uint32_t OutputOffset = SymbolValue.get() - CalleeAddress;
+
+    const uint64_t ParentAddress = BAT->fetchParentAddress(CalleeAddress);
+    const uint64_t HotAddress = ParentAddress ? ParentAddress : CalleeAddress;
+
+    if (const BinaryFunction *ParentBF =
+            BC.getBinaryFunctionAtAddress(HotAddress))
+      CSI.DestId = ParentBF->getFunctionNumber();
+
+    const uint32_t InputOffset =
+        BAT->translate(CalleeAddress, OutputOffset, /*IsBranchSrc*/ false);
+
+    if (!InputOffset)
+      return;
+
+    CSI.EntryDiscriminator =
+        BAT->getSecondaryEntryPointId(HotAddress, InputOffset) + 1;
+  };
+
   if (Symbol) {
     uint64_t EntryID = 0;
     if (const BinaryFunction *const Callee =
             BC.getFunctionForSymbol(Symbol, &EntryID)) {
       CSI.DestId = Callee->getFunctionNumber();
       CSI.EntryDiscriminator = EntryID;
+      if (BAT && BAT->isBATFunction(Callee->getAddress()))
+        setBATSecondaryEntry(Callee);
----------------
aaupov wrote:

I'd rather have `getFunctionForSymbol` stay the same, as we deal with BAT exclusively in profile-related code. 

Moving BAT to BC from RI would be an ergonomic change regardless.

https://github.com/llvm/llvm-project/pull/86219