[BOLT] Generate sections for multiple fragments

This patch adds support to generate any number of sections that are
assigned to fragments of functions that are split more than two-way.
With this, a function's *nth* split fragment goes into section
`.text.cold.n`.

This also changes `FunctionLayout::erase` to make sure, that there are
no empty fragments at the end of the function. This sometimes happens
when blocks are erased from the function. To avoid creating symbols
pointing to these fragments, they need to be removed.

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D130521

bolt/include/bolt/Core/BinaryBasicBlock.h

@@ -681,11 +681,11 @@ public:
   bool isCold() const {
     assert(Fragment.get() < 2 &&
            "Function is split into more than two (hot/cold)-fragments");
-    return Fragment == FragmentNum::cold();
+    return isSplit();
   }
 
   void setIsCold(const bool Flag) {
-    Fragment = Flag ? FragmentNum::cold() : FragmentNum::hot();
+    Fragment = Flag ? FragmentNum::cold() : FragmentNum::main();
   }
 
   /// Return true if the block can be outlined. At the moment we disallow

bolt/include/bolt/Core/BinaryFunction.h

@@ -1078,8 +1078,8 @@ public:
 
   /// Return MC symbol associated with the function.
   /// All references to the function should use this symbol.
-  MCSymbol *getSymbol(const FragmentNum Fragment = FragmentNum::hot()) {
-    if (Fragment == FragmentNum::hot())
+  MCSymbol *getSymbol(const FragmentNum Fragment = FragmentNum::main()) {
+    if (Fragment == FragmentNum::main())
       return Symbols[0];
 
     size_t ColdSymbolIndex = Fragment.get() - 1;
@@ -1316,35 +1316,29 @@ public:
   }
 
   /// Return internal section name for this function.
-  StringRef getCodeSectionName() const { return StringRef(CodeSectionName); }
+  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);
+  }
 
   /// Assign a code section name to the function.
-  void setCodeSectionName(StringRef Name) {
-    CodeSectionName = std::string(Name);
+  void setCodeSectionName(const StringRef Name) {
+    CodeSectionName = Name.str();
   }
 
   /// Get output code section.
-  ErrorOr<BinarySection &> getCodeSection() const {
-    return BC.getUniqueSectionByName(getCodeSectionName());
-  }
-
-  /// Return cold code section name for the function.
-  std::string getColdCodeSectionName(const FragmentNum Fragment) const {
-    std::string Result = ColdCodeSectionName;
-    if (Fragment != FragmentNum::cold())
-      Result.append(".").append(std::to_string(Fragment.get() - 1));
-    return Result;
+  ErrorOr<BinarySection &>
+  getCodeSection(const FragmentNum Fragment = FragmentNum::main()) const {
+    return BC.getUniqueSectionByName(getCodeSectionName(Fragment));
   }
 
   /// Assign a section name for the cold part of the function.
-  void setColdCodeSectionName(StringRef Name) {
-    ColdCodeSectionName = std::string(Name);
-  }
-
-  /// Get output code section for cold code of this function.
-  ErrorOr<BinarySection &>
-  getColdCodeSection(const FragmentNum Fragment) const {
-    return BC.getUniqueSectionByName(getColdCodeSectionName(Fragment));
+  void setColdCodeSectionName(const StringRef Name) {
+    ColdCodeSectionName = Name.str();
   }
 
   /// Return true iif the function will halt execution on entry.
@@ -1878,11 +1872,9 @@ public:
     if (ColdCallSites.empty())
       return nullptr;
 
-    SmallString<8> SymbolSuffix;
-    if (Fragment != FragmentNum::cold())
-      SymbolSuffix = formatv(".{0}", Fragment.get());
-    ColdLSDASymbol = BC.Ctx->getOrCreateSymbol(formatv(
-        "GCC_cold_except_table{0:x-}{1}", getFunctionNumber(), SymbolSuffix));
+    ColdLSDASymbol =
+        BC.Ctx->getOrCreateSymbol(formatv("GCC_cold_except_table{0:x-}.{1}",
+                                          getFunctionNumber(), Fragment.get()));
 
     return ColdLSDASymbol;
   }

bolt/include/bolt/Core/FunctionLayout.h

@@ -47,7 +47,6 @@ public:
   }
 
   static constexpr FragmentNum main() { return FragmentNum(0); }
-  static constexpr FragmentNum hot() { return FragmentNum(0); }
   static constexpr FragmentNum cold() { return FragmentNum(1); }
 };
 
