Add support for attaching branch weight metadata directly from the IRBuilder.

Added a basic unit test for this with CreateCondBr. I didn't go all the
way and test the switch side as the boilerplate for setting up the
switch IRBuilder unit tests is a lot more. Fortunately, the two share
all the interesting code paths.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@160251 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 30 insertions, 4 deletions

diff --git a/include/llvm/IRBuilder.h b/include/llvm/IRBuilder.h
index c6200273b1..d5b6f47f8a 100644
--- a/include/llvm/IRBuilder.h
+++ b/include/llvm/IRBuilder.h
@@ -411,6 +411,17 @@ public:
   // Instruction creation methods: Terminators
   //===--------------------------------------------------------------------===//
 
+private:
+  /// \brief Helper to add branch weight metadata onto an instruction.
+  /// \returns The annotated instruction.
+  template <typename InstTy>
+  InstTy *addBranchWeights(InstTy *I, MDNode *Weights) {
+    if (Weights)
+      I->setMetadata(LLVMContext::MD_prof, Weights);
+    return I;
+  }
+
+public:
   /// CreateRetVoid - Create a 'ret void' instruction.
   ReturnInst *CreateRetVoid() {
     return Insert(ReturnInst::Create(Context));
@@ -444,15 +455,19 @@ public:
 
   /// CreateCondBr - Create a conditional 'br Cond, TrueDest, FalseDest'
   /// instruction.
-  BranchInst *CreateCondBr(Value *Cond, BasicBlock *True, BasicBlock *False) {
-    return Insert(BranchInst::Create(True, False, Cond));
+  BranchInst *CreateCondBr(Value *Cond, BasicBlock *True, BasicBlock *False,
+                           MDNode *BranchWeights = 0) {
+    return Insert(addBranchWeights(BranchInst::Create(True, False, Cond),
+                                   BranchWeights));
   }
 
   /// CreateSwitch - 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) {
-    return Insert(SwitchInst::Create(V, Dest, NumCases));
+  SwitchInst *CreateSwitch(Value *V, BasicBlock *Dest, unsigned NumCases = 10,
+                           MDNode *BranchWeights = 0) {
+    return Insert(addBranchWeights(SwitchInst::Create(V, Dest, NumCases),
+                                   BranchWeights));
   }
 
   /// CreateIndirectBr - Create an indirect branch instruction with the
diff --git a/unittests/VMCore/IRBuilderTest.cpp b/unittests/VMCore/IRBuilderTest.cpp
index 3eb5926d54..b6a3795fd0 100644
--- a/unittests/VMCore/IRBuilderTest.cpp
+++ b/unittests/VMCore/IRBuilderTest.cpp
@@ -12,6 +12,7 @@
 #include "llvm/IRBuilder.h"
 #include "llvm/IntrinsicInst.h"
 #include "llvm/LLVMContext.h"
+#include "llvm/MDBuilder.h"
 #include "llvm/Module.h"
 #include "llvm/ADT/OwningPtr.h"
 
@@ -83,6 +84,16 @@ TEST_F(IRBuilderTest, CreateCondBr) {
   EXPECT_EQ(2u, TI->getNumSuccessors());
   EXPECT_EQ(TBB, TI->getSuccessor(0));
   EXPECT_EQ(FBB, TI->getSuccessor(1));
+
+  BI->eraseFromParent();
+  MDNode *Weights = MDBuilder(getGlobalContext()).createBranchWeights(42, 13);
+  BI = Builder.CreateCondBr(Builder.getTrue(), TBB, FBB, Weights);
+  TI = BB->getTerminator();
+  EXPECT_EQ(BI, TI);
+  EXPECT_EQ(2u, TI->getNumSuccessors());
+  EXPECT_EQ(TBB, TI->getSuccessor(0));
+  EXPECT_EQ(FBB, TI->getSuccessor(1));
+  EXPECT_EQ(Weights, TI->getMetadata(LLVMContext::MD_prof));
 }
 
 }