[BOLT] Use aggregated FuncBranchData in writeBATYAML

Switch from FuncBranchData intermediate maps (Intra/InterIndex)
to aggregated Data, same as one used by DataReader:
https://github.com/llvm/llvm-project/blob/e62ce1f8842cca36eb14126d79dcca0a85bf6d36/bolt/lib/Profile/DataReader.cpp#L385-L389
This aligns the order of the output between YAMLProfileWriter and
writeBATYAML.

Test Plan: updated bolt-address-translation-yaml.test

Reviewers: rafaelauler, dcci, ayermolo, maksfb

Reviewed By: ayermolo, maksfb

Pull Request: #91289

Author: aaupov

bolt/lib/Profile/DataAggregator.cpp

@@ -2355,30 +2355,6 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       for (auto BI = BlockMap.begin(), BE = BlockMap.end(); BI != BE; ++BI)
         YamlBF.Blocks[BI->second.getBBIndex()].Hash = BI->second.getBBHash();
 
-      auto getSuccessorInfo = [&](uint32_t SuccOffset, unsigned SuccDataIdx) {
-        const llvm::bolt::BranchInfo &BI = Branches.Data.at(SuccDataIdx);
-        yaml::bolt::SuccessorInfo SI;
-        SI.Index = BlockMap.getBBIndex(SuccOffset);
-        SI.Count = BI.Branches;
-        SI.Mispreds = BI.Mispreds;
-        return SI;
-      };
-
-      auto getCallSiteInfo = [&](Location CallToLoc, unsigned CallToIdx,
-                                 uint32_t Offset) {
-        const llvm::bolt::BranchInfo &BI = Branches.Data.at(CallToIdx);
-        yaml::bolt::CallSiteInfo CSI;
-        CSI.DestId = 0; // designated for unknown functions
-        CSI.EntryDiscriminator = 0;
-        CSI.Count = BI.Branches;
-        CSI.Mispreds = BI.Mispreds;
-        CSI.Offset = Offset;
-        if (BinaryData *BD = BC.getBinaryDataByName(CallToLoc.Name))
-          YAMLProfileWriter::setCSIDestination(BC, CSI, BD->getSymbol(), BAT,
-                                               CallToLoc.Offset);
-        return CSI;
-      };
-
       // Lookup containing basic block offset and index
       auto getBlock = [&BlockMap](uint32_t Offset) {
         auto BlockIt = BlockMap.upper_bound(Offset);
@@ -2390,25 +2366,26 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         return std::pair(BlockIt->first, BlockIt->second.getBBIndex());
       };
 
-      for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
-        const auto &[_, Index] = getBlock(FromOffset);
-        yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[Index];
-        for (const auto &[SuccOffset, SuccDataIdx] : SuccKV)
-          if (BlockMap.isInputBlock(SuccOffset))
-            YamlBB.Successors.emplace_back(
-                getSuccessorInfo(SuccOffset, SuccDataIdx));
-      }
-      for (const auto &[FromOffset, CallTo] : Branches.InterIndex) {
-        const auto &[BlockOffset, BlockIndex] = getBlock(FromOffset);
-        yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[BlockIndex];
-        const uint32_t Offset = FromOffset - BlockOffset;
-        for (const auto &[CallToLoc, CallToIdx] : CallTo)
-          YamlBB.CallSites.emplace_back(
-              getCallSiteInfo(CallToLoc, CallToIdx, Offset));
-        llvm::sort(YamlBB.CallSites, [](yaml::bolt::CallSiteInfo &A,
-                                        yaml::bolt::CallSiteInfo &B) {
-          return A.Offset < B.Offset;
-        });
+      for (const llvm::bolt::BranchInfo &BI : Branches.Data) {
+        using namespace yaml::bolt;
+        const auto &[BlockOffset, BlockIndex] = getBlock(BI.From.Offset);
+        BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[BlockIndex];
+        if (BI.To.IsSymbol && BI.To.Name == BI.From.Name && BI.To.Offset != 0) {
+          // Internal branch
+          const unsigned SuccIndex = getBlock(BI.To.Offset).second;
+          auto &SI = YamlBB.Successors.emplace_back(SuccessorInfo{SuccIndex});
+          SI.Count = BI.Branches;
+          SI.Mispreds = BI.Mispreds;
+        } else {
+          // Call
+          const uint32_t Offset = BI.From.Offset - BlockOffset;
+          auto &CSI = YamlBB.CallSites.emplace_back(CallSiteInfo{Offset});
+          CSI.Count = BI.Branches;
+          CSI.Mispreds = BI.Mispreds;
+          if (const BinaryData *BD = BC.getBinaryDataByName(BI.To.Name))
+            YAMLProfileWriter::setCSIDestination(BC, CSI, BD->getSymbol(), BAT,
+                                                 BI.To.Offset);
+        }
       }
       // Set entry counts, similar to DataReader::readProfile.
       for (const llvm::bolt::BranchInfo &BI : Branches.EntryData) {

bolt/test/X86/Inputs/blarge_new_bat_order.preagg.txt

@@ -0,0 +1,2 @@
+B 800154 401050 20 0
+F 800159 800193 7

bolt/test/X86/bolt-address-translation-yaml.test

@@ -15,6 +15,17 @@ BRANCHENTRY-YAML-CHECK:    - name: SolveCubic
 BRANCHENTRY-YAML-CHECK:      bid: 0
 BRANCHENTRY-YAML-CHECK:      hash: 0x700F19D24600000
 BRANCHENTRY-YAML-CHECK-NEXT: succ: [ { bid: 7, cnt: 1 }
+# Check that the order is correct between BAT YAML and FDATA->YAML.
+RUN: perf2bolt %t.out --pa -p %p/Inputs/blarge_new_bat_order.preagg.txt \
+RUN:   -w %t.yaml -o %t.fdata
+RUN: llvm-bolt %t.exe -data %t.fdata -w %t.yaml-fdata -o %t.null
+RUN: FileCheck --input-file %t.yaml --check-prefix ORDER-YAML-CHECK %s
+RUN: FileCheck --input-file %t.yaml-fdata --check-prefix ORDER-YAML-CHECK %s
+ORDER-YAML-CHECK:    - name: SolveCubic
+ORDER-YAML-CHECK:      bid: 3
+ORDER-YAML-CHECK:      hash: 0xDDA1DC5F69F900AC
+ORDER-YAML-CHECK-NEXT: calls: [ { off: 0x26, fid: [[#]], cnt: 20 } ]
+ORDER-YAML-CHECK-NEXT: succ: [ { bid: 5, cnt: 7 }
 # Large profile test
 RUN: perf2bolt %t.out --pa -p %p/Inputs/blarge_new_bat.preagg.txt -w %t.yaml -o %t.fdata \
 RUN:   2>&1 | FileCheck --check-prefix READ-BAT-CHECK %s