@@ -161,7 +160,7 @@ public:
   /// Get the fragment that contains all entry blocks and other blocks that
   /// cannot be split.
   FunctionFragment getMainFragment() const {
-    return getFragment(FragmentNum::hot());
+    return getFragment(FragmentNum::main());
   }
 
   /// Get the fragment that contains all entry blocks and other blocks that
@@ -181,7 +180,8 @@ public:
   void insertBasicBlocks(BinaryBasicBlock *InsertAfter,
                          ArrayRef<BinaryBasicBlock *> NewBlocks);
 
-  /// Erase all blocks from the layout that are in ToErase.
+  /// Erase all blocks from the layout that are in ToErase. If this method
+  /// erases all blocks of a fragment, it will be removed as well.
   void eraseBasicBlocks(const DenseSet<const BinaryBasicBlock *> ToErase);
 
   /// Make sure fragments' and basic blocks' indices match the current layout.

bolt/lib/Core/BinaryContext.cpp

@@ -2205,7 +2205,7 @@ BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
     SplitLabels.emplace_back(SplitStartLabel, SplitEndLabel);
 
     MCSectionELF *const SplitSection = LocalCtx->getELFSection(
-        BF.getColdCodeSectionName(FF.getFragmentNum()), ELF::SHT_PROGBITS,
+        BF.getCodeSectionName(FF.getFragmentNum()), ELF::SHT_PROGBITS,
         ELF::SHF_EXECINSTR | ELF::SHF_ALLOC);
     SplitSection->setHasInstructions(true);
     Streamer->switchSection(SplitSection);

bolt/lib/Core/BinaryEmitter.cpp

