diff --git a/include/llvm/IRBuilder.h b/include/llvm/IRBuilder.h index c6200273b1c..d5b6f47f8a2 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 + 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 3eb5926d542..b6a3795fd0f 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)); } }