[BOLT] Jump table trampoline insertion pass

bolt/include/bolt/Passes/JumpTableTrampoline.h

@@ -0,0 +1,43 @@
+//===- bolt/Passes/JumpTableTrampoline.h ------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Jump table trampolines insertion pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_PASSES_JUMP_TABLE_TRAMPOLINE_H
+#define BOLT_PASSES_JUMP_TABLE_TRAMPOLINE_H
+
+#include "bolt/Passes/BinaryPasses.h"
+
+namespace llvm {
+namespace bolt {
+
+/// This pass inserts trampolines for entries in cold fragment into the hot
+/// fragment so that offsets fit into the original jump table entry size.
+class JumpTableTrampoline : public BinaryFunctionPass {
+  DenseSet<const BinaryFunction *> Modified;
+
+  /// Run a pass for \p Function
+  void optimizeFunction(BinaryFunction &Function);
+
+public:
+  explicit JumpTableTrampoline(const cl::opt<bool> &PrintPass)
+      : BinaryFunctionPass(PrintPass) {}
+
+  const char *getName() const override { return "jump-table-trampoline"; }
+  bool shouldPrint(const BinaryFunction &BF) const override {
+    return BinaryFunctionPass::shouldPrint(BF) && Modified.count(&BF);
+  }
+  Error runOnFunctions(BinaryContext &BC) override;
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif

bolt/lib/Passes/CMakeLists.txt

@@ -19,6 +19,7 @@ add_llvm_library(LLVMBOLTPasses
   Inliner.cpp
   Instrumentation.cpp
   JTFootprintReduction.cpp
+  JumpTableTrampoline.cpp
   LongJmp.cpp
   LoopInversionPass.cpp
   LivenessAnalysis.cpp

bolt/lib/Passes/JumpTableTrampoline.cpp

@@ -0,0 +1,69 @@
+//===- bolt/Passes/JumpTableTrampoline.cpp --------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements JumpTableTrampoline class.
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Passes/JumpTableTrampoline.h"
+
+#define DEBUG_TYPE "JTT"
+
+namespace llvm {
+namespace bolt {
+
+void JumpTableTrampoline::optimizeFunction(BinaryFunction &Function) {
+  BinaryContext &BC = Function.getBinaryContext();
+  std::vector<std::unique_ptr<BinaryBasicBlock>> NewBBs;
+  InstructionListType Seq;
+  for (JumpTable *const &JT : llvm::make_second_range(Function.jumpTables())) {
+    for (MCSymbol *&Entry : JT->Entries) {
+      BinaryBasicBlock *Target = Function.getBasicBlockForLabel(Entry);
+      assert(Target && "Jump table entry must have corresponding basic block");
+      if (!Target->isCold())
+        continue;
+
+      std::unique_ptr<BinaryBasicBlock> NewBB =
+          Function.createBasicBlock(BC.Ctx->createNamedTempSymbol("JTT"));
+      NewBBs.emplace_back(std::move(NewBB));
+
+      // Copy attributes from the original destination.
+      BinaryBasicBlock &NewBBRef = *NewBBs.back();
+      NewBBRef.setOffset(Target->getOffset());
+
+      // Populate the new basic block.
+      BC.MIB->createLongJmp(Seq, Entry, BC.Ctx.get());
+      NewBBRef.addInstructions(Seq);
+      Seq.clear();
+
+      // Redirect the jump table entry.
+      Entry = NewBBRef.getLabel();
+    }
+  }
+
+  if (NewBBs.empty())
+    return;
+
+  Modified.insert(&Function);
+  BinaryBasicBlock *LastHotBB = Function.getLayout().getMainFragment().back();
+  assert(LastHotBB && "split function must have at least one hot block");
+  Function.insertBasicBlocks(LastHotBB, std::move(NewBBs));
+}
+
+Error JumpTableTrampoline::runOnFunctions(BinaryContext &BC) {
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions()))
+    if (BF.isSplit() && BF.hasJumpTables())
+      optimizeFunction(BF);
+
+  BC.outs() << "BOLT-INFO: inserted jump table trampolines into "
+            << Modified.size() << " functions.\n";
+  return Error::success();
+}
+
+} // namespace bolt
+} // namespace llvm

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -21,6 +21,7 @@
 #include "bolt/Passes/Inliner.h"
 #include "bolt/Passes/Instrumentation.h"
 #include "bolt/Passes/JTFootprintReduction.h"
+#include "bolt/Passes/JumpTableTrampoline.h"
 #include "bolt/Passes/LongJmp.h"
 #include "bolt/Passes/LoopInversionPass.h"
 #include "bolt/Passes/MCF.h"
@@ -74,6 +75,12 @@ static cl::opt<bool> JTFootprintReductionFlag(
              "instructions at jump sites"),
     cl::cat(BoltOptCategory));
 
+static cl::opt<bool> JumpTableTrampoline(
+    "jump-table-trampoline",
+    cl::desc("insert cold target trampolines into the hot section to preserve "
+             "the jump table entry size"),
+    cl::cat(BoltOptCategory));
+
 cl::opt<bool>
     KeepNops("keep-nops",
              cl::desc("keep no-op instructions. By default they are removed."),
@@ -126,6 +133,11 @@ static cl::opt<bool> PrintJTFootprintReduction(
     cl::desc("print function after jt-footprint-reduction pass"), cl::Hidden,
     cl::cat(BoltOptCategory));
 
+static cl::opt<bool> PrintJumpTableTrampoline(
+    "print-after-jump-table-trampoline",
+    cl::desc("print function after jump-table-trampoline pass"), cl::Hidden,
+    cl::cat(BoltOptCategory));
+
 static cl::opt<bool>
     PrintAdrRelaxation("print-adr-relaxation",
                        cl::desc("print functions after ADR Relaxation pass"),
@@ -501,6 +513,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(std::make_unique<AssignSections>());
 
   if (BC.isAArch64()) {
+    Manager.registerPass(
+        std::make_unique<JumpTableTrampoline>(PrintJumpTableTrampoline));
+
     Manager.registerPass(
         std::make_unique<ADRRelaxationPass>(PrintAdrRelaxation));