[Profile] Preserve branch meta data in CGP lowering of select instruction

[davidxl]

Summary

When CGP lowers select back to control flow, the weight meta data gets dropped. This patch fixes the problem.

[davidxl]

The builder interface should copy all branch related meta data if-available. The new interface will also replace similar code in loop unswitch (followup).

Dehao, the new change also copies over debug location meta data to the branch. What is the potential impact of this to Sample PGO?

[danielcdh]

This shouldn't affect sample pgo because we don't use branch instruction for annotation.

llvm/trunk/include/llvm/IR/IRBuilder.h

public:
   /// instruction.
   BranchInst *CreateCondBr(Value *Cond, BasicBlock *True, BasicBlock *False,
                            MDNode *BranchWeights = nullptr,
                            MDNode *Unpredictable = nullptr) {
     return Insert(addBranchMetadata(BranchInst::Create(True, False, Cond),
                                     BranchWeights, Unpredictable));
   }
 
+  /// \brief Create a conditional 'br Cond, TrueDest, FalseDest'
+  /// instruction. Copy branch meta data if available.
+  BranchInst *CreateCondBr(Value *Cond, BasicBlock *True, BasicBlock *False,
+                           Instruction *MDSrc) {
+    BranchInst *Br = BranchInst::Create(True, False, Cond);
+    if (MDSrc) {
+      unsigned WL[4] = {LLVMContext::MD_prof, LLVMContext::MD_unpredictable,
+                        LLVMContext::MD_make_implicit, LLVMContext::MD_dbg};
+      Br->copyMetadata(*MDSrc, makeArrayRef(&WL[0], 4));
+    }
+    return Insert(Br);
+  }
+
   /// \brief Create a switch instruction with the specified value, default dest,
   /// and with a hint for the number of cases that will be added (for efficient
   /// allocation).
   SwitchInst *CreateSwitch(Value *V, BasicBlock *Dest, unsigned NumCases = 10,
                            MDNode *BranchWeights = nullptr,
                            MDNode *Unpredictable = nullptr) {
     return Insert(addBranchMetadata(SwitchInst::Create(V, Dest, NumCases),
                                     BranchWeights, Unpredictable));

llvm/trunk/lib/CodeGen/CodeGenPrepare.cpp

if (TrueBlock == FalseBlock) {
     BranchInst::Create(EndBlock, FalseBlock);
   }
 
   // Insert the real conditional branch based on the original condition.
   // If we did not create a new block for one of the 'true' or 'false' paths
   // of the condition, it means that side of the branch goes to the end block
   // directly and the path originates from the start block from the point of
   // view of the new PHI.
+  BasicBlock *TT, *FT;
   if (TrueBlock == nullptr) {
-    BranchInst::Create(EndBlock, FalseBlock, SI->getCondition(), SI);
+    TT = EndBlock;
+    FT = FalseBlock;
     TrueBlock = StartBlock;
   } else if (FalseBlock == nullptr) {
-    BranchInst::Create(TrueBlock, EndBlock, SI->getCondition(), SI);
+    TT = TrueBlock;
+    FT = EndBlock;
     FalseBlock = StartBlock;
   } else {
-    BranchInst::Create(TrueBlock, FalseBlock, SI->getCondition(), SI);
+    TT = TrueBlock;
+    FT = FalseBlock;
   }
+  IRBuilder<>(SI).CreateCondBr(SI->getCondition(), TT, FT, SI);
 
   // The select itself is replaced with a PHI Node.
   PHINode *PN = PHINode::Create(SI->getType(), 2, "", &EndBlock->front());
   PN->takeName(SI);
   PN->addIncoming(SI->getTrueValue(), TrueBlock);
   PN->addIncoming(SI->getFalseValue(), FalseBlock);
 
   SI->replaceAllUsesWith(PN);

llvm/trunk/test/CodeGen/X86/cmov-into-branch.ll

 ; Note the reversed profile weights: it doesn't matter if it's
 ; obviously true or obviously false.
 ; Either one should become a branch rather than conditional move.
 ; TODO: But likely true vs. likely false should affect basic block placement?
 define i32 @weighted_select3(i32 %a, i32 %b) {
 ; CHECK-LABEL: weighted_select3:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    testl %edi, %edi
-; CHECK-NEXT:    jne [[LABEL_BB6:.*]]
-; CHECK:         movl %esi, %edi
-; CHECK-NEXT:  [[LABEL_BB6]]
-; CHECK-NEXT:    movl %edi, %eax
-; CHECK-NEXT:    retq
+; CHECK-NEXT:    je [[LABEL_BB6:.*]]
+; CHECK:         movl %edi, %eax
+; CHECK:         [[LABEL_BB6]]
+; CHECK-NEXT:    movl %esi, %edi
+; CHECK-NEXT:    jmp
 ;
   %cmp = icmp ne i32 %a, 0
   %sel = select i1 %cmp, i32 %a, i32 %b, !prof !2
   ret i32 %sel
 }
 
 ; Weightlessness is no reason to die.
 define i32 @unweighted_select(i32 %a, i32 %b) {

llvm/trunk/test/CodeGen/X86/select_meta.ll

+; RUN: llc -mtriple=x86_64-unknown-unknown -print-after-all < %s 2>&1 | FileCheck %s
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @foo(i32, i32, i32) {
+  %4 = and i32 %0, 3
+  %5 = icmp eq i32 %4, 1
+  %6 = select i1 %5, i32 %1, i32 %2, !prof  !1
+; CHECK: br {{.*}}label{{.*}}, label{{.*}}, !prof ![[WT:.*]]
+  ret i32 %6
+}
+
+!llvm.ident = !{!0}
+
+!0 = !{!"clang version 4.0.0 (trunk 279683)"}
+!1 = !{!"branch_weights", i32 1000, i32 1 }
+; CHECK ![[WT]] = !{!"branch_weights", i32 1000, i32 1 }