Index: llvm/trunk/lib/Target/X86/X86IndirectBranchTracking.cpp =================================================================== --- llvm/trunk/lib/Target/X86/X86IndirectBranchTracking.cpp +++ llvm/trunk/lib/Target/X86/X86IndirectBranchTracking.cpp @@ -57,7 +57,7 @@ /// The function will not add it if already exists. /// It will add ENDBR32 or ENDBR64 opcode, depending on the target. /// \returns true if the ENDBR was added and false otherwise. - bool addENDBR(MachineBasicBlock &MBB) const; + bool addENDBR(MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const; }; } // end anonymous namespace @@ -68,20 +68,31 @@ return new X86IndirectBranchTrackingPass(); } -bool X86IndirectBranchTrackingPass::addENDBR(MachineBasicBlock &MBB) const { +bool X86IndirectBranchTrackingPass::addENDBR( + MachineBasicBlock &MBB, MachineBasicBlock::iterator I) const { assert(TII && "Target instruction info was not initialized"); assert((X86::ENDBR64 == EndbrOpcode || X86::ENDBR32 == EndbrOpcode) && "Unexpected Endbr opcode"); - auto MI = MBB.begin(); - // If the MBB is empty or the first instruction is not ENDBR, - // add the ENDBR instruction to the beginning of the MBB. - if (MI == MBB.end() || EndbrOpcode != MI->getOpcode()) { - BuildMI(MBB, MI, MBB.findDebugLoc(MI), TII->get(EndbrOpcode)); - NumEndBranchAdded++; + // If the MBB/I is empty or the current instruction is not ENDBR, + // insert ENDBR instruction to the location of I. + if (I == MBB.end() || I->getOpcode() != EndbrOpcode) { + BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(EndbrOpcode)); + ++NumEndBranchAdded; return true; } + return false; +} +bool IsCallReturnTwice(llvm::MachineOperand &MOp) { + if (!MOp.isGlobal()) + return false; + auto *CalleeFn = dyn_cast(MOp.getGlobal()); + if (!CalleeFn) + return false; + AttributeList Attrs = CalleeFn->getAttributes(); + if (Attrs.hasAttribute(AttributeList::FunctionIndex, Attribute::ReturnsTwice)) + return true; return false; } @@ -107,14 +118,21 @@ !MF.getFunction().hasLocalLinkage()) && !MF.getFunction().doesNoCfCheck()) { auto MBB = MF.begin(); - Changed |= addENDBR(*MBB); + Changed |= addENDBR(*MBB, MBB->begin()); } - for (auto &MBB : MF) + for (auto &MBB : MF) { // Find all basic blocks that their address was taken (for example // in the case of indirect jump) and add ENDBR instruction. if (MBB.hasAddressTaken()) - Changed |= addENDBR(MBB); + Changed |= addENDBR(MBB, MBB.begin()); + for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) { + if (!I->isCall()) + continue; + if (IsCallReturnTwice(I->getOperand(0))) + Changed |= addENDBR(MBB, std::next(I)); + } + } return Changed; } Index: llvm/trunk/test/CodeGen/X86/indirect-branch-tracking-r2.ll =================================================================== --- llvm/trunk/test/CodeGen/X86/indirect-branch-tracking-r2.ll +++ llvm/trunk/test/CodeGen/X86/indirect-branch-tracking-r2.ll @@ -0,0 +1,76 @@ +; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s --check-prefix=X64 +; RUN: llc -mtriple=i386-unknown-unknown < %s | FileCheck %s --check-prefix=X86 + +;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; +;; foo +;; ----- +;; Checks ENDBR insertion after return twice functions. +;; setjmp, sigsetjmp, savectx, vfork, getcontext +;; setzx is not return twice function, should not followed by endbr. +;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; + +; X64: callq setjmp +; X64-NEXT: endbr64 +; X64: callq setzx +; X64-NEXT: xorl +; X64: callq sigsetjmp +; X64-NEXT: endbr64 +; X64: callq savectx +; X64-NEXT: endbr64 +; X64: callq vfork +; X64-NEXT: endbr64 +; X64: callq getcontext +; X64-NEXT: endbr64 + +; X86: calll setjmp +; X86-NEXT: endbr32 +; X86: calll setzx +; X86-NEXT: xorl +; X86: calll sigsetjmp +; X86-NEXT: endbr32 +; X86: calll savectx +; X86-NEXT: endbr32 +; X86: calll vfork +; X86-NEXT: endbr32 +; X86: calll getcontext +; X86-NEXT: endbr32 + +; Function Attrs: noinline nounwind optnone uwtable +define dso_local void @foo() #0 { +entry: + %call = call i32 (i32, ...) bitcast (i32 (...)* @setjmp to i32 (i32, ...)*)(i32 0) #1 + %call1 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call2 = call i32 (i32, ...) bitcast (i32 (...)* @sigsetjmp to i32 (i32, ...)*)(i32 0) #1 + %call3 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call4 = call i32 (i32, ...) bitcast (i32 (...)* @savectx to i32 (i32, ...)*)(i32 0) #1 + %call5 = call i32 @vfork() #1 + %call6 = call i32 (i32, ...) bitcast (i32 (...)* @setzx to i32 (i32, ...)*)(i32 0) + %call7 = call i32 (i32, ...) bitcast (i32 (...)* @getcontext to i32 (i32, ...)*)(i32 0) #1 + ret void +} + +; Function Attrs: returns_twice +declare dso_local i32 @setjmp(...) #1 + +declare dso_local i32 @setzx(...) + +; Function Attrs: returns_twice +declare dso_local i32 @sigsetjmp(...) #1 + +; Function Attrs: returns_twice +declare dso_local i32 @savectx(...) #1 + +; Function Attrs: returns_twice +declare dso_local i32 @vfork() #1 + +; Function Attrs: returns_twice +declare dso_local i32 @getcontext(...) #1 + +attributes #0 = { noinline nounwind optnone uwtable } +attributes #1 = { returns_twice } + +!llvm.module.flags = !{!0, !1, !2} + +!0 = !{i32 1, !"wchar_size", i32 4} +!1 = !{i32 4, !"cf-protection-return", i32 1} +!2 = !{i32 4, !"cf-protection-branch", i32 1}