[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

bolt/include/bolt/Core/BinaryBasicBlock.h

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

@@ -570,11 +570,8 @@ private:
 
   SmallVector<MCSymbol *, 0> ColdSymbols;
 
-  /// Symbol at the end of the function.
+  /// Symbol at the end of each fragment of a split function.
-  mutable MCSymbol *FunctionEndLabel{nullptr};
+  mutable SmallVector<MCSymbol *, 0> FunctionEndLabels;
-
-  /// Symbol at the end of the cold part of split function.
-  mutable MCSymbol *FunctionColdEndLabel{nullptr};
 
   /// Unique number associated with the function.
   uint64_t FunctionNumber;
@@ -1152,24 +1149,27 @@ public:
   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");
+
+    size_t LabelIndex = Fragment.get();
+    if (LabelIndex >= FunctionEndLabels.size()) {
+      FunctionEndLabels.resize(LabelIndex + 1);
+    }
+
+    MCSymbol *&FunctionEndLabel = FunctionEndLabels[LabelIndex];
     if (!FunctionEndLabel) {
       std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
-      FunctionEndLabel = BC.Ctx->createNamedTempSymbol("func_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 FunctionEndLabel;
   }
 
-  /// Return MC symbol associated with the end of the cold part of the function.
-  MCSymbol *getFunctionColdEndLabel() const {
-    if (!FunctionColdEndLabel) {
-      std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
-      FunctionColdEndLabel = BC.Ctx->createNamedTempSymbol("func_cold_end");
-    }
-    return FunctionColdEndLabel;
-  }
-
   /// Return a label used to identify where the constant island was emitted
   /// (AArch only). This is used to update the symbol table accordingly,
   /// emitting data marker symbols as required by the ABI.
@@ -1872,14 +1872,17 @@ public:
   }
 
   /// 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(
+    SmallString<8> SymbolSuffix;
-        Twine("GCC_cold_except_table") + Twine::utohexstr(getFunctionNumber()));
+    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 @@ public:
     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()
+    MCSymbol *LSDASymbol = FF.isSplitFragment()
-                                                : Function.getLSDASymbol();
+                               ? 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()
+  MCSymbol *EndSymbol = Function.getFunctionEndLabel(FF.getFragmentNum());
-                            ? Function.getFunctionColdEndLabel()
-                            : Function.getFunctionEndLabel();
   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 =
+  MCSymbol *LSDASymbol = EmitColdPart
-      EmitColdPart ? BF.getColdLSDASymbol() : BF.getLSDASymbol();
+                             ? 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)
+    for (MCSymbol *const EndLabel : FunctionEndLabels)
-      BC.UndefinedSymbols.insert(FunctionEndLabel);
+      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)
+    for (MCSymbol *const EndLabel : FunctionEndLabels)
-      BC.UndefinedSymbols.insert(FunctionEndLabel);
+      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);

bolt/lib/Core/Exceptions.cpp

@@ -367,114 +367,96 @@ void BinaryFunction::updateEHRanges() {
     uint64_t Action;
   };
 
