Index: include/llvm/Analysis/BranchProbabilityInfo.h =================================================================== --- include/llvm/Analysis/BranchProbabilityInfo.h +++ include/llvm/Analysis/BranchProbabilityInfo.h @@ -111,6 +111,10 @@ void setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors, uint32_t Weight); + static uint32_t getBranchWeightStackProtector(bool Likely) { + return Likely ? 1024*1024 - 1 : 1; + } + private: // Since we allow duplicate edges from one basic block to another, we use // a pair (PredBlock and an index in the successors) to specify an edge. Index: lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h =================================================================== --- lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h +++ lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h @@ -416,8 +416,8 @@ assert(!shouldEmitStackProtector() && "Stack Protector Descriptor is " "already initialized!"); ParentMBB = MBB; - SuccessMBB = AddSuccessorMBB(BB, MBB); - FailureMBB = AddSuccessorMBB(BB, MBB, FailureMBB); + SuccessMBB = AddSuccessorMBB(BB, MBB, true); + FailureMBB = AddSuccessorMBB(BB, MBB, false, FailureMBB); if (!Guard) Guard = StackProtCheckCall.getArgOperand(0); } @@ -488,7 +488,7 @@ /// has not been created yet (i.e. if SuccMBB = 0), then the machine basic /// block will be created. MachineBasicBlock *AddSuccessorMBB(const BasicBlock *BB, - MachineBasicBlock *ParentMBB, + MachineBasicBlock *ParentMBB, bool Likely, MachineBasicBlock *SuccMBB = nullptr); }; Index: lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp =================================================================== --- lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp +++ lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp @@ -7739,6 +7739,7 @@ SelectionDAGBuilder::StackProtectorDescriptor:: AddSuccessorMBB(const BasicBlock *BB, MachineBasicBlock *ParentMBB, + bool Likely, MachineBasicBlock *SuccMBB) { // If SuccBB has not been created yet, create it. if (!SuccMBB) { @@ -7748,6 +7749,7 @@ MF->insert(++BBI, SuccMBB); } // Add it as a successor of ParentMBB. - ParentMBB->addSuccessor(SuccMBB); + ParentMBB->addSuccessor( + SuccMBB, BranchProbabilityInfo::getBranchWeightStackProtector(Likely)); return SuccMBB; } Index: lib/CodeGen/StackProtector.cpp =================================================================== --- lib/CodeGen/StackProtector.cpp +++ lib/CodeGen/StackProtector.cpp @@ -17,6 +17,7 @@ #include "llvm/CodeGen/StackProtector.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/Statistic.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" #include "llvm/Analysis/ValueTracking.h" #include "llvm/CodeGen/Analysis.h" #include "llvm/CodeGen/Passes.h" @@ -31,6 +32,7 @@ #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/Intrinsics.h" +#include "llvm/IR/MDBuilder.h" #include "llvm/IR/Module.h" #include "llvm/Support/CommandLine.h" #include "llvm/Target/TargetSubtargetInfo.h" @@ -459,7 +461,13 @@ LoadInst *LI1 = B.CreateLoad(StackGuardVar); LoadInst *LI2 = B.CreateLoad(AI); Value *Cmp = B.CreateICmpEQ(LI1, LI2); - B.CreateCondBr(Cmp, NewBB, FailBB); + unsigned SuccessWeight = + BranchProbabilityInfo::getBranchWeightStackProtector(true); + unsigned FailureWeight = + BranchProbabilityInfo::getBranchWeightStackProtector(false); + MDNode *Weights = MDBuilder(F->getContext()) + .createBranchWeights(SuccessWeight, FailureWeight); + B.CreateCondBr(Cmp, NewBB, FailBB, Weights); } } Index: test/CodeGen/X86/stack-protector-weight.ll =================================================================== --- /dev/null +++ test/CodeGen/X86/stack-protector-weight.ll @@ -0,0 +1,36 @@ +; RUN: llc -mtriple=x86_64-apple-darwin -print-machineinstrs=expand-isel-pseudos -enable-selectiondag-sp=true %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SELDAG +; RUN: llc -mtriple=x86_64-apple-darwin -print-machineinstrs=expand-isel-pseudos -enable-selectiondag-sp=false %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR + +; SELDAG: # Machine code for function test_branch_weights: +; SELDAG: Successors according to CFG: BB#[[SUCCESS:[0-9]+]](1048575) BB#[[FAILURE:[0-9]+]](1) +; SELDAG: BB#[[FAILURE]]: +; SELDAG: CALL64pcrel32 +; SELDAG: BB#[[SUCCESS]]: + +; IR: # Machine code for function test_branch_weights: +; IR: Successors according to CFG: BB#[[SUCCESS:[0-9]+]](1048575) BB#[[FAILURE:[0-9]+]](1) +; IR: BB#[[SUCCESS]]: +; IR: BB#[[FAILURE]]: +; IR: CALL64pcrel32 + +define i32 @test_branch_weights(i32 %n) #0 { +entry: + %a = alloca [128 x i32], align 16 + %0 = bitcast [128 x i32]* %a to i8* + call void @llvm.lifetime.start(i64 512, i8* %0) + %arraydecay = getelementptr inbounds [128 x i32]* %a, i64 0, i64 0 + call void @foo2(i32* %arraydecay) + %idxprom = sext i32 %n to i64 + %arrayidx = getelementptr inbounds [128 x i32]* %a, i64 0, i64 %idxprom + %1 = load i32* %arrayidx, align 4 + call void @llvm.lifetime.end(i64 512, i8* %0) + ret i32 %1 +} + +declare void @llvm.lifetime.start(i64, i8* nocapture) + +declare void @foo2(i32*) + +declare void @llvm.lifetime.end(i64, i8* nocapture) + +attributes #0 = { ssp "stack-protector-buffer-size"="8" }