[BOLT] Insert EH trampolines for multiple fragments

[FPar]

Summary

This patch adds exception handling trampolines when a function is split
into more than two fragments. Trampolines are tracked per-fragment, such
that they can be removed if splitting is reversed.

[rafauler]

LGTM

bolt/include/bolt/Core/FunctionLayout.h

 
 class FragmentNum {
   unsigned Value{0};
 
 public:
   constexpr FragmentNum() = default;
   constexpr explicit FragmentNum(unsigned Value) : Value(Value) {}
   constexpr unsigned get() const { return Value; }
+
   constexpr bool operator==(const FragmentNum Other) const {
     return Value == Other.Value;
   }
   constexpr bool operator!=(const FragmentNum Other) const {
-    return !(*this == Other);
+    return Value != Other.Value;
+  }
+  constexpr bool operator<(const FragmentNum Other) const {
+    return Value < Other.Value;
+  }
+  constexpr bool operator<=(const FragmentNum Other) const {
+    return Value <= Other.Value;
+  }
+  constexpr bool operator>=(const FragmentNum Other) const {
+    return Value >= Other.Value;
+  }
+  constexpr bool operator>(const FragmentNum Other) const {
+    return Value > Other.Value;
   }
 
   static constexpr FragmentNum main() { return FragmentNum(0); }
   static constexpr FragmentNum cold() { return FragmentNum(1); }
 };
 
 /// A freestanding subset of contiguous blocks of a function.
 class FunctionFragment {

bolt/include/bolt/Passes/SplitFunctions.h

 //===- bolt/Passes/SplitFunctions.h - Split function code -------*- 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_SPLIT_FUNCTIONS_H
 #define BOLT_PASSES_SPLIT_FUNCTIONS_H
 
+#include "bolt/Core/FunctionLayout.h"
 #include "bolt/Passes/BinaryPasses.h"
+#include "llvm/ADT/Hashing.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
 
 namespace llvm {
 namespace bolt {
 
 /// Strategy used to partition blocks into fragments.
 enum SplitFunctionsStrategy : char {

 
 /// Split function code in multiple parts.
 class SplitFunctions : public BinaryFunctionPass {
 private:
   /// Split function body into fragments.
   template <typename Strategy>
   void splitFunction(BinaryFunction &Function, Strategy S = {});
 
+  struct TrampolineKey {
+    FragmentNum SourceFN = FragmentNum::main();
+    const MCSymbol *Target = nullptr;
+
+    TrampolineKey() = default;
+    TrampolineKey(const FragmentNum SourceFN, const MCSymbol *const Target)
+        : SourceFN(SourceFN), Target(Target) {}
+
+    static inline TrampolineKey getEmptyKey() { return TrampolineKey(); };
+    static inline TrampolineKey getTombstoneKey() {
+      return TrampolineKey(FragmentNum(UINT_MAX), nullptr);
+    };
+    static unsigned getHashValue(const TrampolineKey &Val) {
+      return llvm::hash_combine(Val.SourceFN.get(), Val.Target);
+    }
+    static bool isEqual(const TrampolineKey &LHS, const TrampolineKey &RHS) {
+      return LHS.SourceFN == RHS.SourceFN && LHS.Target == RHS.Target;
+    }
+  };
+
   /// Map basic block labels to their trampoline block labels.
-  using TrampolineSetType = DenseMap<const MCSymbol *, const MCSymbol *>;
+  using TrampolineSetType =
+      DenseMap<TrampolineKey, const MCSymbol *, TrampolineKey>;
 
   using BasicBlockOrderType = BinaryFunction::BasicBlockOrderType;
 
   /// Create trampoline landing pads for exception handling code to guarantee
   /// that every landing pad is placed in the same function fragment as the
   /// corresponding thrower block. The trampoline landing pad, when created,
   /// will redirect the execution to the real landing pad in a different
   /// fragment.

bolt/lib/Passes/SplitFunctions.cpp

SplitFunctions::createEHTrampolines(BinaryFunction &BF) const {
   for (BinaryBasicBlock *BB : Blocks) {
     for (MCInst &Instr : *BB) {
       const Optional<MCPlus::MCLandingPad> EHInfo = MIB->getEHInfo(Instr);
       if (!EHInfo || !EHInfo->first)
         continue;
 
       const MCSymbol *LPLabel = EHInfo->first;
       BinaryBasicBlock *LPBlock = BF.getBasicBlockForLabel(LPLabel);
-      if (BB->isCold() == LPBlock->isCold())
+      if (BB->getFragmentNum() == LPBlock->getFragmentNum())
         continue;
 
       const MCSymbol *TrampolineLabel = nullptr;
-      auto Iter = LPTrampolines.find(LPLabel);
+      const TrampolineKey Key(BB->getFragmentNum(), LPLabel);
+      auto Iter = LPTrampolines.find(Key);
       if (Iter != LPTrampolines.end()) {
         TrampolineLabel = Iter->second;
       } else {
         // Create a trampoline basic block in the same fragment as the thrower.
         // Note: there's no need to insert the jump instruction, it will be
         // added by fixBranches().
         BinaryBasicBlock *TrampolineBB = BF.addBasicBlock();
-        TrampolineBB->setIsCold(BB->isCold());
+        TrampolineBB->setFragmentNum(BB->getFragmentNum());
         TrampolineBB->setExecutionCount(LPBlock->getExecutionCount());
         TrampolineBB->addSuccessor(LPBlock, TrampolineBB->getExecutionCount());
         TrampolineBB->setCFIState(LPBlock->getCFIState());
         TrampolineLabel = TrampolineBB->getLabel();
-        LPTrampolines.insert(std::make_pair(LPLabel, TrampolineLabel));
+        LPTrampolines.insert(std::make_pair(Key, TrampolineLabel));
       }
 
       // Substitute the landing pad with the trampoline.
       MIB->updateEHInfo(Instr,
                         MCPlus::MCLandingPad(TrampolineLabel, EHInfo->second));
     }
   }
 
   if (LPTrampolines.empty())
     return LPTrampolines;
 
   // All trampoline blocks were added to the end of the function. Place them at
   // the end of corresponding fragments.
   BinaryFunction::BasicBlockOrderType NewLayout(BF.getLayout().block_begin(),
                                                 BF.getLayout().block_end());
   stable_sort(NewLayout, [&](BinaryBasicBlock *A, BinaryBasicBlock *B) {
-    return A->isCold() < B->isCold();
+    return A->getFragmentNum() < B->getFragmentNum();
   });
   BF.getLayout().update(NewLayout);
 
   // Conservatively introduce branch instructions.
   BF.fixBranches();
 
   // Update exception-handling CFG for the function.
   BF.recomputeLandingPads();
 
   return LPTrampolines;
 }
 
 SplitFunctions::BasicBlockOrderType SplitFunctions::mergeEHTrampolines(
     BinaryFunction &BF, SplitFunctions::BasicBlockOrderType &Layout,
     const SplitFunctions::TrampolineSetType &Trampolines) const {
+  DenseMap<const MCSymbol *, SmallVector<const MCSymbol *, 0>>
+      IncomingTrampolines;
+  for (const auto &Entry : Trampolines) {
+    IncomingTrampolines[Entry.getFirst().Target].emplace_back(
+        Entry.getSecond());
+  }
+
   BasicBlockOrderType MergedLayout;
   for (BinaryBasicBlock *BB : Layout) {
-    auto Iter = Trampolines.find(BB->getLabel());
-    if (Iter != Trampolines.end()) {
-      BinaryBasicBlock *LPBlock = BF.getBasicBlockForLabel(Iter->second);
+    auto Iter = IncomingTrampolines.find(BB->getLabel());
+    if (Iter != IncomingTrampolines.end()) {
+      for (const MCSymbol *const Trampoline : Iter->getSecond()) {
+        BinaryBasicBlock *LPBlock = BF.getBasicBlockForLabel(Trampoline);
         assert(LPBlock && "Could not find matching landing pad block.");
         MergedLayout.push_back(LPBlock);
       }
+    }
     MergedLayout.push_back(BB);
   }
 
   return MergedLayout;
 }
 
 } // namespace bolt
 } // namespace llvm