2 files changed, 66 insertions, 2 deletions

diff --git a/llvm/lib/CodeGen/TailDuplication.cpp b/llvm/lib/CodeGen/TailDuplication.cpp
index ff0181e0f0c..170a8e19ed4 100644
--- a/llvm/lib/CodeGen/TailDuplication.cpp
+++ b/llvm/lib/CodeGen/TailDuplication.cpp
@@ -19,6 +19,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
@@ -61,6 +62,7 @@ namespace {
   class TailDuplicatePass : public MachineFunctionPass {
     const TargetInstrInfo *TII;
     const TargetRegisterInfo *TRI;
+    const MachineBranchProbabilityInfo *MBPI;
     MachineModuleInfo *MMI;
     MachineRegisterInfo *MRI;
     OwningPtr<RegScavenger> RS;
@@ -80,6 +82,8 @@ namespace {
 
     virtual bool runOnMachineFunction(MachineFunction &MF);
 
+    virtual void getAnalysisUsage(AnalysisUsage &AU) const;
+
   private:
     void AddSSAUpdateEntry(unsigned OrigReg, unsigned NewReg,
                            MachineBasicBlock *BB);
@@ -132,6 +136,8 @@ bool TailDuplicatePass::runOnMachineFunction(MachineFunction &MF) {
   TRI = MF.getTarget().getRegisterInfo();
   MRI = &MF.getRegInfo();
   MMI = getAnalysisIfAvailable<MachineModuleInfo>();
+  MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
+
   PreRegAlloc = MRI->isSSA();
   RS.reset();
   if (MRI->tracksLiveness() && TRI->trackLivenessAfterRegAlloc(MF))
@@ -144,6 +150,11 @@ bool TailDuplicatePass::runOnMachineFunction(MachineFunction &MF) {
   return MadeChange;
 }
 
+void TailDuplicatePass::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<MachineBranchProbabilityInfo>();
+  MachineFunctionPass::getAnalysisUsage(AU);
+}
+
 static void VerifyPHIs(MachineFunction &MF, bool CheckExtra) {
   for (MachineFunction::iterator I = ++MF.begin(), E = MF.end(); I != E; ++I) {
     MachineBasicBlock *MBB = I;
@@ -721,11 +732,12 @@ TailDuplicatePass::duplicateSimpleBB(MachineBasicBlock *TailBB,
     if (PredTBB)
       TII->InsertBranch(*PredBB, PredTBB, PredFBB, PredCond, DebugLoc());
 
+    uint32_t Weight = MBPI->getEdgeWeight(PredBB, TailBB);
     PredBB->removeSuccessor(TailBB);
     unsigned NumSuccessors = PredBB->succ_size();
     assert(NumSuccessors <= 1);
     if (NumSuccessors == 0 || *PredBB->succ_begin() != NewTarget)
-      PredBB->addSuccessor(NewTarget);
+      PredBB->addSuccessor(NewTarget, Weight);
 
     TDBBs.push_back(PredBB);
   }
@@ -836,7 +848,7 @@ TailDuplicatePass::TailDuplicate(MachineBasicBlock *TailBB,
            "TailDuplicate called on block with multiple successors!");
     for (MachineBasicBlock::succ_iterator I = TailBB->succ_begin(),
            E = TailBB->succ_end(); I != E; ++I)
-      PredBB->addSuccessor(*I);
+      PredBB->addSuccessor(*I, MBPI->getEdgeWeight(TailBB, I));
 
     Changed = true;
     ++NumTailDups;
diff --git a/llvm/test/CodeGen/ARM/taildup-branch-weight.ll b/llvm/test/CodeGen/ARM/taildup-branch-weight.ll
new file mode 100644
index 00000000000..80bbc8c3238
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/taildup-branch-weight.ll
@@ -0,0 +1,52 @@
+; RUN: llc < %s -march=arm -print-machineinstrs=tailduplication -tail-dup-size=100 -enable-tail-merge=false -disable-cgp -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: Machine code for function test0:
+; CHECK: Successors according to CFG: BB#1(4) BB#2(124)
+
+define void @test0(i32 %a, i32 %b, i32* %c, i32* %d) {
+entry:
+  store i32 3, i32* %d
+  br label %B1
+
+B2:
+  store i32 2, i32* %c
+  br label %B4
+
+B3:
+  store i32 2, i32* %c
+  br label %B4
+
+B1:
+  store i32 1, i32* %d
+  %test0 = icmp slt i32 %a, %b
+  br i1 %test0, label %B2, label %B3, !prof !0
+
+B4:
+  ret void
+}
+
+!0 = metadata !{metadata !"branch_weights", i32 4, i32 124}
+
+; CHECK: Machine code for function test1:
+; CHECK: Successors according to CFG: BB#1(8) BB#2(248)
+
+@g0 = common global i32 0, align 4
+
+define void @test1(i32 %a, i32 %b, i32* %c, i32* %d, i32* %e) {
+
+  %test0 = icmp slt i32 %a, %b
+  br i1 %test0, label %B1, label %B2, !prof !1
+
+B1:
+  br label %B3
+
+B2:
+  store i32 2, i32* %c
+  br label %B3
+
+B3:
+  store i32 3, i32* %e
+  ret void
+}
+
+!1 = metadata !{metadata !"branch_weights", i32 248, i32 8}