[BOLT] Optimize the three way branch

Summary:
Three way branches commonly appear
in HHVM. They have one test and then two jumps.  The
jump's destinations are not currently optimized.
This pass attempts to optimize which is the first branch.

(cherry picked from FBD30460441)

bolt/src/BinaryPassManager.cpp

@@ -29,6 +29,7 @@
 #include "Passes/SplitFunctions.h"
 #include "Passes/StokeInfo.h"
 #include "Passes/TailDuplication.h"
+#include "Passes/ThreeWayBranch.h"
 #include "Passes/ValidateInternalCalls.h"
 #include "Passes/VeneerElimination.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -83,6 +84,11 @@ static cl::opt<bool> TailDuplicationFlag(
     cl::desc("duplicate unconditional branches that cross a cache line"),
     cl::ZeroOrMore, cl::ReallyHidden, cl::cat(BoltOptCategory));
 
+static cl::opt<bool> ThreeWayBranchFlag("three-way-branch",
+                                        cl::desc("reorder three way branches"),
+                                        cl::ZeroOrMore, cl::ReallyHidden,
+                                        cl::cat(BoltOptCategory));
+
 static cl::opt<bool>
 PrintJTFootprintReduction("print-after-jt-footprint-reduction",
   cl::desc("print function after jt-footprint-reduction pass"),
@@ -446,6 +452,9 @@ void BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
 
   Manager.registerPass(std::make_unique<PLTCall>(PrintPLT));
 
+  Manager.registerPass(std::make_unique<ThreeWayBranch>(),
+                       opts::ThreeWayBranchFlag);
+
   Manager.registerPass(std::make_unique<ReorderBasicBlocks>(PrintReordered));
 
   Manager.registerPass(

bolt/src/MCPlusBuilder.h

@@ -583,6 +583,11 @@ class MCPlusBuilder {
     return false;
   }
 
+  virtual bool isPacked(const MCInst &Inst) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
   /// If non-zero, this is used to fill the executable space with instructions
   /// that will trap. Defaults to 0.
   virtual unsigned getTrapFillValue() const { return 0; }
@@ -1572,6 +1577,27 @@ class MCPlusBuilder {
     return false;
   }
 
+  virtual bool replaceBranchCondition(MCInst &Inst, const MCSymbol *TBB,
+                                      MCContext *Ctx, unsigned CC) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  virtual unsigned getInvertedCondCode(unsigned CC) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  virtual unsigned getCondCodesLogicalOr(unsigned CC1, unsigned CC2) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  virtual bool isValidCondCode(unsigned CC) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
   /// Return the conditional code used in a conditional jump instruction.
   /// Returns invalid code if not conditional jump.
   virtual unsigned getCondCode(const MCInst &Inst) const {

bolt/src/Passes/CMakeLists.txt

@@ -38,6 +38,7 @@ add_llvm_library(LLVMBOLTPasses
   StackReachingUses.cpp
   StokeInfo.cpp
   TailDuplication.cpp
+  ThreeWayBranch.cpp
   ValidateInternalCalls.cpp
   VeneerElimination.cpp
   RetpolineInsertion.cpp

bolt/src/Passes/ThreeWayBranch.cpp

@@ -0,0 +1,168 @@
+//===--------- Passes/ThreeWayBranch.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
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "ThreeWayBranch.h"
+
+#include <numeric>
+
+using namespace llvm;
+
+namespace llvm {
+namespace bolt {
+
+bool ThreeWayBranch::shouldRunOnFunction(BinaryFunction &Function) {
+  BinaryContext &BC = Function.getBinaryContext();
+  BinaryFunction::BasicBlockOrderType BlockLayout = Function.getLayout();
+  for (BinaryBasicBlock *BB : BlockLayout) {
+    for (MCInst &Inst : *BB) {
+      if (BC.MIB->isPacked(Inst))
+        return false;
+    }
+  }
+  return true;
+}
+
+void ThreeWayBranch::runOnFunction(BinaryFunction &Function) {
+  BinaryContext &BC = Function.getBinaryContext();
+  MCContext *Ctx = BC.Ctx.get();
+  // New blocks will be added and layout will change,
+  // so make a copy here to iterate over the original layout
+  BinaryFunction::BasicBlockOrderType BlockLayout = Function.getLayout();
+  for (BinaryBasicBlock *BB : BlockLayout) {
+    // The block must be hot
+    if (BB->getExecutionCount() == 0 ||
+        BB->getExecutionCount() == BinaryBasicBlock::COUNT_NO_PROFILE)
+      continue;
+    // with two successors
+    if (BB->succ_size() != 2)
+      continue;
+    // no jump table
+    if (BB->hasJumpTable())
+      continue;
+
+    BinaryBasicBlock *FalseSucc = BB->getConditionalSuccessor(false);
+    BinaryBasicBlock *TrueSucc = BB->getConditionalSuccessor(true);
+
+    // One of BB's successors must have only one instruction that is a
+    // conditional jump
+    if ((FalseSucc->succ_size() != 2 || FalseSucc->size() != 1) &&
+        (TrueSucc->succ_size() != 2 || TrueSucc->size() != 1))
+      continue;
+
+    // SecondBranch has the second conditional jump
+    BinaryBasicBlock *SecondBranch = FalseSucc;
+    BinaryBasicBlock *FirstEndpoint = TrueSucc;
+    if (FalseSucc->succ_size() != 2) {
+      SecondBranch = TrueSucc;
+      FirstEndpoint = FalseSucc;
+    }
+
+    BinaryBasicBlock *SecondEndpoint =
+        SecondBranch->getConditionalSuccessor(false);
+    BinaryBasicBlock *ThirdEndpoint =
+        SecondBranch->getConditionalSuccessor(true);
+
+    // Make sure we can modify the jump in SecondBranch without disturbing any
+    // other paths
+    if (SecondBranch->pred_size() != 1)
+      continue;
+
+    // Get Jump Instructions
+    MCInst *FirstJump = BB->getLastNonPseudoInstr();
+    MCInst *SecondJump = SecondBranch->getLastNonPseudoInstr();
+
+    // Get condition codes
+    unsigned FirstCC = BC.MIB->getCondCode(*FirstJump);
+    if (SecondBranch != FalseSucc)
+      FirstCC = BC.MIB->getInvertedCondCode(FirstCC);
+    // ThirdCC = ThirdCond && !FirstCC = !(!ThirdCond ||
+    // !(!FirstCC)) = !(!ThirdCond || FirstCC)
+    unsigned ThirdCC =
+        BC.MIB->getInvertedCondCode(BC.MIB->getCondCodesLogicalOr(
+            BC.MIB->getInvertedCondCode(BC.MIB->getCondCode(*SecondJump)),
+            FirstCC));
+    // SecondCC = !ThirdCond && !FirstCC = !(!(!ThirdCond) ||
+    // !(!FirstCC)) = !(ThirdCond || FirstCC)
+    unsigned SecondCC =
+        BC.MIB->getInvertedCondCode(BC.MIB->getCondCodesLogicalOr(
+            BC.MIB->getCondCode(*SecondJump), FirstCC));
+
+    if (!BC.MIB->isValidCondCode(FirstCC) ||
+        !BC.MIB->isValidCondCode(ThirdCC) || !BC.MIB->isValidCondCode(SecondCC))
+      continue;
+
+    std::vector<std::pair<BinaryBasicBlock *, unsigned>> Blocks;
+    Blocks.push_back(std::make_pair(FirstEndpoint, FirstCC));
+    Blocks.push_back(std::make_pair(SecondEndpoint, SecondCC));
+    Blocks.push_back(std::make_pair(ThirdEndpoint, ThirdCC));
+
+    std::sort(Blocks.begin(), Blocks.end(),
+              [&](const std::pair<BinaryBasicBlock *, unsigned> A,
+                  const std::pair<BinaryBasicBlock *, unsigned> B) {
+                return A.first->getExecutionCount() <
+                       B.first->getExecutionCount();
+              });
+
+    uint64_t NewSecondBranchCount = Blocks[1].first->getExecutionCount() +
+                                    Blocks[0].first->getExecutionCount();
+    bool SecondBranchBigger =
+        NewSecondBranchCount > Blocks[2].first->getExecutionCount();
+
+    BB->removeAllSuccessors();
+    if (SecondBranchBigger) {
+      BB->addSuccessor(Blocks[2].first, Blocks[2].first->getExecutionCount());
+      BB->addSuccessor(SecondBranch, NewSecondBranchCount);
+    } else {
+      BB->addSuccessor(SecondBranch, NewSecondBranchCount);
+      BB->addSuccessor(Blocks[2].first, Blocks[2].first->getExecutionCount());
+    }
+
+    // Remove and add so there is no duplicate successors
+    SecondBranch->removeAllSuccessors();
+    SecondBranch->addSuccessor(Blocks[0].first,
+                               Blocks[0].first->getExecutionCount());
+    SecondBranch->addSuccessor(Blocks[1].first,
+                               Blocks[1].first->getExecutionCount());
+
+    SecondBranch->setExecutionCount(NewSecondBranchCount);
+
+    // Replace the branch condition to fallthrough for the most common block
+    if (SecondBranchBigger) {
+      BC.MIB->replaceBranchCondition(*FirstJump, Blocks[2].first->getLabel(),
+                                     Ctx, Blocks[2].second);
+    } else {
+      BC.MIB->replaceBranchCondition(
+          *FirstJump, SecondBranch->getLabel(), Ctx,
+          BC.MIB->getInvertedCondCode(Blocks[2].second));
+    }
+
+    // Replace the branch condition to fallthrough for the second most common
+    // block
+    BC.MIB->replaceBranchCondition(*SecondJump, Blocks[0].first->getLabel(),
+                                   Ctx, Blocks[0].second);
+
+    ++BranchesAltered;
+  }
+}
+
+void ThreeWayBranch::runOnFunctions(BinaryContext &BC) {
+  for (auto &It : BC.getBinaryFunctions()) {
+    BinaryFunction &Function = It.second;
+    if (!shouldRunOnFunction(Function))
+      continue;
+    runOnFunction(Function);
+  }
+
+  outs() << "BOLT-INFO: number of three way branches order changed: "
+         << BranchesAltered << "\n";
+}
+
+} // end namespace bolt
+} // end namespace llvm

bolt/src/Passes/ThreeWayBranch.h

@@ -0,0 +1,43 @@
+//===--------- Passes/ThreeWayBranch.h ------------------------------------===//
+//
+// 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 LLVM_TOOLS_LLVM_BOLT_PASSES_THREEWAYBRANCH_H
+#define LLVM_TOOLS_LLVM_BOLT_PASSES_THREEWAYBRANCH_H
+
+#include "BinaryPasses.h"
+
+namespace llvm {
+namespace bolt {
+
+/// Pass for optimizing a three way branch namely a single comparison and 2
+/// conditional jumps by reordering blocks, replacing successors, and replacing
+/// jump conditions and destinations
+class ThreeWayBranch : public BinaryFunctionPass {
+  /// Record how many 3 way branches were adjusted
+  uint64_t BranchesAltered = 0;
+
+  /// Returns true if this pass should run on Function
+  bool shouldRunOnFunction(BinaryFunction &Function);
+
+  /// Runs pass on Function
+  void runOnFunction(BinaryFunction &Function);
+
+public:
+  explicit ThreeWayBranch() : BinaryFunctionPass(false) {}
+
+  const char *getName() const override { return "three way branch"; }
+
+  void runOnFunctions(BinaryContext &BC) override;
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif