diff --git a/bolt/include/bolt/Core/BinaryBasicBlock.h b/bolt/include/bolt/Core/BinaryBasicBlock.h --- a/bolt/include/bolt/Core/BinaryBasicBlock.h +++ b/bolt/include/bolt/Core/BinaryBasicBlock.h @@ -31,6 +31,7 @@ namespace bolt { class BinaryFunction; +class JumpTable; class BinaryBasicBlock { public: @@ -623,6 +624,10 @@ /// remove the conditional successor and branch instruction. void removeDuplicateConditionalSuccessor(MCInst *CondBranch); + /// Update successors of the basic block based on the jump table instruction. + /// The block must end with a jump table instruction. + void updateJumpTableSuccessors(); + /// Test if BB is a predecessor of this block. bool isPredecessor(const BinaryBasicBlock *BB) const { auto Itr = std::find(Predecessors.begin(), Predecessors.end(), BB); @@ -909,7 +914,12 @@ return Index; } - bool hasJumpTable() const; + /// Return jump table if the block contains a jump table instruction or + /// nullptr otherwise. + const JumpTable *getJumpTable() const; + + /// Check if the block has a jump table instruction. + bool hasJumpTable() const { return getJumpTable() != nullptr; } private: void adjustNumPseudos(const MCInst &Inst, int Sign); diff --git a/bolt/lib/Core/BinaryBasicBlock.cpp b/bolt/lib/Core/BinaryBasicBlock.cpp --- a/bolt/lib/Core/BinaryBasicBlock.cpp +++ b/bolt/lib/Core/BinaryBasicBlock.cpp @@ -39,10 +39,10 @@ return getParent()->hasInstructions(); } -bool BinaryBasicBlock::hasJumpTable() const { +const JumpTable *BinaryBasicBlock::getJumpTable() const { const MCInst *Inst = getLastNonPseudoInstr(); const JumpTable *JT = Inst ? Function->getJumpTable(*Inst) : nullptr; - return (JT != nullptr); + return JT; } void BinaryBasicBlock::adjustNumPseudos(const MCInst &Inst, int Sign) { @@ -351,6 +351,36 @@ BranchInfo.push_back({Count, 0}); } +void BinaryBasicBlock::updateJumpTableSuccessors() { + const JumpTable *JT = getJumpTable(); + assert(JT && "Expected jump table instruction."); + + // Clear existing successors. + removeAllSuccessors(); + + // Generate the list of successors in deterministic order without duplicates. + SmallVector SuccessorBBs; + for (const MCSymbol *Label : JT->Entries) { + BinaryBasicBlock *BB = getFunction()->getBasicBlockForLabel(Label); + // Ignore __builtin_unreachable() + if (!BB) { + assert(Label == getFunction()->getFunctionEndLabel() && + "JT label should match a block or end of function."); + continue; + } + SuccessorBBs.emplace_back(BB); + } + llvm::sort(SuccessorBBs, + [](const BinaryBasicBlock *BB1, const BinaryBasicBlock *BB2) { + return BB1->getInputOffset() < BB2->getInputOffset(); + }); + SuccessorBBs.erase(std::unique(SuccessorBBs.begin(), SuccessorBBs.end()), + SuccessorBBs.end()); + + for (BinaryBasicBlock *BB : SuccessorBBs) + addSuccessor(BB); +} + void BinaryBasicBlock::adjustExecutionCount(double Ratio) { auto adjustedCount = [&](uint64_t Count) -> uint64_t { double NewCount = Count * Ratio; diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp --- a/bolt/lib/Core/BinaryFunction.cpp +++ b/bolt/lib/Core/BinaryFunction.cpp @@ -1908,21 +1908,7 @@ BC.MIB->setJumpTable(*LastIndirectJump, LastJT, LastJTIndexReg, AllocId); HasUnknownControlFlow = false; - // re-populate successors based on the jump table. - std::set JTLabels; - LastIndirectJumpBB->removeAllSuccessors(); - const JumpTable *JT = getJumpTableContainingAddress(LastJT); - for (const MCSymbol *Label : JT->Entries) - JTLabels.emplace(Label); - for (const MCSymbol *Label : JTLabels) { - BinaryBasicBlock *BB = getBasicBlockForLabel(Label); - // Ignore __builtin_unreachable() - if (!BB) { - assert(Label == getFunctionEndLabel() && "if no BB found, must be end"); - continue; - } - LastIndirectJumpBB->addSuccessor(BB); - } + LastIndirectJumpBB->updateJumpTableSuccessors(); } if (HasFixedIndirectBranch) diff --git a/bolt/test/X86/Inputs/jump-table-pic.s b/bolt/test/X86/Inputs/jump-table-pic.s new file mode 100644 --- /dev/null +++ b/bolt/test/X86/Inputs/jump-table-pic.s @@ -0,0 +1,52 @@ +# Test case with a simple PIC-style jump table where the register containing +# the jump table address is defined outside of the basic block containing the +# jump on register. +# +# One of the destinations of the jump table points past the end of the function +# similar to the code generated for __builtin_unreachable(). + + .text + .globl main + .type main, @function +main: + .cfi_startproc + pushq %rbx + .cfi_def_cfa_offset 16 + .cfi_offset 3, -16 + leaq .LJUMPTABLE(%rip), %rcx +.L13: + cmpl $3, %edi + ja .L2 + + movslq (%rcx,%rdi,4), %rax + addq %rcx, %rax + jmp *%rax + +.L12: + movq %rdi, %rax + popq %rbx + .cfi_remember_state + .cfi_def_cfa_offset 8 + ret + +.L5: + .cfi_restore_state + addq $9, %rdi + jmp .L12 +.L6: + addq $8, %rdi + jmp .L12 +.L2: + addq $2, %rdi + jmp .L12 +.LUNREACHABLE: + .cfi_endproc + .size main, .-main + + .section .rodata + .align 4 +.LJUMPTABLE: + .long .L2-.LJUMPTABLE + .long .L6-.LJUMPTABLE + .long .L5-.LJUMPTABLE + .long .LUNREACHABLE-.LJUMPTABLE diff --git a/bolt/test/X86/jump-table-pic-order.test b/bolt/test/X86/jump-table-pic-order.test new file mode 100644 --- /dev/null +++ b/bolt/test/X86/jump-table-pic-order.test @@ -0,0 +1,12 @@ +# Check that successors of a basic block with jump table are generated +# in the same order as they appear in the input code. + +RUN: %clang %cflags %S/Inputs/jump-table-pic.s -o %t.exe -Wl,-q +RUN: llvm-bolt %t.exe -strict -print-cfg -print-only=main -o /dev/null \ +RUN: | FileCheck %s + +CHECK: BB Layout : {{.*, .*, .*,}} [[BB4to6:.*, .*, .*]] + +# Check that successors appear in the order matching the input layout. +CHECK: jmpq *%rax # JUMPTABLE +CHECK-NEXT: Successors: [[BB4to6]]