[BOLT-AArch64] Support reordering bzip2 no relocs

Summary:
Add functionality to support reordering bzip2 compiled to
AArch64, with function splitting but without relocations:

 * Expand the AArch64 backend to support inverting branches and
analyzing branches so BOLT reordering machinery is able to shuffle
blocks and fix branches correctly;
 * Add a new pass named LongJmp to add stubs whenever code needs to
jump to the cold area, when using function splitting, because of the
limited target encoding capability in AArch64 (as a RISC architecture).

(cherry picked from FBD5748184)

Author: rafaelauler

bolt/BinaryBasicBlock.cpp

@@ -138,6 +138,21 @@ BinaryBasicBlock *BinaryBasicBlock::getSuccessor(const MCSymbol *Label) const {
   return nullptr;
 }
 
+BinaryBasicBlock *
+BinaryBasicBlock::getSuccessor(const MCSymbol *Label,
+                               BinaryBranchInfo &BI) const {
+  auto BIIter = branch_info_begin();
+  for (BinaryBasicBlock *BB : successors()) {
+    if (BB->getLabel() == Label) {
+      BI = *BIIter;
+      return BB;
+    }
+    ++BIIter;
+  }
+
+  return nullptr;
+}
+
 BinaryBasicBlock *BinaryBasicBlock::getLandingPad(const MCSymbol *Label) const {
   for (BinaryBasicBlock *BB : landing_pads()) {
     if (BB->getLabel() == Label)

bolt/BinaryBasicBlock.h

@@ -333,6 +333,10 @@ class BinaryBasicBlock {
   /// return the successor.
   BinaryBasicBlock *getSuccessor(const MCSymbol *Label = nullptr) const;
 
+  /// Return the related branch info as well as the successor.
+  BinaryBasicBlock *getSuccessor(const MCSymbol *Label,
+                                 BinaryBranchInfo &BI) const;
+
   /// If the basic block ends with a conditional branch (possibly followed by
   /// an unconditional branch) and thus has 2 successors, return a successor
   /// corresponding to a jump condition which could be true or false.
@@ -637,6 +641,12 @@ class BinaryBasicBlock {
     return Instructions.erase(II);
   }
 
+  /// Erase all instructions
+  void clear() {
+    Instructions.clear();
+    NumPseudos = 0;
+  }
+
   /// Retrieve iterator for \p Inst or return end iterator if instruction is not
   /// from this basic block.
   decltype(Instructions)::iterator findInstruction(const MCInst *Inst) {

bolt/BinaryContext.cpp

@@ -28,6 +28,7 @@ namespace opts {
 extern cl::OptionCategory BoltCategory;
 
 extern cl::opt<bool> Relocs;
+extern cl::opt<BinaryFunction::ReorderType> ReorderFunctions;
 
 static cl::opt<bool>
 PrintDebugInfo("print-debug-info",
@@ -186,6 +187,28 @@ unsigned BinaryContext::addDebugFilenameToUnit(const uint32_t DestCUID,
   return Ctx->getDwarfFile(Dir, FileNames[FileIndex - 1].Name, 0, DestCUID);
 }
 
+std::vector<BinaryFunction *> BinaryContext::getSortedFunctions(
+    std::map<uint64_t, BinaryFunction> &BinaryFunctions) {
+  std::vector<BinaryFunction *> SortedFunctions(BinaryFunctions.size());
+  std::transform(BinaryFunctions.begin(), BinaryFunctions.end(),
+                 SortedFunctions.begin(),
+                 [](std::pair<const uint64_t, BinaryFunction> &BFI) {
+                   return &BFI.second;
+                 });
+
+  if (opts::ReorderFunctions != BinaryFunction::RT_NONE) {
+    std::stable_sort(SortedFunctions.begin(), SortedFunctions.end(),
+                     [](const BinaryFunction *A, const BinaryFunction *B) {
+                       if (A->hasValidIndex() && B->hasValidIndex()) {
+                         return A->getIndex() < B->getIndex();
+                       } else {
+                         return A->hasValidIndex();
+                       }
+                     });
+  }
+  return SortedFunctions;
+}
+
 void BinaryContext::preprocessDebugInfo(
     std::map<uint64_t, BinaryFunction> &BinaryFunctions) {
   // Populate MCContext with DWARF files.

bolt/BinaryContext.h

@@ -149,6 +149,11 @@ class BinaryContext {
   /// Number of functions with profile information
   uint64_t NumProfiledFuncs{0};
 
+  /// Track next available address for new allocatable sections. RewriteInstance
+  /// sets this prior to running BOLT passes, so layout passes are aware of the
+  /// final addresses functions will have.
+  uint64_t LayoutStartAddress{0};
+
   /// True if the binary requires immediate relocation processing.
   bool RequiresZNow{false};
 
@@ -272,6 +277,10 @@ class BinaryContext {
                                   const uint32_t SrcCUID,
                                   unsigned FileIndex);
 
+  /// Return functions in output layout order
+  static std::vector<BinaryFunction *>
+  getSortedFunctions(std::map<uint64_t, BinaryFunction> &BinaryFunctions);
+
   /// Compute the native code size for a range of instructions.
   /// Note: this can be imprecise wrt the final binary since happening prior to
   /// relaxation, as well as wrt the original binary because of opcode

bolt/BinaryFunction.cpp

@@ -3969,7 +3969,7 @@ DynoStats BinaryFunction::getDynoStats() const {
           BC.MIA->getAnnotationWithDefault<uint64_t>(Instr, "CTCTakenCount");
       }
       Stats[DynoStats::FUNCTION_CALLS] += CallFreq;
-      if (BC.MIA->getMemoryOperandNo(Instr) != -1) {
+      if (BC.MIA->isIndirectCall(Instr)) {
         Stats[DynoStats::INDIRECT_CALLS] += CallFreq;
       } else if (const auto *CallSymbol = BC.MIA->getTargetSymbol(Instr)) {
         const auto *BF = BC.getFunctionForSymbol(CallSymbol);

bolt/BinaryFunction.h

@@ -711,7 +711,7 @@ class BinaryFunction {
     assert(InstructionOffsets.size() == Instructions.size() &&
            "There must be one instruction at every offset.");
     Instructions.emplace_back(std::forward<MCInst>(Instruction));
-    InstructionOffsets[Offset] = Instructions.size() - 1; 
+    InstructionOffsets[Offset] = Instructions.size() - 1;
   }
 
   /// Return instruction at a given offset in the function. Valid before
@@ -1863,6 +1863,18 @@ class BinaryFunction {
     return Estimate;
   }
 
+  size_t estimateColdSize() const {
+    if (!isSplit())
+      return estimateSize();
+    size_t Estimate = 0;
+    for (const auto *BB : BasicBlocksLayout) {
+      if (BB->isCold()) {
+        Estimate += BC.computeCodeSize(BB->begin(), BB->end());
+      }
+    }
+    return Estimate;
+  }
+
   size_t estimateSize() const {
     size_t Estimate = 0;
     for (const auto *BB : BasicBlocksLayout) {

bolt/BinaryPassManager.cpp

@@ -14,6 +14,7 @@
 #include "Passes/FrameOptimizer.h"
 #include "Passes/IndirectCallPromotion.h"
 #include "Passes/Inliner.h"
+#include "Passes/LongJmp.h"
 #include "Passes/PLTCall.h"
 #include "Passes/ReorderFunctions.h"
 #include "Passes/StokeInfo.h"
@@ -105,6 +106,13 @@ PrintFinalized("print-finalized",
   cl::Hidden,
   cl::cat(BoltOptCategory));
 
+static cl::opt<bool>
+PrintLongJmp("print-longjmp",
+  cl::desc("print functions after longjmp pass"),
+  cl::ZeroOrMore,
+  cl::Hidden,
+  cl::cat(BoltOptCategory));
+
 static cl::opt<bool>
 PrintICF("print-icf",
   cl::desc("print functions after ICF optimization"),
@@ -396,6 +404,12 @@ void BinaryFunctionPassManager::runAllPasses(
 
   Manager.registerPass(llvm::make_unique<AllocCombinerPass>(PrintFOP));
 
+  // Thighten branches according to offset differences between branch and
+  // targets. No extra instructions after this pass, otherwise we may have
+  // relocations out of range and crash during linking.
+  if (BC.TheTriple->getArch() == llvm::Triple::aarch64)
+    Manager.registerPass(llvm::make_unique<LongJmpPass>(PrintLongJmp));
+
   // This pass turns tail calls into jumps which makes them invisible to
   // function reordering. It's unsafe to use any CFG or instruction analysis
   // after this point.

bolt/Passes/BinaryPasses.cpp

@@ -233,7 +233,7 @@ void OptimizeBodylessFunctions::optimizeCalls(BinaryFunction &BF,
                    << ": replacing call to " << OriginalTarget->getName()
                    << " by call to " << Target->getName()
                    << " while folding " << CallSites << " call sites\n");
-      BC.MIA->replaceCallTargetOperand(Inst, Target, BC.Ctx.get());
+      BC.MIA->replaceBranchTarget(Inst, Target, BC.Ctx.get());
 
       NumOptimizedCallSites += CallSites;
       if (BB->hasProfile()) {

bolt/Passes/CMakeLists.txt

@@ -13,6 +13,7 @@ add_llvm_library(LLVMBOLTPasses
   IndirectCallPromotion.cpp
   Inliner.cpp
   LivenessAnalysis.cpp
+  LongJmp.cpp
   MCF.cpp
   PettisAndHansen.cpp
   PLTCall.cpp