[llvm] [BOLT] Initialize auxiliary variables for CDSplit (PR #73895)

[via llvm-commits]

https://github.com/ShatianWang updated https://github.com/llvm/llvm-project/pull/73895

>From d1bc831fbf741c5d7fc8d8371a4d912171d508c6 Mon Sep 17 00:00:00 2001
From: Shatian Wang <shatian at meta.com>
Date: Sun, 26 Nov 2023 13:03:53 -0800
Subject: [PATCH] [BOLT] Initialize auxiliary variables for CDSplit

This diff defines and initializes auxiliary variables used by
CDSplit. NFC.
---
 bolt/lib/Passes/SplitFunctions.cpp | 123 ++++++++++++++++++++++++++++-
 1 file changed, 122 insertions(+), 1 deletion(-)

diff --git a/bolt/lib/Passes/SplitFunctions.cpp b/bolt/lib/Passes/SplitFunctions.cpp
index 79da4cc3b04ba95..9709d5ea15f9777 100644
--- a/bolt/lib/Passes/SplitFunctions.cpp
+++ b/bolt/lib/Passes/SplitFunctions.cpp
@@ -146,6 +146,11 @@ struct SplitCacheDirected final : public SplitStrategy {
     return BF.hasValidProfile() && hasFullProfile(BF) && !allBlocksCold(BF);
   }
 
+  explicit SplitCacheDirected(BinaryContext &BC) {
+    initializeAuxiliaryVariables(BC);
+    buildCallGraph(BC);
+  }
+
   // When some functions are hot-warm split and others are hot-warm-cold split,
   // we do not want to change the fragment numbers of the blocks in the hot-warm
   // split functions.
@@ -173,6 +178,122 @@ struct SplitCacheDirected final : public SplitStrategy {
   }
 
 private:
+  struct JumpInfo {
+    bool HasUncondBranch = false;
+    BinaryBasicBlock *CondSuccessor = nullptr;
+    BinaryBasicBlock *UncondSuccessor = nullptr;
+  };
+
+  // Auxiliary variables used by the algorithm.
+  size_t TotalNumBlocks{0};
+  size_t OrigHotSectionSize{0};
+  DenseMap<const BinaryBasicBlock *, size_t> GlobalIndices;
+  DenseMap<const BinaryBasicBlock *, size_t> BBSizes;
+  DenseMap<const BinaryBasicBlock *, size_t> BBOffsets;
+  DenseMap<const BinaryBasicBlock *, JumpInfo> JumpInfos;
+
+  // Call graph.
+  std::vector<SmallVector<const BinaryBasicBlock *, 0>> Callers;
+  std::vector<SmallVector<const BinaryBasicBlock *, 0>> Callees;
+
+  bool shouldConsiderForCallGraph(const BinaryFunction *BF) {
+    // Only a subset of the functions in the binary will be considered
+    // for initializing auxiliary variables and building call graph.
+    return BF->hasValidIndex() && BF->hasValidProfile() && !BF->empty();
+  }
+
+  void initializeAuxiliaryVariables(BinaryContext &BC) {
+    // Gather information about conditional and unconditional successors of
+    // each basic block; this information will be used to estimate block size
+    // increase due to hot-warm splitting.
+    auto analyzeBranches = [&](BinaryBasicBlock *BB) {
+      JumpInfo BBJumpInfo;
+      const MCSymbol *TBB = nullptr;
+      const MCSymbol *FBB = nullptr;
+      MCInst *CondBranch = nullptr;
+      MCInst *UncondBranch = nullptr;
+      if (BB->analyzeBranch(TBB, FBB, CondBranch, UncondBranch)) {
+        BBJumpInfo.HasUncondBranch = UncondBranch != nullptr;
+        if (BB->succ_size() == 1) {
+          BBJumpInfo.UncondSuccessor = BB->getSuccessor();
+        } else if (BB->succ_size() == 2) {
+          BBJumpInfo.CondSuccessor = BB->getConditionalSuccessor(true);
+          BBJumpInfo.UncondSuccessor = BB->getConditionalSuccessor(false);
+        }
+      }
+      return BBJumpInfo;
+    };
+
+    for (BinaryFunction *BF : BC.getSortedFunctions()) {
+      if (!shouldConsiderForCallGraph(BF))
+        continue;
+
+      // Calculate the size of each BB after hot-cold splitting.
+      // This populates BinaryBasicBlock::OutputAddressRange which
+      // can be used to compute the size of each BB.
+      BC.calculateEmittedSize(*BF, /*FixBranches=*/true);
+
+      for (BinaryBasicBlock *BB : BF->getLayout().blocks()) {
+        // Unique global index.
+        GlobalIndices[BB] = TotalNumBlocks;
+        TotalNumBlocks++;
+
+        // Block size after hot-cold splitting.
+        BBSizes[BB] = BB->getOutputSize();
+
+        // Hot block offset after hot-cold splitting.
+        BBOffsets[BB] = OrigHotSectionSize;
+        if (!BB->isSplit())
+          OrigHotSectionSize += BBSizes[BB];
+
+        // (Un)Conditional branch instruction information.
+        JumpInfos[BB] = analyzeBranches(BB);
+      }
+    }
+  }
+
+  void buildCallGraph(BinaryContext &BC) {
+    Callers.resize(TotalNumBlocks);
+    Callees.resize(TotalNumBlocks);
+    for (BinaryFunction *SrcFunction : BC.getSortedFunctions()) {
+      if (!shouldConsiderForCallGraph(SrcFunction))
+        continue;
+
+      for (BinaryBasicBlock &SrcBB : SrcFunction->blocks()) {
+        // Skip blocks that are not executed
+        if (SrcBB.getKnownExecutionCount() == 0)
+          continue;
+
+        // Find call instructions and extract target symbols from each one
+        for (const MCInst &Inst : SrcBB) {
+          if (!BC.MIB->isCall(Inst))
+            continue;
+
+          // Call info
+          const MCSymbol *DstSym = BC.MIB->getTargetSymbol(Inst);
+          // Ignore calls w/o information
+          if (!DstSym)
+            continue;
+
+          const BinaryFunction *DstFunction = BC.getFunctionForSymbol(DstSym);
+          // Ignore calls that do not have a valid target, but do not ignore
+          // recursive calls, because caller block could be moved to warm.
+          if (!DstFunction || DstFunction->getLayout().block_empty())
+            continue;
+
+          const BinaryBasicBlock *DstBB = &(DstFunction->front());
+
+          // Record the call only if DstBB is also in functions to consider for
+          // call graph.
+          if (GlobalIndices.contains(DstBB)) {
+            Callers[GlobalIndices[DstBB]].push_back(&SrcBB);
+            Callees[GlobalIndices[&SrcBB]].push_back(DstBB);
+          }
+        }
+      }
+    }
+  }
+
   /// Find the best index for splitting. The returned value is the index of the
   /// last hot basic block. Hence, "no splitting" is equivalent to returning the
   /// value which is one less than the size of the function.
@@ -308,7 +429,7 @@ void SplitFunctions::runOnFunctions(BinaryContext &BC) {
     // before function reordering and hot-warm-cold splitting
     // (SplitCacheDirected) after function reordering.
     if (BC.HasFinalizedFunctionOrder)
-      Strategy = std::make_unique<SplitCacheDirected>();
+      Strategy = std::make_unique<SplitCacheDirected>(BC);
     else
       Strategy = std::make_unique<SplitProfile2>();
     opts::AggressiveSplitting = true;