[PATCH] Set edge weights in stack protector pass

[Duncan P. N. Exon Smith]

> On 2014-Nov-03, at 13:04, Akira Hatanaka <ahatanak at gmail.com> wrote:
> 
> ping
> 
> http://reviews.llvm.org/D5766
> 

LGTM with a few nitpicks.

> 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;

`IsLikely` is a better name, and `1024*1024` should have spaces around
the `*`.

Bike shed: I personally find `(1u << 20) - 1` less magical :).

> +  }
> +
>  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.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,

Same name change here.

>                  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/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);

I think these should have `/* IsLikely */` comments:

    SuccessMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ true);
    FailureMBB = AddSuccessorMBB(BB, MBB, /* IsLikely */ 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,

Same name change.  Also, you should update the doxygen to explain the
variable.

>                                         MachineBasicBlock *SuccMBB = nullptr);
>    };
>  
> 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 <es:__stack_chk_fail>
> +; 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 <ga:@__stack_chk_fail>
> +
> +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" }
>