[Profile] Lower expect intrinsic properly for select instruction

[davidxl]

Summary

builtin expect passed to select instruction gets ignored. This patch fixes the issue.

[davidxl]

Make the implementation more efficient -- avoid traversing IR twice.

[spatel]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

86

Neat trick. :)
But can you add a comment and/or rename BI, so it is clear that this isn't strictly for a 'BI' (BranchInst) any more?

165–167

Fix comment to indicate why we are using reverse iterator (and that we're also handling selects in this loop).

168

Use 'auto' to avoid 80-col ugliness?

test/Transforms/LowerExpectIntrinsic/basic.ll

276–287

Minimize test?

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 42, i32 43
  ret i32 %t9
}

[davidxl]

Addressed review comments.

[spatel]

LGTM.

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

static bool handleSwitchExpect(SwitchInst &SI) {
 
   SI.setMetadata(LLVMContext::MD_prof,
                  MDBuilder(CI->getContext()).createBranchWeights(Weights));
 
   SI.setCondition(ArgValue);
   return true;
 }
 
-static bool handleBranchExpect(BranchInst &BI) {
-  if (BI.isUnconditional())
-    return false;
+template<class BrSelInst> static bool handleBrSelExpect(BrSelInst &BI) {
 
   // Handle non-optimized IR code like:
   //   %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1)
   //   %tobool = icmp ne i64 %expval, 0
   //   br i1 %tobool, label %if.then, label %if.end
   //
   // Or the following simpler case:
   //   %expval = call i1 @llvm.expect.i1(i1 %cmp, i1 1)
template<class BrSelInst> static bool handleBrSelExpect(BrSelInst &BI) {
 
   if (CmpI)
     CmpI->setOperand(0, ArgValue);
   else
     BI.setCondition(ArgValue);
   return true;
 }
 
+static bool handleBranchExpect(BranchInst &BI) {
+  if (BI.isUnconditional())
+    return false;
+
+  return handleBrSelExpect<BranchInst>(BI);
+}
+
 static bool lowerExpectIntrinsic(Function &F) {
   bool Changed = false;
 
   for (BasicBlock &BB : F) {
     // Create "block_weights" metadata.
     if (BranchInst *BI = dyn_cast<BranchInst>(BB.getTerminator())) {
       if (handleBranchExpect(*BI))
         ExpectIntrinsicsHandled++;
     } else if (SwitchInst *SI = dyn_cast<SwitchInst>(BB.getTerminator())) {
       if (handleSwitchExpect(*SI))
         ExpectIntrinsicsHandled++;
     }
 
+    for (BasicBlock::iterator BI = BB.begin(), BE = BB.end(); BI != BE;) {
+      if (SelectInst *SI = dyn_cast<SelectInst>(BI++))
+        if (handleBrSelExpect(*SI))
+          ExpectIntrinsicsHandled++;
+    }
+
     // Remove llvm.expect intrinsics.
     for (BasicBlock::iterator BI = BB.begin(), BE = BB.end(); BI != BE;) {
       CallInst *CI = dyn_cast<CallInst>(BI++);
       if (!CI)
         continue;
 
       Function *Fn = CI->getCalledFunction();
       if (Fn && Fn->getIntrinsicID() == Intrinsic::expect) {
         Value *Exp = CI->getArgOperand(0);
         CI->replaceAllUsesWith(Exp);
         CI->eraseFromParent();
         Changed = true;

test/Transforms/LowerExpectIntrinsic/basic.ll

   store i32 1, i32* %retval
   br label %return
 
 return:                                           ; preds = %if.end, %if.then
   %0 = load i32, i32* %retval
   ret i32 %0
 }
 
+; CHECK-LABEL: @test10(
+define i32 @test10(i32, i32, i32) #0 {
+  %4 = and i32 %0, 3
+  %5 = icmp eq i32 %4, 1
+  %6 = zext i1 %5 to i64
+  %7 = call i64 @llvm.expect.i64(i64 %6, i64 0)
+  %8 = icmp ne i64 %7, 0
+  %9 = select i1 %8, i32 %1, i32 %2
+; CHECK: select{{.*}}, !prof !1
+  ret i32 %9
+}
+
+
 declare i1 @llvm.expect.i1(i1, i1) nounwind readnone
 
 ; CHECK: !0 = !{!"branch_weights", i32 2000, i32 1}
 ; CHECK: !1 = !{!"branch_weights", i32 1, i32 2000}
 ; CHECK: !2 = !{!"branch_weights", i32 1, i32 2000, i32 1}
 ; CHECK: !3 = !{!"branch_weights", i32 2000, i32 1, i32 1}