@@ -287,10 +287,8 @@ bool BinaryEmitter::emitFunction(BinaryFunction &Function,
   if (Function.getState() == BinaryFunction::State::Empty)
     return false;
 
-  MCSection *Section = BC.getCodeSection(
-      FF.isSplitFragment()
-          ? Function.getColdCodeSectionName(FF.getFragmentNum())
-          : Function.getCodeSectionName());
+  MCSection *Section =
+      BC.getCodeSection(Function.getCodeSectionName(FF.getFragmentNum()));
   Streamer.switchSection(Section);
   Section->setHasInstructions(true);
   BC.Ctx->addGenDwarfSection(Section);
@@ -408,7 +406,7 @@ void BinaryEmitter::emitFunctionBody(BinaryFunction &BF,
                                      const FunctionFragment &FF,
                                      bool EmitCodeOnly) {
   if (!EmitCodeOnly && FF.isSplitFragment() && BF.hasConstantIsland()) {
-    assert(FF.getFragmentNum() == FragmentNum::cold() &&
+    assert(BF.getLayout().isHotColdSplit() &&
            "Constant island support only with hot/cold split");
     BF.duplicateConstantIslands();
   }
@@ -918,7 +916,7 @@ void BinaryEmitter::emitLSDA(BinaryFunction &BF, bool EmitColdPart) {
 
   // Corresponding FDE start.
   const MCSymbol *StartSymbol =
-      EmitColdPart ? BF.getSymbol(FragmentNum::cold()) : BF.getSymbol();
+      BF.getSymbol(EmitColdPart ? FragmentNum::cold() : FragmentNum::main());
 
   // Emit the LSDA header.
 

bolt/lib/Core/BinaryFunction.cpp

@@ -4038,10 +4038,6 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
   }
 
   const uint64_t BaseAddress = getCodeSection()->getOutputAddress();
-  ErrorOr<BinarySection &> ColdSection =
-      getColdCodeSection(FragmentNum::cold());
-  const uint64_t ColdBaseAddress =
-      isSplit() ? ColdSection->getOutputAddress() : 0;
   if (BC.HasRelocations || isInjected()) {
     const uint64_t StartOffset = Layout.getSymbolOffset(*getSymbol());
     const uint64_t EndOffset = Layout.getSymbolOffset(*getFunctionEndLabel());
@@ -4053,20 +4049,35 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
       setOutputDataAddress(BaseAddress + DataOffset);
     }
     if (isSplit()) {
-      const MCSymbol *ColdStartSymbol = getSymbol(FragmentNum::cold());
-      assert(ColdStartSymbol && ColdStartSymbol->isDefined() &&
-             "split function should have defined cold symbol");
-      const MCSymbol *ColdEndSymbol = getFunctionEndLabel(FragmentNum::cold());
-      assert(ColdEndSymbol && ColdEndSymbol->isDefined() &&
-             "split function should have defined cold end symbol");
-      const uint64_t ColdStartOffset = Layout.getSymbolOffset(*ColdStartSymbol);
-      const uint64_t ColdEndOffset = Layout.getSymbolOffset(*ColdEndSymbol);
-      cold().setAddress(ColdBaseAddress + ColdStartOffset);
-      cold().setImageSize(ColdEndOffset - ColdStartOffset);
-      if (hasConstantIsland()) {
-        const uint64_t DataOffset =
-            Layout.getSymbolOffset(*getFunctionColdConstantIslandLabel());
-        setOutputColdDataAddress(ColdBaseAddress + DataOffset);
+      for (const FunctionFragment &FF : getLayout().getSplitFragments()) {
+        ErrorOr<BinarySection &> ColdSection =
+            getCodeSection(FF.getFragmentNum());
+        // If fragment is empty, cold section might not exist
+        if (FF.empty() && ColdSection.getError())
+          continue;
+        const uint64_t ColdBaseAddress = ColdSection->getOutputAddress();
+
+        const MCSymbol *ColdStartSymbol = getSymbol(FF.getFragmentNum());
+        // If fragment is empty, symbol might have not been emitted
+        if (FF.empty() && (!ColdStartSymbol || !ColdStartSymbol->isDefined()) &&
+            !hasConstantIsland())
+          continue;
+        assert(ColdStartSymbol && ColdStartSymbol->isDefined() &&
+               "split function should have defined cold symbol");
+        const MCSymbol *ColdEndSymbol =
+            getFunctionEndLabel(FF.getFragmentNum());
+        assert(ColdEndSymbol && ColdEndSymbol->isDefined() &&
+               "split function should have defined cold end symbol");
+        const uint64_t ColdStartOffset =
+            Layout.getSymbolOffset(*ColdStartSymbol);
+        const uint64_t ColdEndOffset = Layout.getSymbolOffset(*ColdEndSymbol);
+        cold().setAddress(ColdBaseAddress + ColdStartOffset);
+        cold().setImageSize(ColdEndOffset - ColdStartOffset);
+        if (hasConstantIsland()) {
+          const uint64_t DataOffset =
+              Layout.getSymbolOffset(*getFunctionColdConstantIslandLabel());
+          setOutputColdDataAddress(ColdBaseAddress + DataOffset);
+        }
       }
     }
   } else {
@@ -4088,32 +4099,42 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
   if (getLayout().block_empty())
     return;
 
+  assert((getLayout().isHotColdSplit() ||
+          (cold().getAddress() == 0 && cold().getImageSize() == 0 &&
+           BC.HasRelocations)) &&
+         "Function must be split two ways or cold fragment must have no "
+         "address (only in relocation mode)");
+
   BinaryBasicBlock *PrevBB = nullptr;
-  for (BinaryBasicBlock *BB : this->Layout.blocks()) {
-    assert(BB->getLabel()->isDefined() && "symbol should be defined");
-    const uint64_t BBBaseAddress = BB->isCold() ? ColdBaseAddress : BaseAddress;
-    if (!BC.HasRelocations) {
-      if (BB->isCold()) {
-        assert(BBBaseAddress == cold().getAddress());
-      } else {
-        assert(BBBaseAddress == getOutputAddress());
+  for (const FunctionFragment &FF : getLayout().fragments()) {
+    const uint64_t FragmentBaseAddress =
+        getCodeSection(isSimple() ? FF.getFragmentNum() : FragmentNum::main())
+            ->getOutputAddress();
+    for (BinaryBasicBlock *const BB : FF) {
+      assert(BB->getLabel()->isDefined() && "symbol should be defined");
+      if (!BC.HasRelocations) {
+        if (BB->isSplit()) {
+          assert(FragmentBaseAddress == cold().getAddress());
+        } else {
+          assert(FragmentBaseAddress == getOutputAddress());
+        }
       }
-    }
-    const uint64_t BBOffset = Layout.getSymbolOffset(*BB->getLabel());
-    const uint64_t BBAddress = BBBaseAddress + BBOffset;
-    BB->setOutputStartAddress(BBAddress);
+      const uint64_t BBOffset = Layout.getSymbolOffset(*BB->getLabel());
+      const uint64_t BBAddress = FragmentBaseAddress + BBOffset;
+      BB->setOutputStartAddress(BBAddress);
 
-    if (PrevBB) {
-      uint64_t PrevBBEndAddress = BBAddress;
-      if (BB->isCold() != PrevBB->isCold())
-        PrevBBEndAddress = getOutputAddress() + getOutputSize();
-      PrevBB->setOutputEndAddress(PrevBBEndAddress);
-    }
-    PrevBB = BB;
+      if (PrevBB) {
+        uint64_t PrevBBEndAddress = BBAddress;
+        if (BB->isSplit() != PrevBB->isSplit())
+          PrevBBEndAddress = getOutputAddress() + getOutputSize();
+        PrevBB->setOutputEndAddress(PrevBBEndAddress);
+      }
+      PrevBB = BB;
 
-    BB->updateOutputValues(Layout);
+      BB->updateOutputValues(Layout);
+    }
   }
-  PrevBB->setOutputEndAddress(PrevBB->isCold()
+  PrevBB->setOutputEndAddress(PrevBB->isSplit()
                                   ? cold().getAddress() + cold().getImageSize()
                                   : getOutputAddress() + getOutputSize());
 }

bolt/lib/Core/FunctionLayout.cpp

@@ -3,6 +3,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/edit_distance.h"
 #include <algorithm>
+#include <cstddef>
 #include <functional>
 
 using namespace llvm;
@@ -61,18 +62,35 @@ void FunctionLayout::eraseBasicBlocks(
   NewFragments.emplace_back(0);
   for (const FunctionFragment FF : fragments()) {
     unsigned ErasedBlocks = count_if(FF, IsErased);
+    // Only add the fragment if it is non-empty after removing blocks.
     unsigned NewFragment = NewFragments.back() + FF.size() - ErasedBlocks;
     NewFragments.emplace_back(NewFragment);
   }
   llvm::erase_if(Blocks, IsErased);
   Fragments = std::move(NewFragments);
+
+  // Remove empty fragments at the end
+  const_iterator EmptyTailBegin =
+      llvm::find_if_not(reverse(fragments()), [](const FunctionFragment &FF) {
+        return FF.empty();
+      }).base();
+  if (EmptyTailBegin != fragment_end()) {
+    // Add +1 for one-past-the-end entry
+    const FunctionFragment TailBegin = *EmptyTailBegin;
+    unsigned NewFragmentSize = TailBegin.getFragmentNum().get() + 1;
+    Fragments.resize(NewFragmentSize);
+  }
+
+  updateLayoutIndices();
 }
 
 void FunctionLayout::updateLayoutIndices() const {
   unsigned BlockIndex = 0;
   for (const FunctionFragment FF : fragments()) {
-    for (BinaryBasicBlock *const BB : FF)
+    for (BinaryBasicBlock *const BB : FF) {
       BB->setLayoutIndex(BlockIndex++);
+      BB->setFragmentNum(FF.getFragmentNum());
+    }
   }
 }
 

bolt/lib/Passes/SplitFunctions.cpp

@@ -309,7 +309,7 @@ void SplitFunctions::splitFunction(BinaryFunction &BF, SplitStrategy Strategy) {
         PreSplitLayout = mergeEHTrampolines(BF, PreSplitLayout, Trampolines);
 
       for (BinaryBasicBlock &BB : BF)
-        BB.setIsCold(false);
+        BB.setFragmentNum(FragmentNum::main());
       BF.getLayout().update(PreSplitLayout);
     } else {
       SplitBytesHot += HotSize;

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -3186,11 +3186,11 @@ void RewriteInstance::emitAndLink() {
     for (const FunctionFragment FF :
          Function->getLayout().getSplitFragments()) {
       if (ErrorOr<BinarySection &> ColdSection =
-              Function->getColdCodeSection(FF.getFragmentNum()))
+              Function->getCodeSection(FF.getFragmentNum()))
         BC->deregisterSection(*ColdSection);
     }
     if (Function->getLayout().isSplit())
-      Function->ColdCodeSectionName = getBOLTTextSectionName().str();
+      Function->setColdCodeSectionName(getBOLTTextSectionName());
   }
 
   if (opts::PrintCacheMetrics) {
@@ -3728,7 +3728,7 @@ void RewriteInstance::mapCodeSections(RuntimeDyld &RTDyld) {
 
     for (const FunctionFragment FF : Function.getLayout().getSplitFragments()) {
       ErrorOr<BinarySection &> ColdSection =
-          Function.getColdCodeSection(FF.getFragmentNum());
+          Function.getCodeSection(FF.getFragmentNum());
       assert(ColdSection && "cannot find section for cold part");
       // Cold fragments are aligned at 16 bytes.
       NextAvailableAddress = alignTo(NextAvailableAddress, 16);
@@ -4521,12 +4521,12 @@ void RewriteInstance::updateELFSymbolTable(
       for (const FunctionFragment FF :
            Function.getLayout().getSplitFragments()) {
         ELFSymTy NewColdSym = FunctionSymbol;
-        const SmallString<256> Buf = formatv(
+        const SmallString<256> SymbolName = formatv(
             "{0}.cold.{1}", cantFail(FunctionSymbol.getName(StringSection)),
             FF.getFragmentNum().get() - 1);
-        NewColdSym.st_name = AddToStrTab(Buf);
+        NewColdSym.st_name = AddToStrTab(SymbolName);
         NewColdSym.st_shndx =
-            Function.getColdCodeSection(FF.getFragmentNum())->getIndex();
+            Function.getCodeSection(FF.getFragmentNum())->getIndex();
         NewColdSym.st_value = Function.cold().getAddress();
         NewColdSym.st_size = Function.cold().getImageSize();
         NewColdSym.setBindingAndType(ELF::STB_LOCAL, ELF::STT_FUNC);
@@ -4658,7 +4658,7 @@ void RewriteInstance::updateELFSymbolTable(
         NewSymbol.st_shndx =
             OutputAddress >= Function->cold().getAddress() &&
                     OutputAddress < Function->cold().getImageSize()
-                ? Function->getColdCodeSection(FragmentNum::cold())->getIndex()
+                ? Function->getCodeSection(FragmentNum::cold())->getIndex()
                 : Function->getCodeSection()->getIndex();
       } else {
         // Check if the symbol belongs to moved data object and update it.
@@ -4756,6 +4756,9 @@ void RewriteInstance::updateELFSymbolTable(
     Symbols.emplace_back(NewSymbol);
 
     if (Function->isSplit()) {
+      assert(Function->getLayout().isHotColdSplit() &&
+             "Adding symbols based on cold fragment when there are more than "
+             "2 fragments");
       ELFSymTy NewColdSym = NewSymbol;
       NewColdSym.setType(ELF::STT_NOTYPE);
       SmallVector<char, 256> Buf;