[llvm] [BOLT] Extend calculateEmittedSize for Block Size Calculation (PR #72218)

[via llvm-commits]

https://github.com/ShatianWang created https://github.com/llvm/llvm-project/pull/72218

This commit modifies BinaryContext::calculateEmittedSize to update the
BinaryBasicBlock::OutputAddressRange for each basic block in the input BF.
The modification is done in place, where BB.OutputAddressRange.second less
BB.OutputAddressRange.first now gives the emitted size of the basic block.

>From 1dbdd1c998493989eaa1dc32f298a014523f6497 Mon Sep 17 00:00:00 2001
From: Shatian Wang <shatian at meta.com>
Date: Thu, 2 Nov 2023 12:26:49 -0700
Subject: [PATCH] [BOLT] Extend calculateEmittedSize for Block Size Calculation

---
 bolt/lib/Core/BinaryContext.cpp | 45 ++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 6 deletions(-)

diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index fd4d09964b725e1..a68168a2303159c 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -2305,14 +2305,47 @@ BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
   MCAsmLayout Layout(Assembler);
   Assembler.layout(Layout);
 
+  // Obtain main fragment size.
   const uint64_t HotSize =
       Layout.getSymbolOffset(*EndLabel) - Layout.getSymbolOffset(*StartLabel);
-  const uint64_t ColdSize =
-      std::accumulate(SplitLabels.begin(), SplitLabels.end(), 0ULL,
-                      [&](const uint64_t Accu, const LabelRange &Labels) {
-                        return Accu + Layout.getSymbolOffset(*Labels.second) -
-                               Layout.getSymbolOffset(*Labels.first);
-                      });
+  // Populate new start and end offsets of each basic block in main.
+  BinaryBasicBlock *PrevBB = nullptr;
+  for (BinaryBasicBlock *BB : BF.getLayout().getMainFragment()) {
+    uint64_t BBStartOffset = Layout.getSymbolOffset(*(BB->getLabel()));
+    BB->setOutputStartAddress(BBStartOffset);
+    if (PrevBB)
+      PrevBB->setOutputEndAddress(BBStartOffset);
+    PrevBB = BB;
+  }
+  if (PrevBB)
+    PrevBB->setOutputEndAddress(HotSize);
+
+  // Obtain split fragment sizes.
+  std::vector<uint64_t> SplitFragmentSizes;
+  uint64_t ColdSize = 0;
+  for (const auto &Labels : SplitLabels) {
+    uint64_t Size = Layout.getSymbolOffset(*Labels.second) -
+                    Layout.getSymbolOffset(*Labels.first);
+    SplitFragmentSizes.push_back(Size);
+    ColdSize += Size;
+  }
+
+  // Populate new start and end offsets of each basic block in split fragments.
+  PrevBB = nullptr;
+  uint64_t FragmentInd = 0;
+  for (FunctionFragment &FF : BF.getLayout().getSplitFragments()) {
+    for (BinaryBasicBlock *BB : FF) {
+      uint64_t BBStartOffset = Layout.getSymbolOffset(*(BB->getLabel()));
+      BB->setOutputStartAddress(BBStartOffset);
+      if (PrevBB)
+        PrevBB->setOutputEndAddress(BBStartOffset);
+      PrevBB = BB;
+    }
+    if (PrevBB)
+      PrevBB->setOutputEndAddress(SplitFragmentSizes[FragmentInd]);
+    FragmentInd++;
+    PrevBB = nullptr;
+  }
 
   // Clean-up the effect of the code emission.
   for (const MCSymbol &Symbol : Assembler.symbols()) {