From ca780b4578b72e847b555577c085741b8105ca26 Mon Sep 17 00:00:00 2001 From: Akira Hatanaka Date: Mon, 1 Dec 2014 04:27:03 +0000 Subject: [PATCH] [stack protector] Set edge weights for newly created basic blocks. This commit fixes a bug in stack protector pass where edge weights were not set when new basic blocks were added to lists of successor basic blocks. Differential Revision: http://reviews.llvm.org/D5766 git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@222987 91177308-0d34-0410-b5e6-96231b3b80d8 --- include/llvm/Analysis/BranchProbabilityInfo.h | 4 +++ .../SelectionDAG/SelectionDAGBuilder.cpp | 4 ++- .../SelectionDAG/SelectionDAGBuilder.h | 7 ++-- lib/CodeGen/StackProtector.cpp | 10 +++++- test/CodeGen/X86/stack-protector-weight.ll | 36 +++++++++++++++++++ 5 files changed, 56 insertions(+), 5 deletions(-) create mode 100644 test/CodeGen/X86/stack-protector-weight.ll diff --git a/include/llvm/Analysis/BranchProbabilityInfo.h b/include/llvm/Analysis/BranchProbabilityInfo.h index 4414c84f6b7..89eef68d843 100644 --- a/include/llvm/Analysis/BranchProbabilityInfo.h +++ b/include/llvm/Analysis/BranchProbabilityInfo.h @@ -111,6 +111,10 @@ public: void setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors, uint32_t Weight); + static uint32_t getBranchWeightStackProtector(bool IsLikely) { + return IsLikely ? (1u << 20) - 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. diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp index 3ab89f6de59..fbc1502d500 100644 --- a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp +++ b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp @@ -7763,6 +7763,7 @@ MachineBasicBlock * SelectionDAGBuilder::StackProtectorDescriptor:: AddSuccessorMBB(const BasicBlock *BB, MachineBasicBlock *ParentMBB, + bool IsLikely, MachineBasicBlock *SuccMBB) { // If SuccBB has not been created yet, create it. if (!SuccMBB) { @@ -7772,6 +7773,7 @@ AddSuccessorMBB(const BasicBlock *BB, MF->insert(++BBI, SuccMBB); } // Add it as a successor of ParentMBB. - ParentMBB->addSuccessor(SuccMBB); + ParentMBB->addSuccessor( + SuccMBB, BranchProbabilityInfo::getBranchWeightStackProtector(IsLikely)); return SuccMBB; } diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h index f74e6525b0c..b2215f22ab1 100644 --- a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h +++ b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h @@ -417,8 +417,8 @@ private: assert(!shouldEmitStackProtector() && "Stack Protector Descriptor is " "already initialized!"); ParentMBB = MBB; - SuccessMBB = AddSuccessorMBB(BB, MBB); - FailureMBB = AddSuccessorMBB(BB, MBB, FailureMBB); + SuccessMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ true); + FailureMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ false, FailureMBB); if (!Guard) Guard = StackProtCheckCall.getArgOperand(0); } @@ -487,9 +487,10 @@ private: /// Add a successor machine basic block to ParentMBB. If the successor mbb /// has not been created yet (i.e. if SuccMBB = 0), then the machine basic - /// block will be created. + /// block will be created. Assign a large weight if IsLikely is true. MachineBasicBlock *AddSuccessorMBB(const BasicBlock *BB, MachineBasicBlock *ParentMBB, + bool IsLikely, MachineBasicBlock *SuccMBB = nullptr); }; diff --git a/lib/CodeGen/StackProtector.cpp b/lib/CodeGen/StackProtector.cpp index 3fc0a36dd22..58e4ad971e8 100644 --- a/lib/CodeGen/StackProtector.cpp +++ b/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 @@ bool StackProtector::InsertStackProtectors() { 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); } } diff --git a/test/CodeGen/X86/stack-protector-weight.ll b/test/CodeGen/X86/stack-protector-weight.ll new file mode 100644 index 00000000000..c5bf49134e4 --- /dev/null +++ b/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" }