-  // If previous call can throw, this is its exception handler.
+  for (const FunctionFragment FF : getLayout().fragments()) {
-  EHInfo PreviousEH = {nullptr, 0};
+    // Sites to update - either regular or cold.
+    CallSitesType &Sites = FF.isMainFragment() ? CallSites : ColdCallSites;
 
-  // Marker for the beginning of exceptions range.
+    // If previous call can throw, this is its exception handler.
-  const MCSymbol *StartRange = nullptr;
+    EHInfo PreviousEH = {nullptr, 0};
 
-  // Indicates whether the start range is located in a cold part.
+    // Marker for the beginning of exceptions range.
-  bool IsStartInCold = false;
+    const MCSymbol *StartRange = nullptr;
 
-  // Have we crossed hot/cold border for split functions?
+    for (BinaryBasicBlock *const BB : FF) {
-  bool SeenCold = false;
+      for (auto II = BB->begin(); II != BB->end(); ++II) {
+        if (!BC.MIB->isCall(*II))
+          continue;
 
-  // Sites to update - either regular or cold.
+        // Instruction can throw an exception that should be handled.
-  CallSitesType *Sites = &CallSites;
+        const bool Throws = BC.MIB->isInvoke(*II);
 
-  for (BinaryBasicBlock *BB : getLayout().blocks()) {
+        // Ignore the call if it's a continuation of a no-throw gap.
+        if (!Throws && !StartRange)
+          continue;
 
-    if (BB->isCold() && !SeenCold) {
+        assert(getLayout().isHotColdSplit() &&
-      SeenCold = true;
+               "Exceptions only supported for hot/cold splitting");
 
-      // Close the range (if any) and change the target call sites.
+        // Extract exception handling information from the instruction.
-      if (StartRange) {
+        const MCSymbol *LP = nullptr;
-        Sites->emplace_back(CallSite{StartRange, getFunctionEndLabel(),
+        uint64_t Action = 0;
-                                     PreviousEH.LP, PreviousEH.Action});
+        if (const Optional<MCPlus::MCLandingPad> EHInfo =
-      }
+                BC.MIB->getEHInfo(*II))
-      Sites = &ColdCallSites;
+          std::tie(LP, Action) = *EHInfo;
 
-      // Reset the range.
+        // No action if the exception handler has not changed.
-      StartRange = nullptr;
+        if (Throws && StartRange && PreviousEH.LP == LP &&
-      PreviousEH = {nullptr, 0};
+            PreviousEH.Action == Action)
-    }
+          continue;
 
-    for (auto II = BB->begin(); II != BB->end(); ++II) {
+        // Same symbol is used for the beginning and the end of the range.
-      if (!BC.MIB->isCall(*II))
+        const MCSymbol *EHSymbol;
-        continue;
+        MCInst EHLabel;
+        {
+          std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
+          EHSymbol = BC.Ctx->createNamedTempSymbol("EH");
+          BC.MIB->createEHLabel(EHLabel, EHSymbol, BC.Ctx.get());
+        }
 
-      // Instruction can throw an exception that should be handled.
+        II = std::next(BB->insertPseudoInstr(II, EHLabel));
-      const bool Throws = BC.MIB->isInvoke(*II);
 
-      // Ignore the call if it's a continuation of a no-throw gap.
+        // At this point we could be in one of the following states:
-      if (!Throws && !StartRange)
+        //
-        continue;
+        // I. Exception handler has changed and we need to close previous range
+        //    and start a new one.
+        //
+        // II. Start a new exception range after the gap.
+        //
+        // III. Close current exception range and start a new gap.
+        const MCSymbol *EndRange;
+        if (StartRange) {
+          // I, III:
+          EndRange = EHSymbol;
+        } else {
+          // II:
+          StartRange = EHSymbol;
+          EndRange = nullptr;
+        }
 
-      // Extract exception handling information from the instruction.
+        // Close the previous range.
-      const MCSymbol *LP = nullptr;
+        if (EndRange) {
-      uint64_t Action = 0;
+          Sites.emplace_back(
-      if (const Optional<MCPlus::MCLandingPad> EHInfo = BC.MIB->getEHInfo(*II))
+              CallSite{StartRange, EndRange, PreviousEH.LP, PreviousEH.Action});
-        std::tie(LP, Action) = *EHInfo;
+        }
 
-      // No action if the exception handler has not changed.
+        if (Throws) {
-      if (Throws && StartRange && PreviousEH.LP == LP &&
+          // I, II:
-          PreviousEH.Action == Action)
+          StartRange = EHSymbol;
-        continue;
+          PreviousEH = EHInfo{LP, Action};
-
+        } else {
-      // Same symbol is used for the beginning and the end of the range.
+          StartRange = nullptr;
-      const MCSymbol *EHSymbol;
+        }
-      MCInst EHLabel;
-      {
-        std::unique_lock<std::shared_timed_mutex> Lock(BC.CtxMutex);
-        EHSymbol = BC.Ctx->createNamedTempSymbol("EH");
-        BC.MIB->createEHLabel(EHLabel, EHSymbol, BC.Ctx.get());
-      }
-
-      II = std::next(BB->insertPseudoInstr(II, EHLabel));
-
-      // At this point we could be in one of the following states:
-      //
-      // I. Exception handler has changed and we need to close previous range
-      //    and start a new one.
-      //
-      // II. Start a new exception range after the gap.
-      //
-      // III. Close current exception range and start a new gap.
-      const MCSymbol *EndRange;
-      if (StartRange) {
-        // I, III:
-        EndRange = EHSymbol;
-      } else {
-        // II:
-        StartRange = EHSymbol;
-        IsStartInCold = SeenCold;
-        EndRange = nullptr;
-      }
-
-      // Close the previous range.
-      if (EndRange) {
-        Sites->emplace_back(
-            CallSite{StartRange, EndRange, PreviousEH.LP, PreviousEH.Action});
-      }
-
-      if (Throws) {
-        // I, II:
-        StartRange = EHSymbol;
-        IsStartInCold = SeenCold;
-        PreviousEH = EHInfo{LP, Action};
-      } else {
-        StartRange = nullptr;
       }
     }
-  }
 
-  // Check if we need to close the range.
+    // Check if we need to close the range.
-  if (StartRange) {
+    if (StartRange) {
-    assert((!isSplit() || Sites == &ColdCallSites) && "sites mismatch");
+      assert((FF.isMainFragment() || &Sites == &ColdCallSites) &&
-    const MCSymbol *EndRange =
+             "sites mismatch");
-        IsStartInCold ? getFunctionColdEndLabel() : getFunctionEndLabel();
+      const MCSymbol *EndRange = getFunctionEndLabel(FF.getFragmentNum());
-    Sites->emplace_back(
+      Sites.emplace_back(
-        CallSite{StartRange, EndRange, PreviousEH.LP, PreviousEH.Action});
+          CallSite{StartRange, EndRange, PreviousEH.LP, PreviousEH.Action});
+    }
   }
 }
 

bolt/lib/Passes/Aligner.cpp

@@ -84,7 +84,7 @@ void alignCompact(BinaryFunction &Function, const MCCodeEmitter *Emitter) {
   size_t ColdSize = 0;
 
   for (const BinaryBasicBlock &BB : Function)
-    if (BB.isCold())
+    if (BB.isSplit())
       ColdSize += BC.computeCodeSize(BB.begin(), BB.end(), Emitter);
     else
       HotSize += BC.computeCodeSize(BB.begin(), BB.end(), Emitter);

bolt/lib/Passes/BinaryPasses.cpp

@@ -592,41 +592,39 @@ void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
 
   for (auto &It : BC.getBinaryFunctions()) {
     BinaryFunction &BF = It.second;
-    int64_t CurrentGnuArgsSize = 0;
 
-    // Have we crossed hot/cold border for split functions?
+    for (const FunctionFragment FF : BF.getLayout().fragments()) {
-    bool SeenCold = false;
+      int64_t CurrentGnuArgsSize = 0;
 
-    for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
+      for (BinaryBasicBlock *const BB : FF) {
-      if (BB->isCold() && !SeenCold) {
+        // First convert GnuArgsSize annotations into CFIs. This may change
-        SeenCold = true;
+        // instr pointers, so do it before recording ptrs for preserved
-        CurrentGnuArgsSize = 0;
+        // annotations
-      }
+        if (BF.usesGnuArgsSize()) {
-
+          for (auto II = BB->begin(); II != BB->end(); ++II) {
-      // First convert GnuArgsSize annotations into CFIs. This may change instr
+            if (!BC.MIB->isInvoke(*II))
-      // pointers, so do it before recording ptrs for preserved annotations
+              continue;
-      if (BF.usesGnuArgsSize()) {
+            const int64_t NewGnuArgsSize = BC.MIB->getGnuArgsSize(*II);
-        for (auto II = BB->begin(); II != BB->end(); ++II) {
+            assert(NewGnuArgsSize >= 0 &&
-          if (!BC.MIB->isInvoke(*II))
+                   "expected non-negative GNU_args_size");
-            continue;
+            if (NewGnuArgsSize != CurrentGnuArgsSize) {
-          const int64_t NewGnuArgsSize = BC.MIB->getGnuArgsSize(*II);
+              auto InsertII = BF.addCFIInstruction(
-          assert(NewGnuArgsSize >= 0 && "expected non-negative GNU_args_size");
+                  BB, II,
-          if (NewGnuArgsSize != CurrentGnuArgsSize) {
+                  MCCFIInstruction::createGnuArgsSize(nullptr, NewGnuArgsSize));
-            auto InsertII = BF.addCFIInstruction(
+              CurrentGnuArgsSize = NewGnuArgsSize;
-                BB, II,
+              II = std::next(InsertII);
-                MCCFIInstruction::createGnuArgsSize(nullptr, NewGnuArgsSize));
+            }
-            CurrentGnuArgsSize = NewGnuArgsSize;
-            II = std::next(InsertII);
           }
         }
-      }
 
-      // Now record preserved annotations separately and then strip annotations.
+        // Now record preserved annotations separately and then strip
-      for (auto II = BB->begin(); II != BB->end(); ++II) {
+        // annotations.
-        if (BF.requiresAddressTranslation() && BC.MIB->getOffset(*II))
+        for (auto II = BB->begin(); II != BB->end(); ++II) {
-          PreservedOffsetAnnotations.emplace_back(&(*II),
+          if (BF.requiresAddressTranslation() && BC.MIB->getOffset(*II))
-                                                  *BC.MIB->getOffset(*II));
+            PreservedOffsetAnnotations.emplace_back(&(*II),
-        BC.MIB->stripAnnotations(*II);
+                                                    *BC.MIB->getOffset(*II));
+          BC.MIB->stripAnnotations(*II);
+        }
       }
     }
   }

bolt/lib/Passes/IndirectCallPromotion.cpp

@@ -257,9 +257,9 @@ IndirectCallPromotion::getCallTargets(BinaryBasicBlock &BB,
       MCSymbol *Entry = JT->Entries[I];
       assert(BF.getBasicBlockForLabel(Entry) ||
              Entry == BF.getFunctionEndLabel() ||
-             Entry == BF.getFunctionColdEndLabel());
+             Entry == BF.getFunctionEndLabel(FragmentNum::cold()));
       if (Entry == BF.getFunctionEndLabel() ||
-          Entry == BF.getFunctionColdEndLabel())
+          Entry == BF.getFunctionEndLabel(FragmentNum::cold()))
         continue;
       const Location To(Entry);
       const BinaryBasicBlock::BinaryBranchInfo &BI = BB.getBranchInfo(Entry);