Split the indirect jump critical edges.

[hudsontek]

Summary

This diff is for processing the critical edges generated from the switch statements. When there are more than 2 cases in the switch statement, a critical edge will form between the entry block and the epilog block. This critical edge is not taken care of by the machine-sink pass. Therefore we created this diff to make the split happen

[MatzeB]

Looks like llvm/test/CodeGen/X86/switch-phi-const.ll may need some adjustments.

[MatzeB]

llvm/lib/CodeGen/MachineBasicBlock.cpp

1067

I thought we had the magic number 3 abstracted away already. This needs to be abstracted in the findIndirectJumpTable function as it is X86 specific!

1275–1278

Why was the comment removed?

llvm/lib/Target/X86/X86InstrInfo.h

332

I thought this would return just the jump table for the particular instruction, but MachineJumpTableInfo is just the generic interface for all jumptables it seems. So we probably rather something like llvm::Optional<unsigned> to return the operand index with the JTI machine operand...

[MatzeB]

I updated this diff to to no longer hardcode machine specific numbers in MachineBasicBlock::SplitCriticalEdge and update the X86 to support PIC and non-PIC jump table styles in D140975

llvm/include/llvm/CodeGen/MachineBasicBlock.h

public:
   MachineBasicBlock *
   SplitCriticalEdge(MachineBasicBlock *Succ, Pass &P,
                     std::vector<SparseBitVector<>> *LiveInSets = nullptr);
 
   /// Check if the edge between this block and the given successor \p
   /// Succ, can be split. If this returns true a subsequent call to
   /// SplitCriticalEdge is guaranteed to return a valid basic block if
   /// no changes occurred in the meantime.
-  bool canSplitCriticalEdge(const MachineBasicBlock *Succ) const;
+  bool canSplitCriticalEdge(const MachineBasicBlock *Succ);
 
   void pop_front() { Insts.pop_front(); }
   void pop_back() { Insts.pop_back(); }
   void push_back(MachineInstr *MI) { Insts.push_back(MI); }
 
   /// Insert MI into the instruction list before I, possibly inside a bundle.
   ///
   /// If the insertion point is inside a bundle, MI will be added to the bundle,

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   ///
   /// @param MO The MachineOperand in question. MO.isReg() should always
   /// be true.
   /// @return Whether this operand is allowed to be used PC-relatively.
   virtual bool isPCRelRegisterOperandLegal(const MachineOperand &MO) const {
     return false;
   }
 
