[llvm] [BOLT] Refactor SplitFunctions for Function Reuse (PR #73078)

[via llvm-commits]

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

>From 4d47f054990c18738242786e38ab1d2e5c9e7c89 Mon Sep 17 00:00:00 2001
From: Shatian Wang <shatian at meta.com>
Date: Wed, 8 Nov 2023 08:20:11 -0800
Subject: [PATCH] [BOLT] Refactor SplitFunctions for function reuse. NFC.

This commit updates SplitFunctions.h and SplitFunctions.cpp to enable
the reuse of createEHTrampolines, mergeEHTrampolines, hasFullProfile,
and allBlocksCold by a distinct function splitting pass (CDSplit).
---
 bolt/include/bolt/Core/BinaryFunction.h   | 14 ++++++++++
 bolt/include/bolt/Passes/SplitFunctions.h | 32 +++++++++++------------
 bolt/lib/Passes/SplitFunctions.cpp        | 18 +++----------
 3 files changed, 33 insertions(+), 31 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 72c360ca0c2db66..3723cccc50f040c 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -1272,6 +1272,20 @@ class BinaryFunction {
   /// otherwise processed.
   bool isPseudo() const { return IsPseudo; }
 
+  /// Return true if every block in the function has a valid execution count.
+  bool hasFullProfile() const {
+    return llvm::all_of(blocks(), [](const BinaryBasicBlock &BB) {
+      return BB.getExecutionCount() != BinaryBasicBlock::COUNT_NO_PROFILE;
+    });
+  }
+
+  /// Return true if every block in the function has a zero execution count.
+  bool allBlocksCold() const {
+    return llvm::all_of(blocks(), [](const BinaryBasicBlock &BB) {
+      return BB.getExecutionCount() == 0;
+    });
+  }
+
   /// Return true if the function contains explicit or implicit indirect branch
   /// to its split fragments, e.g., split jump table, landing pad in split
   /// fragment.
diff --git a/bolt/include/bolt/Passes/SplitFunctions.h b/bolt/include/bolt/Passes/SplitFunctions.h
index 4058f3317dfbdbb..91b6d5518eaab26 100644
--- a/bolt/include/bolt/Passes/SplitFunctions.h
+++ b/bolt/include/bolt/Passes/SplitFunctions.h
@@ -50,6 +50,19 @@ class SplitFunctions : public BinaryFunctionPass {
   /// Split function body into fragments.
   void splitFunction(BinaryFunction &Function, SplitStrategy &Strategy);
 
+  std::atomic<uint64_t> SplitBytesHot{0ull};
+  std::atomic<uint64_t> SplitBytesCold{0ull};
+
+public:
+  explicit SplitFunctions(const cl::opt<bool> &PrintPass)
+      : BinaryFunctionPass(PrintPass) {}
+
+  bool shouldOptimize(const BinaryFunction &BF) const override;
+
+  const char *getName() const override { return "split-functions"; }
+
+  void runOnFunctions(BinaryContext &BC) override;
+
   struct TrampolineKey {
     FragmentNum SourceFN = FragmentNum::main();
     const MCSymbol *Target = nullptr;
@@ -81,27 +94,14 @@ class SplitFunctions : public BinaryFunctionPass {
   /// corresponding thrower block. The trampoline landing pad, when created,
   /// will redirect the execution to the real landing pad in a different
   /// fragment.
-  TrampolineSetType createEHTrampolines(BinaryFunction &Function) const;
+  static TrampolineSetType createEHTrampolines(BinaryFunction &Function);
 
   /// Merge trampolines into \p Layout without trampolines. The merge will place
   /// a trampoline immediately before its destination. Used to revert the effect
   /// of trampolines after createEHTrampolines().
-  BasicBlockOrderType
+  static BasicBlockOrderType
   mergeEHTrampolines(BinaryFunction &BF, BasicBlockOrderType &Layout,
-                     const TrampolineSetType &Trampolines) const;
-
-  std::atomic<uint64_t> SplitBytesHot{0ull};
-  std::atomic<uint64_t> SplitBytesCold{0ull};
-
-public:
-  explicit SplitFunctions(const cl::opt<bool> &PrintPass)
-      : BinaryFunctionPass(PrintPass) {}
-
-  bool shouldOptimize(const BinaryFunction &BF) const override;
-
-  const char *getName() const override { return "split-functions"; }
-
-  void runOnFunctions(BinaryContext &BC) override;
+                     const TrampolineSetType &Trampolines);
 };
 
 } // namespace bolt
diff --git a/bolt/lib/Passes/SplitFunctions.cpp b/bolt/lib/Passes/SplitFunctions.cpp
index 34973cecdf49161..223f8d17367845d 100644
--- a/bolt/lib/Passes/SplitFunctions.cpp
+++ b/bolt/lib/Passes/SplitFunctions.cpp
@@ -109,21 +109,9 @@ static cl::opt<SplitFunctionsStrategy> SplitStrategy(
 } // namespace opts
 
 namespace {
-bool hasFullProfile(const BinaryFunction &BF) {
-  return llvm::all_of(BF.blocks(), [](const BinaryBasicBlock &BB) {
-    return BB.getExecutionCount() != BinaryBasicBlock::COUNT_NO_PROFILE;
-  });
-}
-
-bool allBlocksCold(const BinaryFunction &BF) {
-  return llvm::all_of(BF.blocks(), [](const BinaryBasicBlock &BB) {
-    return BB.getExecutionCount() == 0;
-  });
-}
-
 struct SplitProfile2 final : public SplitStrategy {
   bool canSplit(const BinaryFunction &BF) override {
-    return BF.hasValidProfile() && hasFullProfile(BF) && !allBlocksCold(BF);
+    return BF.hasValidProfile() && BF.hasFullProfile() && !BF.allBlocksCold();
   }
 
   bool keepEmpty() override { return false; }
@@ -434,7 +422,7 @@ void SplitFunctions::splitFunction(BinaryFunction &BF, SplitStrategy &S) {
 }
 
 SplitFunctions::TrampolineSetType
-SplitFunctions::createEHTrampolines(BinaryFunction &BF) const {
+SplitFunctions::createEHTrampolines(BinaryFunction &BF) {
   const auto &MIB = BF.getBinaryContext().MIB;
 
   // Map real landing pads to the corresponding trampolines.
@@ -501,7 +489,7 @@ SplitFunctions::createEHTrampolines(BinaryFunction &BF) const {
 
 SplitFunctions::BasicBlockOrderType SplitFunctions::mergeEHTrampolines(
     BinaryFunction &BF, SplitFunctions::BasicBlockOrderType &Layout,
-    const SplitFunctions::TrampolineSetType &Trampolines) const {
+    const SplitFunctions::TrampolineSetType &Trampolines) {
   DenseMap<const MCSymbol *, SmallVector<const MCSymbol *, 0>>
       IncomingTrampolines;
   for (const auto &Entry : Trampolines) {