[BOLT] Make exception handling fragment aware

This adds basic fragment awareness in the exception handling passes and
generates the necessary symbols for fragments.

Reviewed By: rafauler

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

Author: Fabian Parzefall

bolt/include/bolt/Core/BinaryBasicBlock.h

@@ -133,9 +133,9 @@ class BinaryBasicBlock {
   /// CFI state at the entry to this basic block.
   int32_t CFIState{-1};
 
-  /// In cases where the parent function has been split, IsCold == true means
-  /// this BB will be allocated outside its parent function.
-  bool IsCold{false};
+  /// In cases where the parent function has been split, FragmentNum > 0 means
+  /// this BB will be allocated in a fragment outside its parent function.
+  FragmentNum Fragment;
 
   /// Indicates if the block could be outlined.
   bool CanOutline{true};
@@ -672,13 +672,21 @@ class BinaryBasicBlock {
 
   void markValid(const bool Valid) { IsValid = Valid; }
 
-  FragmentNum getFragmentNum() const {
-    return IsCold ? FragmentNum::cold() : FragmentNum::hot();
-  }
+  FragmentNum getFragmentNum() const { return Fragment; }
+
+  void setFragmentNum(const FragmentNum Value) { Fragment = Value; }
 
-  bool isCold() const { return IsCold; }
+  bool isSplit() const { return Fragment != FragmentNum::main(); }
 
-  void setIsCold(const bool Flag) { IsCold = Flag; }
+  bool isCold() const {
+    assert(Fragment.get() < 2 &&
+           "Function is split into more than two (hot/cold)-fragments");
+    return Fragment == FragmentNum::cold();
+  }
+
+  void setIsCold(const bool Flag) {
+    Fragment = Flag ? FragmentNum::cold() : FragmentNum::hot();
+  }
 
   /// 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/BinaryFunction.h

@@ -570,11 +570,8 @@ class BinaryFunction {
 
   SmallVector<MCSymbol *, 0> ColdSymbols;
 
-  /// Symbol at the end of the function.
-  mutable MCSymbol *FunctionEndLabel{nullptr};
-
-  /// Symbol at the end of the cold part of split function.
-  mutable MCSymbol *FunctionColdEndLabel{nullptr};
+  /// Symbol at the end of each fragment of a split function.
+  mutable SmallVector<MCSymbol *, 0> FunctionEndLabels;
 
   /// Unique number associated with the function.
   uint64_t FunctionNumber;
@@ -1152,22 +1149,25 @@ class BinaryFunction {
   bool forEachEntryPoint(EntryPointCallbackTy Callback) const;
 
   /// Return MC symbol associated with the end of the function.
-  MCSymbol *getFunctionEndLabel() const {
+  MCSymbol *
+  getFunctionEndLabel(const FragmentNum Fragment = FragmentNum::main()) const {
     assert(BC.Ctx && "cannot be called with empty context");
-    if (!FunctionEndLabel) {
-      std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
-      FunctionEndLabel = BC.Ctx->createNamedTempSymbol("func_end");
+
+    size_t LabelIndex = Fragment.get();
+    if (LabelIndex >= FunctionEndLabels.size()) {
+      FunctionEndLabels.resize(LabelIndex + 1);
     }
-    return FunctionEndLabel;
-  }
 
-  /// Return MC symbol associated with the end of the cold part of the function.
-  MCSymbol *getFunctionColdEndLabel() const {
-    if (!FunctionColdEndLabel) {
+    MCSymbol *&FunctionEndLabel = FunctionEndLabels[LabelIndex];
+    if (!FunctionEndLabel) {
       std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
-      FunctionColdEndLabel = BC.Ctx->createNamedTempSymbol("func_cold_end");
+      if (Fragment == FragmentNum::main())
+        FunctionEndLabel = BC.Ctx->createNamedTempSymbol("func_end");
+      else
+        FunctionEndLabel = BC.Ctx->createNamedTempSymbol(
+            formatv("func_cold_end.{0}", Fragment.get() - 1));
     }
-    return FunctionColdEndLabel;
+    return FunctionEndLabel;
   }
 
   /// Return a label used to identify where the constant island was emitted
@@ -1872,14 +1872,17 @@ class BinaryFunction {
   }
 
   /// Return symbol pointing to function's LSDA for the cold part.
-  MCSymbol *getColdLSDASymbol() {
+  MCSymbol *getColdLSDASymbol(const FragmentNum Fragment) {
     if (ColdLSDASymbol)
       return ColdLSDASymbol;
     if (ColdCallSites.empty())
       return nullptr;
 
-    ColdLSDASymbol = BC.Ctx->getOrCreateSymbol(
-        Twine("GCC_cold_except_table") + Twine::utohexstr(getFunctionNumber()));
+    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));
 
     return ColdLSDASymbol;
   }

bolt/include/bolt/Core/FunctionLayout.h

@@ -46,6 +46,7 @@ class FragmentNum {
     return !(*this == Other);
   }
 
+  static constexpr FragmentNum main() { return FragmentNum(0); }
   static constexpr FragmentNum hot() { return FragmentNum(0); }
   static constexpr FragmentNum cold() { return FragmentNum(1); }
 };

bolt/lib/Core/BinaryBasicBlock.cpp

@@ -102,7 +102,7 @@ bool BinaryBasicBlock::validateSuccessorInvariants() {
     if (Valid) {
       for (const MCSymbol *Sym : UniqueSyms) {
         Valid &= (Sym == Function->getFunctionEndLabel() ||
-                  Sym == Function->getFunctionColdEndLabel());
+                  Sym == Function->getFunctionEndLabel(getFragmentNum()));
         if (!Valid) {
           errs() << "BOLT-WARNING: Jump table contains illegal entry: "
                  << Sym->getName() << "\n";

bolt/lib/Core/BinaryEmitter.cpp

@@ -331,14 +331,13 @@ bool BinaryEmitter::emitFunction(BinaryFunction &Function,
 
   // Emit CFI start
   if (Function.hasCFI()) {
-    assert(Function.getLayout().isHotColdSplit() &&
-           "Exceptions supported only with hot/cold splitting.");
     Streamer.emitCFIStartProc(/*IsSimple=*/false);
     if (Function.getPersonalityFunction() != nullptr)
       Streamer.emitCFIPersonality(Function.getPersonalityFunction(),
                                   Function.getPersonalityEncoding());
-    MCSymbol *LSDASymbol = FF.isSplitFragment() ? Function.getColdLSDASymbol()
-                                                : Function.getLSDASymbol();
+    MCSymbol *LSDASymbol = FF.isSplitFragment()
+                               ? Function.getColdLSDASymbol(FF.getFragmentNum())
+                               : Function.getLSDASymbol();
     if (LSDASymbol)
       Streamer.emitCFILsda(LSDASymbol, BC.LSDAEncoding);
     else
@@ -383,9 +382,7 @@ bool BinaryEmitter::emitFunction(BinaryFunction &Function,
   if (Function.hasCFI())
     Streamer.emitCFIEndProc();
 
-  MCSymbol *EndSymbol = FF.isSplitFragment()
-                            ? Function.getFunctionColdEndLabel()
-                            : Function.getFunctionEndLabel();
+  MCSymbol *EndSymbol = Function.getFunctionEndLabel(FF.getFragmentNum());
   Streamer.emitLabel(EndSymbol);
 
   if (MAI->hasDotTypeDotSizeDirective()) {
@@ -913,8 +910,9 @@ void BinaryEmitter::emitLSDA(BinaryFunction &BF, bool EmitColdPart) {
   Streamer.emitValueToAlignment(TTypeAlignment);
 
   // Emit the LSDA label.
-  MCSymbol *LSDASymbol =
-      EmitColdPart ? BF.getColdLSDASymbol() : BF.getLSDASymbol();
+  MCSymbol *LSDASymbol = EmitColdPart
+                             ? BF.getColdLSDASymbol(FragmentNum::cold())
+                             : BF.getLSDASymbol();
   assert(LSDASymbol && "no LSDA symbol set");
   Streamer.emitLabel(LSDASymbol);
 

bolt/lib/Core/BinaryFunction.cpp

@@ -31,6 +31,7 @@
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/MC/MCSymbol.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -1572,8 +1573,9 @@ bool BinaryFunction::scanExternalRefs() {
   if (BC.HasRelocations) {
     for (std::pair<const uint32_t, MCSymbol *> &LI : Labels)
       BC.UndefinedSymbols.insert(LI.second);
-    if (FunctionEndLabel)
-      BC.UndefinedSymbols.insert(FunctionEndLabel);
+    for (MCSymbol *const EndLabel : FunctionEndLabels)
+      if (EndLabel)
+        BC.UndefinedSymbols.insert(EndLabel);
   }
 
   clearList(Relocations);
@@ -2843,8 +2845,9 @@ void BinaryFunction::clearDisasmState() {
   if (BC.HasRelocations) {
     for (std::pair<const uint32_t, MCSymbol *> &LI : Labels)
       BC.UndefinedSymbols.insert(LI.second);
-    if (FunctionEndLabel)
-      BC.UndefinedSymbols.insert(FunctionEndLabel);
+    for (MCSymbol *const EndLabel : FunctionEndLabels)
+      if (EndLabel)
+        BC.UndefinedSymbols.insert(EndLabel);
   }
 }
 
@@ -4053,7 +4056,7 @@ void BinaryFunction::updateOutputValues(const MCAsmLayout &Layout) {
       const MCSymbol *ColdStartSymbol = getSymbol(FragmentNum::cold());
       assert(ColdStartSymbol && ColdStartSymbol->isDefined() &&
              "split function should have defined cold symbol");
-      const MCSymbol *ColdEndSymbol = getFunctionColdEndLabel();
+      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);