+  /// Return the pointer to the jump table, if given instr has it,
+  /// otherwise return nullptr.
+  virtual MachineJumpTableInfo *
+  getJumpTableFromIndirectJump(MachineInstr &MI) const {
+    return nullptr;
+  }
+
 protected:
   /// Target-dependent implementation for IsCopyInstr.
   /// If the specific machine instruction is a instruction that moves/copies
   /// value from one register to another register return destination and source
   /// registers as machine operands.
   virtual Optional<DestSourcePair>
   isCopyInstrImpl(const MachineInstr &MI) const {
     return None;

llvm/lib/CodeGen/MachineBasicBlock.cpp

 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/LiveVariables.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h" // facebook T46037538
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include <algorithm>
+#include <memory>
 using namespace llvm;
 
 #define DEBUG_TYPE "codegen"
 
 static cl::opt<bool> PrintSlotIndexes(
     "print-slotindexes",
     cl::desc("When printing machine IR, annotate instructions and blocks with "
              "SlotIndexes when available"),

 void MachineBasicBlock::moveBefore(MachineBasicBlock *NewAfter) {
   getParent()->splice(NewAfter->getIterator(), getIterator());
 }
 
 void MachineBasicBlock::moveAfter(MachineBasicBlock *NewBefore) {
   getParent()->splice(++NewBefore->getIterator(), getIterator());
 }
 
+static MachineJumpTableInfo *findIndirectJumpTable(MachineBasicBlock &MBB) {
+  MachineBasicBlock::iterator TerminatorI = MBB.getFirstTerminator();
+  if (TerminatorI == MBB.end())
+    return nullptr;
+  const TargetInstrInfo *TII = MBB.getParent()->getSubtarget().getInstrInfo();
+  return TII->getJumpTableFromIndirectJump(*TerminatorI);
+}
+
 void MachineBasicBlock::updateTerminator(
     MachineBasicBlock *PreviousLayoutSuccessor) {
   LLVM_DEBUG(dbgs() << "Updating terminators on " << printMBBReference(*this)
                     << "\n");
 
   const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo();
   // A block with no successors has no concerns with fall-through edges.
   if (this->succ_empty())
     return;
 
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
   SmallVector<MachineOperand, 4> Cond;
   DebugLoc DL = findBranchDebugLoc();
   bool B = TII->analyzeBranch(*this, TBB, FBB, Cond);
   (void) B;
-  assert(!B && "UpdateTerminators requires analyzable predecessors!");
+  assert((!B || findIndirectJumpTable(*this)) &&
+         "UpdateTerminators requires analyzable predecessors!");
   if (Cond.empty()) {
     if (TBB) {
       // The block has an unconditional branch. If its successor is now its
       // layout successor, delete the branch.
       if (isLayoutSuccessor(TBB))
         TII->removeBranch(*this);
     } else {
       // The block has an unconditional fallthrough, or the end of the block is
MachineBasicBlock *MachineBasicBlock::splitAt(MachineInstr &MI,
 
   return SplitBB;
 }
 
 MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
     MachineBasicBlock *Succ, Pass &P,
     std::vector<SparseBitVector<>> *LiveInSets) {
   if (!canSplitCriticalEdge(Succ))
-    return nullptr;
+    return nullptr; // this is not a critical edge
 
   MachineFunction *MF = getParent();
+  const TargetInstrInfo &TII = *MF->getSubtarget().getInstrInfo();
+  DebugLoc DL;
   MachineBasicBlock *PrevFallthrough = getNextNode();
-  DebugLoc DL;  // FIXME: this is nowhere
-
   MachineBasicBlock *NMBB = MF->CreateMachineBasicBlock();
+
+  if (MachineJumpTableInfo *JTI = findIndirectJumpTable(*this)) {
+    // Check if there is a indirect jump with jump table
+    MachineInstr &MI = *getFirstTerminator();
+    const MachineOperand &MO = MI.getOperand(3);
+    JTI->ReplaceMBBInJumpTable(MO.getIndex(), Succ, NMBB);
+
+    // insert an unconditional branch at the end of the new split BB
+    TII.insertUnconditionalBranch(*NMBB, Succ, DL);
+  }
+
   MF->insert(std::next(MachineFunction::iterator(this)), NMBB);
   LLVM_DEBUG(dbgs() << "Splitting critical edge: " << printMBBReference(*this)
                     << " -- " << printMBBReference(*NMBB) << " -- "
                     << printMBBReference(*Succ) << '\n');
 
   LiveIntervals *LIS = P.getAnalysisIfAvailable<LiveIntervals>();
   SlotIndexes *Indexes = P.getAnalysisIfAvailable<SlotIndexes>();
   if (LIS)
if (MachineLoop *TIL = MLI->getLoopFor(this)) {
             P->addBasicBlockToLoop(NMBB, MLI->getBase());
         }
       }
     }
 
   return NMBB;
 }
 
-bool MachineBasicBlock::canSplitCriticalEdge(
-    const MachineBasicBlock *Succ) const {
-  // Splitting the critical edge to a landing pad block is non-trivial. Don't do
-  // it in this generic function.
+bool MachineBasicBlock::canSplitCriticalEdge(const MachineBasicBlock *Succ) {
   if (Succ->isEHPad())
     return false;
 
   // Splitting the critical edge to a callbr's indirect block isn't advised.
   // Don't do it in this generic function.
   if (Succ->isInlineAsmBrIndirectTarget())
     return false;
 
-  const MachineFunction *MF = getParent();
+  const MachineFunction &MF = *Succ->getParent();
   // Performance might be harmed on HW that implements branching using exec mask
   // where both sides of the branches are always executed.
-  if (MF->getTarget().requiresStructuredCFG())
+  if (MF.getTarget().requiresStructuredCFG())
     return false;
 
+  // Check whether we have indirect jump with a jumptable that we can rewrite.
+  if (findIndirectJumpTable(*this) != nullptr)
+    return true;
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   // We may need to update this's terminator, but we can't do that if
-  // analyzeBranch fails. If this uses a jump table, we won't touch it.
-  const TargetInstrInfo *TII = MF->getSubtarget().getInstrInfo();
+  // analyzeBranch fails.
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
   SmallVector<MachineOperand, 4> Cond;
+
   // AnalyzeBanch should modify this, since we did not allow modification.
   if (TII->analyzeBranch(*const_cast<MachineBasicBlock *>(this), TBB, FBB, Cond,
                          /*AllowModify*/ false))
     return false;
 
   // Avoid bugpoint weirdness: A block may end with a conditional branch but
   // jumps to the same MBB is either case. We have duplicate CFG edges in that
   // case that we can't handle. Since this never happens in properly optimized

llvm/lib/Target/X86/X86InstrInfo.h

void replaceBranchWithTailCall(MachineBasicBlock &MBB,
                                  SmallVectorImpl<MachineOperand> &Cond,
                                  const MachineInstr &TailCall) const override;
 
   bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                      MachineBasicBlock *&FBB,
                      SmallVectorImpl<MachineOperand> &Cond,
                      bool AllowModify) const override;
 
+  /// Return the pointer to the jump table, if given instr has it.
+  MachineJumpTableInfo *
+  getJumpTableFromIndirectJump(MachineInstr &MI) const override;
+
   Optional<ExtAddrMode>
   getAddrModeFromMemoryOp(const MachineInstr &MemI,
                           const TargetRegisterInfo *TRI) const override;
 
   bool getConstValDefinedInReg(const MachineInstr &MI, const Register Reg,
                                int64_t &ImmVal) const override;
 
   bool preservesZeroValueInReg(const MachineInstr *MI,

llvm/lib/Target/X86/X86InstrInfo.cpp

while (I != MBB.begin()) {
     // analysis.
     if (!I->isBranch())
       return true;
 
     // Handle unconditional branches.
     if (I->getOpcode() == X86::JMP_1) {
       UnCondBrIter = I;
 
-      if (!AllowModify) {
-        TBB = I->getOperand(0).getMBB();
+      // TBB is used to indicate the unconditional destination.
+      if (!AllowModify) {
         continue;
       }
 
       // If the block has any instructions after a JMP, delete them.
       MBB.erase(std::next(I), MBB.end());
 
       Cond.clear();
       FBB = nullptr;
 
       // Delete the JMP if it's equivalent to a fall-through.
       if (MBB.isLayoutSuccessor(I->getOperand(0).getMBB())) {
         TBB = nullptr;
         I->eraseFromParent();
         I = MBB.end();
         UnCondBrIter = MBB.end();
-        continue;
       }
 
-      // TBB is used to indicate the unconditional destination.
-      TBB = I->getOperand(0).getMBB();
       continue;
     }
 
     // Handle conditional branches.
     X86::CondCode BranchCode = X86::getCondFromBranch(*I);
     if (BranchCode == X86::COND_INVALID)
       return true;  // Can't handle indirect branch.
 
bool X86InstrInfo::analyzeBranch(MachineBasicBlock &MBB,
                                  MachineBasicBlock *&TBB,
                                  MachineBasicBlock *&FBB,
                                  SmallVectorImpl<MachineOperand> &Cond,
                                  bool AllowModify) const {
   SmallVector<MachineInstr *, 4> CondBranches;
   return AnalyzeBranchImpl(MBB, TBB, FBB, Cond, CondBranches, AllowModify);
 }
 
+MachineJumpTableInfo *
+X86InstrInfo::getJumpTableFromIndirectJump(MachineInstr &MI) const {
+  // Currently we only take care of JMP64m jump instrs.
+  if (MI.getOpcode() != X86::JMP64m)
+    return nullptr;
+
+  // For a JMP64m instr, the 4th operand is the jump table.
+  MachineOperand &MO = MI.getOperand(3);
+  if (!MO.isJTI())
+    return nullptr;
+
+  MachineFunction *MF = MI.getParent()->getParent();
+  return MF->getJumpTableInfo();
+}
+
 bool X86InstrInfo::analyzeBranchPredicate(MachineBasicBlock &MBB,
                                           MachineBranchPredicate &MBP,
                                           bool AllowModify) const {
   using namespace std::placeholders;
 
   SmallVector<MachineOperand, 4> Cond;
   SmallVector<MachineInstr *, 4> CondBranches;
   if (AnalyzeBranchImpl(MBB, MBP.TrueDest, MBP.FalseDest, Cond, CondBranches,

llvm/test/CodeGen/X86/switch-jmp-edge-split.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=x86_64-none-linux-gnu -run-pass=machine-sink -verify-machineinstrs  -o - %s | FileCheck %s
+
+# This test was originally generated from the following sample:
+#
+# int foo(int i) {
+#       int j = 2;
+#       switch(i) {
+#               case 0:
+#                       j+= 1;
+#                       break;
+#               case 1:
+#                       j*=3;
+#                       break;
+#               case 2:
+#                       j/=2;
+#                       break;
+#               case 3:
+#                       j-=10;
+#               default:
+#                       j+=100;
+#       }
+#
+#       return j;
+# }
+
+# The added code is to split the critical edge formed by the switch jumps.
+
+---
+name:            _Z3fooi
+tracksRegLiveness: true
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.5', '%bb.1', '%bb.2', '%bb.3' ]
+body:             |
+  ; CHECK-LABEL: name: _Z3fooi
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.7(0x19999998), %bb.1(0x66666668)
+  ; CHECK-NEXT:   liveins: $edi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr32 = COPY $edi
+  ; CHECK-NEXT:   [[SUB32ri8_:%[0-9]+]]:gr32 = SUB32ri8 [[COPY]], 3, implicit-def $eflags
+  ; CHECK-NEXT:   JCC_1 %bb.1, 6, implicit $eflags
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.7:
+  ; CHECK-NEXT:   successors: %bb.5(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32ri:%[0-9]+]]:gr32 = MOV32ri 102
+  ; CHECK-NEXT:   JMP_1 %bb.5
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1.entry:
+  ; CHECK-NEXT:   successors: %bb.8(0x20000000), %bb.2(0x20000000), %bb.3(0x20000000), %bb.4(0x20000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32rr:%[0-9]+]]:gr32 = MOV32rr [[COPY]]
+  ; CHECK-NEXT:   [[SUBREG_TO_REG:%[0-9]+]]:gr64_nosp = SUBREG_TO_REG 0, [[MOV32rr]], %subreg.sub_32bit
+  ; CHECK-NEXT:   JMP64m $noreg, 8, [[SUBREG_TO_REG]], %jump-table.0, $noreg :: (load (s64) from jump-table)
+  ; CHECK-NEXT:   JMP_1 %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.8:
+  ; CHECK-NEXT:   successors: %bb.6(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32ri1:%[0-9]+]]:gr32 = MOV32ri 3
+  ; CHECK-NEXT:   JMP_1 %bb.6
+  ; CHECK-NEXT:   JMP_1 %bb.6
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.6(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32ri2:%[0-9]+]]:gr32 = MOV32ri 6
+  ; CHECK-NEXT:   JMP_1 %bb.6
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3:
+  ; CHECK-NEXT:   successors: %bb.6(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32ri3:%[0-9]+]]:gr32 = MOV32ri 1
+  ; CHECK-NEXT:   JMP_1 %bb.6
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.4:
+  ; CHECK-NEXT:   successors: %bb.5(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[MOV32ri4:%[0-9]+]]:gr32 = MOV32ri 92
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.5:
+  ; CHECK-NEXT:   successors: %bb.6(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[PHI:%[0-9]+]]:gr32 = PHI [[MOV32ri]], %bb.7, [[MOV32ri]], %bb.7, [[MOV32ri4]], %bb.4
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.6:
+  ; CHECK-NEXT:   [[PHI1:%[0-9]+]]:gr32 = PHI [[MOV32ri1]], %bb.8, [[MOV32ri3]], %bb.3, [[MOV32ri2]], %bb.2, [[PHI]], %bb.5
+  ; CHECK-NEXT:   $eax = COPY [[PHI1]]
+  ; CHECK-NEXT:   RET 0, $eax
+  bb.0.entry:
+    successors: %bb.4(0x19999998), %bb.6(0x66666668)
+    liveins: $edi
+
+    %2:gr32 = COPY $edi
+    %4:gr32 = MOV32ri 3
+    %3:gr32 = MOV32ri 102
+    %6:gr32 = MOV32rr %2
+    %5:gr64_nosp = SUBREG_TO_REG 0, killed %6, %subreg.sub_32bit
+    %7:gr32 = SUB32ri8 %2, 3, implicit-def $eflags
+    JCC_1 %bb.4, 7, implicit $eflags
+
+  bb.6.entry:
+    successors: %bb.5(0x20000000), %bb.1(0x20000000), %bb.2(0x20000000), %bb.3(0x20000000)
+
+    JMP64m $noreg, 8, %5, %jump-table.0, $noreg :: (load (s64) from jump-table)
+
+  bb.1:
+    successors: %bb.5(0x80000000)
+
+    %10:gr32 = MOV32ri 6
+    JMP_1 %bb.5
+
+  bb.2:
+    successors: %bb.5(0x80000000)
+
+    %9:gr32 = MOV32ri 1
+    JMP_1 %bb.5
+
+  bb.3:
+    successors: %bb.4(0x80000000)
+
+    %8:gr32 = MOV32ri 92
+
+  bb.4:
+    successors: %bb.5(0x80000000)
+
+    %0:gr32 = PHI %3, %bb.0, %3, %bb.0, %8, %bb.3
+
+  bb.5:
+    %1:gr32 = PHI %4, %bb.6, %9, %bb.2, %10, %bb.1, %0, %bb.4
+    $eax = COPY %1
+    RET 0, $eax
+
+...