[BOLT] Introduce .text.warm for -use-cdsplit=1

bolt/include/bolt/Core/BinaryBasicBlock.h

@@ -677,9 +677,7 @@ class BinaryBasicBlock {
     return isSplit();
   }
 
-  void setIsCold(const bool Flag) {
-    Fragment = Flag ? FragmentNum::cold() : FragmentNum::main();
-  }
+  void setIsCold(const bool Flag);
 
   /// Return true if the block can be outlined. At the moment we disallow
   /// outlining of blocks that can potentially throw exceptions or are

bolt/include/bolt/Core/BinaryContext.h

@@ -927,6 +927,8 @@ class BinaryContext {
 
   const char *getMainCodeSectionName() const { return ".text"; }
 
+  const char *getWarmCodeSectionName() const { return ".text.warm"; }
+
   const char *getColdCodeSectionName() const { return ".text.cold"; }
 
   const char *getHotTextMoverSectionName() const { return ".text.mover"; }
@@ -1230,6 +1232,9 @@ class BinaryContext {
   ///
   /// Return the pair where the first size is for the main part, and the second
   /// size is for the cold one.
+  /// Modify BinaryBasicBlock::OutputAddressRange for each basic block in the
+  /// function in place so that BB.OutputAddressRange.second less
+  /// BB.OutputAddressRange.first gives the emitted size of BB.
   std::pair<size_t, size_t> calculateEmittedSize(BinaryFunction &BF,
                                                  bool FixBranches = true);
 

bolt/include/bolt/Core/BinaryFunction.h

@@ -355,6 +355,9 @@ class BinaryFunction {
   /// Name for the section this function code should reside in.
   std::string CodeSectionName;
 
+  /// Name for the corresponding warm code section.
+  std::string WarmCodeSectionName;
+
   /// Name for the corresponding cold code section.
   std::string ColdCodeSectionName;
 
@@ -1231,13 +1234,7 @@ class BinaryFunction {
 
   /// Return internal section name for this function.
   SmallString<32>
-  getCodeSectionName(const FragmentNum Fragment = FragmentNum::main()) const {
-    if (Fragment == FragmentNum::main())
-      return SmallString<32>(CodeSectionName);
-    if (Fragment == FragmentNum::cold())
-      return SmallString<32>(ColdCodeSectionName);
-    return formatv("{0}.{1}", ColdCodeSectionName, Fragment.get() - 1);
-  }
+  getCodeSectionName(const FragmentNum Fragment = FragmentNum::main()) const;
 
   /// Assign a code section name to the function.
   void setCodeSectionName(const StringRef Name) {
@@ -1250,6 +1247,11 @@ class BinaryFunction {
     return BC.getUniqueSectionByName(getCodeSectionName(Fragment));
   }
 
+  /// Assign a section name for the warm part of the function.
+  void setWarmCodeSectionName(const StringRef Name) {
+    WarmCodeSectionName = Name.str();
+  }
+
   /// Assign a section name for the cold part of the function.
   void setColdCodeSectionName(const StringRef Name) {
     ColdCodeSectionName = Name.str();
@@ -1272,6 +1274,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.

bolt/include/bolt/Core/FunctionLayout.h

@@ -62,7 +62,10 @@ class FragmentNum {
   }
 
   static constexpr FragmentNum main() { return FragmentNum(0); }
-  static constexpr FragmentNum cold() { return FragmentNum(1); }
+  static constexpr FragmentNum warm() { return FragmentNum(1); }
+  static constexpr FragmentNum cold(bool Flag = false) {
+    return FragmentNum(Flag ? 2 : 1);
+  }
 };
 
 /// A freestanding subset of contiguous blocks of a function.

bolt/include/bolt/Passes/CDSplit.h

@@ -0,0 +1,63 @@
+//===- bolt/Passes/CDSplit.h - Split functions into hot/warm/cold
+// after function reordering pass -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_PASSES_CDSPLIT
+#define BOLT_PASSES_CDSPLIT
+
+#include "bolt/Passes/SplitFunctions.h"
+#include <atomic>
+
+namespace llvm {
+namespace bolt {
+
+using BasicBlockOrder = BinaryFunction::BasicBlockOrderType;
+
+class CDSplit : public BinaryFunctionPass {
+private:
+  /// Overall stats.
+  std::atomic<uint64_t> SplitBytesHot{0ull};
+  std::atomic<uint64_t> SplitBytesCold{0ull};
+
+  /// List of functions to be considered.
+  /// All functions in the list are used to construct a call graph.
+  /// A subset of functions in this list are considered for splitting.
+  std::vector<BinaryFunction *> FunctionsToConsider;
+
+  /// Helper functions to initialize global variables.
+  void initialize(BinaryContext &BC);
+
+  /// Split function body into 3 fragments: hot / warm / cold.
+  void runOnFunction(BinaryFunction &BF);
+
+  /// Assign each basic block in the given function to either hot, cold,
+  /// or warm fragment using the CDSplit algorithm.
+  void assignFragmentThreeWay(const BinaryFunction &BF,
+                              const BasicBlockOrder &BlockOrder);
+
+  /// Find the best split index that separates hot from warm.
+  /// The basic block whose index equals the returned split index will be the
+  /// last hot block.
+  size_t findSplitIndex(const BinaryFunction &BF,
+                        const BasicBlockOrder &BlockOrder);
+
+public:
+  explicit CDSplit(const cl::opt<bool> &PrintPass)
+      : BinaryFunctionPass(PrintPass) {}
+
+  bool shouldOptimize(const BinaryFunction &BF) const override;
+
+  const char *getName() const override { return "cdsplit"; }
+
+  void runOnFunctions(BinaryContext &BC) override;
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif

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

bolt/lib/Core/BinaryBasicBlock.cpp

@@ -15,17 +15,25 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/MC/MCInst.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 
 #define DEBUG_TYPE "bolt"
 
-namespace llvm {
-namespace bolt {
+using namespace llvm;
+using namespace bolt;
+namespace opts {
+extern cl::opt<bool> UseCDSplit;
+}
 
 constexpr uint32_t BinaryBasicBlock::INVALID_OFFSET;
 
-bool operator<(const BinaryBasicBlock &LHS, const BinaryBasicBlock &RHS) {
-  return LHS.Index < RHS.Index;
+bool bolt::operator<(const BinaryBasicBlock &LHS, const BinaryBasicBlock &RHS) {
+  return LHS.getIndex() < RHS.getIndex();
+}
+
+void BinaryBasicBlock::setIsCold(const bool Flag) {
+  Fragment = Flag ? FragmentNum::cold(opts::UseCDSplit) : FragmentNum::main();
 }
 
 bool BinaryBasicBlock::hasCFG() const { return getParent()->hasCFG(); }
@@ -611,6 +619,3 @@ BinaryBasicBlock *BinaryBasicBlock::splitAt(iterator II) {
 
   return NewBlock;
 }
-
-} // namespace bolt
-} // namespace llvm

bolt/lib/Core/BinaryContext.cpp

@@ -2331,14 +2331,37 @@ BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
   MCAsmLayout Layout(Assembler);
   Assembler.layout(Layout);
 
+  // Obtain fragment sizes.
+  std::vector<uint64_t> FragmentSizes;
+  // 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);
-                      });
+  FragmentSizes.push_back(HotSize);
+  // Split fragment sizes.
+  uint64_t ColdSize = 0;
+  for (const auto &Labels : SplitLabels) {
+    uint64_t Size = Layout.getSymbolOffset(*Labels.second) -
+                    Layout.getSymbolOffset(*Labels.first);
+    FragmentSizes.push_back(Size);
+    ColdSize += Size;
+  }
+
+  // Populate new start and end offsets of each basic block.
+  BinaryBasicBlock *PrevBB = nullptr;
+  uint64_t FragmentIndex = 0;
+  for (FunctionFragment &FF : BF.getLayout().fragments()) {
+    for (BinaryBasicBlock *BB : FF) {
+      const uint64_t BBStartOffset = Layout.getSymbolOffset(*(BB->getLabel()));
+      BB->setOutputStartAddress(BBStartOffset);
+      if (PrevBB)
+        PrevBB->setOutputEndAddress(BBStartOffset);
+      PrevBB = BB;
+    }
+    if (PrevBB)
+      PrevBB->setOutputEndAddress(FragmentSizes[FragmentIndex]);
+    FragmentIndex++;
+    PrevBB = nullptr;
+  }
 
   // Clean-up the effect of the code emission.
   for (const MCSymbol &Symbol : Assembler.symbols()) {

bolt/lib/Core/BinaryEmitter.cpp

@@ -34,6 +34,7 @@ namespace opts {
 
 extern cl::opt<JumpTableSupportLevel> JumpTables;
 extern cl::opt<bool> PreserveBlocksAlignment;
+extern cl::opt<bool> UseCDSplit;
 
 cl::opt<bool> AlignBlocks("align-blocks", cl::desc("align basic blocks"),
                           cl::cat(BoltOptCategory));
@@ -287,7 +288,10 @@ void BinaryEmitter::emitFunctions() {
 
   // Mark the end of hot text.
   if (opts::HotText) {
-    Streamer.switchSection(BC.getTextSection());
+    if (opts::UseCDSplit)
+      Streamer.switchSection(BC.getCodeSection(BC.getWarmCodeSectionName()));
+    else
+      Streamer.switchSection(BC.getTextSection());
     Streamer.emitLabel(BC.getHotTextEndSymbol());
   }
 }

bolt/lib/Core/BinaryFunction.cpp

@@ -59,6 +59,7 @@ extern cl::opt<bool> EnableBAT;
 extern cl::opt<bool> Instrument;
 extern cl::opt<bool> StrictMode;
 extern cl::opt<bool> UpdateDebugSections;
+extern cl::opt<bool> UseCDSplit;
 extern cl::opt<unsigned> Verbosity;
 
 extern bool processAllFunctions();
@@ -167,6 +168,18 @@ template <typename R> static bool emptyRange(const R &Range) {
   return Range.begin() == Range.end();
 }
 
+/// Return internal section name for this function.
+SmallString<32>
+BinaryFunction::getCodeSectionName(const FragmentNum Fragment) const {
+  if (Fragment == FragmentNum::main())
+    return SmallString<32>(CodeSectionName);
+  if (Fragment == FragmentNum::cold(opts::UseCDSplit))
+    return SmallString<32>(ColdCodeSectionName);
+  if (Fragment == FragmentNum::warm())
+    return SmallString<32>(WarmCodeSectionName);
+  return formatv("{0}.{1}", ColdCodeSectionName, Fragment.get() - 1);
+}
+
 /// Gets debug line information for the instruction located at the given
 /// address in the original binary. The SMLoc's pointer is used
 /// to point to this information, which is represented by a

bolt/lib/Passes/BinaryPasses.cpp

@@ -1244,8 +1244,10 @@ void AssignSections::runOnFunctions(BinaryContext &BC) {
     else
       Function.setCodeSectionName(BC.getColdCodeSectionName());
 
-    if (Function.isSplit())
+    if (Function.isSplit()) {
+      Function.setWarmCodeSectionName(BC.getWarmCodeSectionName());
       Function.setColdCodeSectionName(BC.getColdCodeSectionName());
+    }
   }
 }