Index: llvm/trunk/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp =================================================================== --- llvm/trunk/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp +++ llvm/trunk/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp @@ -83,9 +83,8 @@ return true; } -static bool handleBranchExpect(BranchInst &BI) { - if (BI.isUnconditional()) - return false; +// Handle both BranchInst and SelectInst. +template static bool handleBrSelExpect(BrSelInst &BSI) { // Handle non-optimized IR code like: // %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1) @@ -98,9 +97,9 @@ CallInst *CI; - ICmpInst *CmpI = dyn_cast(BI.getCondition()); + ICmpInst *CmpI = dyn_cast(BSI.getCondition()); if (!CmpI) { - CI = dyn_cast(BI.getCondition()); + CI = dyn_cast(BSI.getCondition()); } else { if (CmpI->getPredicate() != CmpInst::ICMP_NE) return false; @@ -129,15 +128,22 @@ else Node = MDB.createBranchWeights(UnlikelyBranchWeight, LikelyBranchWeight); - BI.setMetadata(LLVMContext::MD_prof, Node); + BSI.setMetadata(LLVMContext::MD_prof, Node); if (CmpI) CmpI->setOperand(0, ArgValue); else - BI.setCondition(ArgValue); + BSI.setCondition(ArgValue); return true; } +static bool handleBranchExpect(BranchInst &BI) { + if (BI.isUnconditional()) + return false; + + return handleBrSelExpect(BI); +} + static bool lowerExpectIntrinsic(Function &F) { bool Changed = false; @@ -151,11 +157,19 @@ ExpectIntrinsicsHandled++; } - // Remove llvm.expect intrinsics. - for (BasicBlock::iterator BI = BB.begin(), BE = BB.end(); BI != BE;) { - CallInst *CI = dyn_cast(BI++); - if (!CI) + // Remove llvm.expect intrinsics. Iterate backwards in order + // to process select instructions before the intrinsic gets + // removed. + for (auto BI = BB.rbegin(), BE = BB.rend(); BI != BE;) { + Instruction *Inst = &*BI++; + CallInst *CI = dyn_cast(Inst); + if (!CI) { + if (SelectInst *SI = dyn_cast(Inst)) { + if (handleBrSelExpect(*SI)) + ExpectIntrinsicsHandled++; + } continue; + } Function *Fn = CI->getCalledFunction(); if (Fn && Fn->getIntrinsicID() == Intrinsic::expect) { Index: llvm/trunk/test/Transforms/LowerExpectIntrinsic/basic.ll =================================================================== --- llvm/trunk/test/Transforms/LowerExpectIntrinsic/basic.ll +++ llvm/trunk/test/Transforms/LowerExpectIntrinsic/basic.ll @@ -273,6 +273,16 @@ ret i32 %0 } +; CHECK-LABEL: @test10( +define i32 @test10(i64 %t6) { + %t7 = call i64 @llvm.expect.i64(i64 %t6, i64 0) + %t8 = icmp ne i64 %t7, 0 + %t9 = select i1 %t8, i32 1, i32 2 +; CHECK: select{{.*}}, !prof !1 + ret i32 %t9 +} + + declare i1 @llvm.expect.i1(i1, i1) nounwind readnone ; CHECK: !0 = !{!"branch_weights", i32 2000, i32 1}