diff --git a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp --- a/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp +++ b/llvm/lib/Target/X86/X86IndirectBranchTracking.cpp @@ -18,6 +18,7 @@ #include "X86.h" #include "X86InstrInfo.h" #include "X86Subtarget.h" +#include "X86TargetMachine.h" #include "llvm/ADT/Statistic.h" #include "llvm/CodeGen/MachineFunctionPass.h" #include "llvm/CodeGen/MachineInstrBuilder.h" @@ -102,7 +103,16 @@ // Check that the cf-protection-branch is enabled. Metadata *isCFProtectionSupported = MF.getMMI().getModule()->getModuleFlag("cf-protection-branch"); - if (!isCFProtectionSupported && !IndirectBranchTracking) + // NB: We need to enable IBT in jitted code if JIT compiler is CET + // enabled. + const X86TargetMachine *TM = + static_cast(&MF.getTarget()); +#ifdef __CET__ + bool isJITwithCET = TM->isJIT(); +#else + bool isJITwithCET = false; +#endif + if (!isCFProtectionSupported && !IndirectBranchTracking && !isJITwithCET) return false; // True if the current MF was changed and false otherwise. @@ -111,10 +121,11 @@ TII = SubTarget.getInstrInfo(); EndbrOpcode = SubTarget.is64Bit() ? X86::ENDBR64 : X86::ENDBR32; - // Non-internal function or function whose address was taken, can be - // accessed through indirect calls. Mark the first BB with ENDBR instruction - // unless nocf_check attribute is used. - if ((MF.getFunction().hasAddressTaken() || + // Large code model, non-internal function or function whose address + // was taken, can be accessed through indirect calls. Mark the first + // BB with ENDBR instruction unless nocf_check attribute is used. + if ((TM->getCodeModel() == CodeModel::Large || + MF.getFunction().hasAddressTaken() || !MF.getFunction().hasLocalLinkage()) && !MF.getFunction().doesNoCfCheck()) { auto MBB = MF.begin(); @@ -136,8 +147,8 @@ Changed |= addENDBR(MBB, std::next(I)); if (EHPadIBTNeeded && I->isEHLabel()) { - Changed |= addENDBR(MBB, std::next(I)); - EHPadIBTNeeded = false; + Changed |= addENDBR(MBB, std::next(I)); + EHPadIBTNeeded = false; } } } diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h --- a/llvm/lib/Target/X86/X86TargetMachine.h +++ b/llvm/lib/Target/X86/X86TargetMachine.h @@ -30,6 +30,8 @@ class X86TargetMachine final : public LLVMTargetMachine { std::unique_ptr TLOF; mutable StringMap> SubtargetMap; + // True if this is used in JIT. + bool IsJIT; public: X86TargetMachine(const Target &T, const Triple &TT, StringRef CPU, @@ -52,6 +54,8 @@ TargetLoweringObjectFile *getObjFileLowering() const override { return TLOF.get(); } + + bool isJIT() const { return IsJIT; } }; } // end namespace llvm diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp --- a/llvm/lib/Target/X86/X86TargetMachine.cpp +++ b/llvm/lib/Target/X86/X86TargetMachine.cpp @@ -214,7 +214,7 @@ getEffectiveRelocModel(TT, JIT, RM), getEffectiveX86CodeModel(CM, JIT, TT.getArch() == Triple::x86_64), OL), - TLOF(createTLOF(getTargetTriple())) { + TLOF(createTLOF(getTargetTriple())), IsJIT(JIT) { // On PS4, the "return address" of a 'noreturn' call must still be within // the calling function, and TrapUnreachable is an easy way to get that. if (TT.isPS4() || TT.isOSBinFormatMachO()) { diff --git a/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll b/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/X86/indirect-branch-tracking-cm-lager.ll @@ -0,0 +1,36 @@ +; RUN: llc -mtriple=x86_64-unknown-unknown -code-model=large < %s | FileCheck %s + +; In large code model indirect branches are needed when branching to addresses +; whose offset from the current instruction pointer is unknown. +;CHECK-COUNT-3: endbr + +@a = dso_local local_unnamed_addr global i32 1, align 4 + +; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly +define dso_local void @ext() local_unnamed_addr #0 { +entry: + store i32 0, i32* @a, align 4 + ret void +} + +; Function Attrs: nofree norecurse nounwind uwtable +define dso_local i32 @main() local_unnamed_addr #1 { +entry: + tail call fastcc void @foo() + %0 = load i32, i32* @a, align 4 + ret i32 %0 +} + +; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly +define internal fastcc void @foo() unnamed_addr #0 { +entry: + tail call void @ext() + ret void +} + +!llvm.module.flags = !{!0, !1, !2, !3} + +!0 = !{i32 1, !"wchar_size", i32 4} +!1 = !{i32 4, !"cf-protection-return", i32 1} +!2 = !{i32 4, !"cf-protection-branch", i32 1} +!3 = !{i32 1, !"Code Model", i32 4} diff --git a/llvm/test/ExecutionEngine/MCJIT/cet-code-model-lager.ll b/llvm/test/ExecutionEngine/MCJIT/cet-code-model-lager.ll new file mode 100644 --- /dev/null +++ b/llvm/test/ExecutionEngine/MCJIT/cet-code-model-lager.ll @@ -0,0 +1,33 @@ +; RUN: %lli -code-model=large %s > /dev/null +target triple = "x86_64-unknown-linux-gnu" + +@a = dso_local local_unnamed_addr global i32 1, align 4 + +; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly +define dso_local void @ext() local_unnamed_addr #0 { +entry: + store i32 0, i32* @a, align 4 + ret void +} + +; Function Attrs: nofree norecurse nounwind uwtable +define dso_local i32 @main() local_unnamed_addr #1 { +entry: + tail call fastcc void @foo() + %0 = load i32, i32* @a, align 4 + ret i32 %0 +} + +; Function Attrs: nofree noinline norecurse nounwind uwtable writeonly +define internal fastcc void @foo() unnamed_addr #0 { +entry: + tail call void @ext() + ret void +} + +!llvm.module.flags = !{!0, !1, !2, !3} + +!0 = !{i32 1, !"wchar_size", i32 4} +!1 = !{i32 4, !"cf-protection-return", i32 1} +!2 = !{i32 4, !"cf-protection-branch", i32 1} +!3 = !{i32 1, !"Code Model", i32 4}