[Profile] Fix builtin_expect lowering bug

[davidxl]

Summary

There is an obvious bug in the lowering code when non-boolean integer value is expected. The existing code wrongly assumes the NE comparison's constant operand is always zero and also skips EQ predicate. This either produces wrong profile annotation or misses annotation.

Other predicate kinds should also be supported (separately).

[danielcdh]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

134 ↗

(On Diff #100963)

Maybe inline the ExpectedValue->getZExtValue() so that there is no need to have ExpectedIntValue.

llvm/trunk/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
   //
   // Or the following simpler case:
   //   %expval = call i1 @llvm.expect.i1(i1 %cmp, i1 1)
   //   br i1 %expval, label %if.then, label %if.end
 
   CallInst *CI;
 
   ICmpInst *CmpI = dyn_cast<ICmpInst>(BSI.getCondition());
+  CmpInst::Predicate Predicate;
+  uint64_t ValueComparedTo = 0;
   if (!CmpI) {
     CI = dyn_cast<CallInst>(BSI.getCondition());
+    Predicate = CmpInst::ICMP_NE;
+    ValueComparedTo = 0;
   } else {
-    if (CmpI->getPredicate() != CmpInst::ICMP_NE)
+    Predicate = CmpI->getPredicate();
+    if (Predicate != CmpInst::ICMP_NE && Predicate != CmpInst::ICMP_EQ)
       return false;
+    ConstantInt *CmpConstOperand = dyn_cast<ConstantInt>(CmpI->getOperand(1));
+    if (!CmpConstOperand)
+      return false;
+    ValueComparedTo = CmpConstOperand->getZExtValue();
     CI = dyn_cast<CallInst>(CmpI->getOperand(0));
   }
 
   if (!CI)
     return false;
 
   Function *Fn = CI->getCalledFunction();
   if (!Fn || Fn->getIntrinsicID() != Intrinsic::expect)
     return false;
 
   Value *ArgValue = CI->getArgOperand(0);
   ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(CI->getArgOperand(1));
   if (!ExpectedValue)
     return false;
 
   MDBuilder MDB(CI->getContext());
   MDNode *Node;
 
-  // If expect value is equal to 1 it means that we are more likely to take
-  // branch 0, in other case more likely is branch 1.
-  if (ExpectedValue->isOne())
+  if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
+      (Predicate == CmpInst::ICMP_EQ))
     Node = MDB.createBranchWeights(LikelyBranchWeight, UnlikelyBranchWeight);
   else
     Node = MDB.createBranchWeights(UnlikelyBranchWeight, LikelyBranchWeight);
 
   BSI.setMetadata(LLVMContext::MD_prof, Node);
 
   if (CmpI)
     CmpI->setOperand(0, ArgValue);

llvm/trunk/test/Transforms/LowerExpectIntrinsic/expect_nonboolean.ll

+; RUN: opt -lower-expect  -S -o - < %s | FileCheck %s
+; RUN: opt -S -passes='function(lower-expect)' < %s | FileCheck %s
+
+define i32 @foo(i32 %arg) #0 {
+; CHECK-LABEL: @foo(i32{{.*}})
+bb:
+  %tmp = sext i32 %arg to i64
+  %tmp1 = call i64 @llvm.expect.i64(i64 %tmp, i64 4)
+  %tmp2 = icmp ne i64 %tmp1, 0
+  br i1 %tmp2, label %bb3, label %bb5
+; CHECK: br i1 %tmp2{{.*}}!prof [[LIKELY:![0-9]+]]
+
+bb3:                                              ; preds = %bb
+  %tmp4 = call i32 (...) @bar()
+  br label %bb5
+
+bb5:                                              ; preds = %bb3, %bb
+  ret i32 1
+}
+
+define i32 @foo2(i32 %arg) #0 {
+; CHECK-LABEL: @foo2
+bb:
+  %tmp = sext i32 %arg to i64
+  %tmp1 = call i64 @llvm.expect.i64(i64 %tmp, i64 4)
+  %tmp2 = icmp eq i64 %tmp1, 2
+  br i1 %tmp2, label %bb3, label %bb5
+; CHECK: br i1 %tmp2{{.*}}!prof [[UNLIKELY:![0-9]+]]
+
+bb3:                                              ; preds = %bb
+  %tmp4 = call i32 (...) @bar()
+  br label %bb5
+
+bb5:                                              ; preds = %bb3, %bb
+  ret i32 1
+}
+
+define i32 @foo3(i32 %arg) #0 {
+; CHECK-LABEL: @foo3
+bb:
+  %tmp = sext i32 %arg to i64
+  %tmp1 = call i64 @llvm.expect.i64(i64 %tmp, i64 4)
+  %tmp2 = icmp eq i64 %tmp1, 4
+  br i1 %tmp2, label %bb3, label %bb5
+; CHECK: br i1 %tmp2{{.*}}!prof [[LIKELY]]
+
+bb3:                                              ; preds = %bb
+  %tmp4 = call i32 (...) @bar()
+  br label %bb5
+
+bb5:                                              ; preds = %bb3, %bb
+  ret i32 1
+}
+
+define i32 @foo4(i32 %arg) #0 {
+; CHECK-LABEL: @foo4
+bb:
+  %tmp = sext i32 %arg to i64
+  %tmp1 = call i64 @llvm.expect.i64(i64 %tmp, i64 4)
+  %tmp2 = icmp ne i64 %tmp1, 2
+  br i1 %tmp2, label %bb3, label %bb5
+; CHECK: br i1 %tmp2{{.*}}!prof [[LIKELY]]
+
+bb3:                                              ; preds = %bb
+  %tmp4 = call i32 (...) @bar()
+  br label %bb5
+
+bb5:                                              ; preds = %bb3, %bb
+  ret i32 1
+}
+
+define i32 @foo5(i32 %arg, i32 %arg1) #0 {
+; CHECK-LABEL: @foo5
+bb:
+  %tmp = sext i32 %arg1 to i64
+  %tmp2 = call i64 @llvm.expect.i64(i64 %tmp, i64 4)
+  %tmp3 = sext i32 %arg to i64
+  %tmp4 = icmp ne i64 %tmp2, %tmp3
+  br i1 %tmp4, label %bb5, label %bb7
+; CHECK-NOT: !prof
+
+bb5:                                              ; preds = %bb
+  %tmp6 = call i32 (...) @bar()
+  br label %bb7
+
+bb7:                                              ; preds = %bb5, %bb
+  ret i32 1
+}
+
+declare i64 @llvm.expect.i64(i64, i64) #1
+
+declare i32 @bar(...) local_unnamed_addr #0
+
+attributes #0 = { nounwind uwtable }
+attributes #1 = { nounwind readnone }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 5.0.0 (trunk 304373)"}
+; CHECK: [[LIKELY]] = !{!"branch_weights", i32 2000, i32 1}
+; CHECK: [[UNLIKELY]] = !{!"branch_weights", i32 1, i32 2000}
+