Propagate branchweight through phi values when ExpectedValue is unlikely in LowerExpectIntrinsic

[LukeZhuang]

Summary

Currently handlePhiDef does not take the probability operand of expect_with_probability into consideration.
When the probability is low, meaning the "expected value" is unlikely, we need to mark the branches from which the phi value is the same as "expected value" as unlikely.

The existing code treats __builtin_expect((a0==1)||(a1==1)||(a2==1), 1) and __builtin_expect_with_probability((a0==1)||(a1==1)||(a2==1), 1, 0) the same, when in fact they are very different: They both "expect" value 1, but the first says it is very likely and the second says it has probability 0, i.e. very unlikely.

An example of IR code we need to handle is:

b0:
  %c0 = icmp eq i8 %a0, 1
  br i1 %c0, label %b3, label %b1
  
b1:
  %c1 = icmp eq i8 %a1, 1
  br i1 %c1, label %b3, label %b2
  
b2:
  %c2 = icmp eq i8 %a2, 1
  br label %b3

b3:
  %cond0 = phi i1 [ true, %b0 ], [ true, %b1 ], [ %c2, %b2 ]
  %cond1 = extend %cond0 to i64
  %expval = call i64 @llvm.expect.with.probability.i64(i64 %cond1, i64 1, double 0.0)

We should also annotate the BranchInst of b0 and b1 because we know that "true" is unlikely. The only path that does not result in an unlikely outcome is through b2.

This helps __builtin_expect(x, 0) and __builtin_expect_with_probability(x, 1, 0) have the same behavior. For reference, GCC does already behave consistently like this.

In addition, we also added an ImmArg attribute to the expect.with.probability intrinsic to disallow non-constant probability arguments.

This issue was introduced by https://reviews.llvm.org/D79830, which did not update handlePhiDef correctly for the expect.with.probability case.

[davidxl]

llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

137

This should probably be made more biased.

[LukeZhuang]

llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

137

Hi David, thanks for the comment! May I ask is "more biased" meant by a higher threshold of considering a branch likely/unlikely? Or could you give an example of which it ought to behave differently than this code(and why)? Thanks!

[davidxl]

never mind. I misunderstood the usage of this variable. There is no need to adjust.

[davidxl]

lgtm

[LukeZhuang]

Hi David, thank you very much! Seems I do not have access to push to the repo, may I ask could you please help me do it? Here is the commit info:

Commit message title:
[LowerExpectIntrinsic] Propagate branch weights through phi values when ExpectedValue is unlikely in LowerExpectIntrinsic

Detailed info:
Update handlePhiDef to consider the probability argument in an expect.with.probability intrinsic when annotating BranchInsts.
In addition, we also disallow non-constant probability arguments in this intrinsic.

Differential Revsion: https://reviews.llvm.org/D140337

[davidxl]

It might be better to ask for commit access. It is needed for potential follow-ups.

[LukeZhuang]

In D140337#4011964, @davidxl wrote:

It might be better to ask for commit access. It is needed for potential follow-ups.

Got it, I have just asked for commit access. Thanks!

llvm/include/llvm/IR/Intrinsics.td

 
 //===------------------------- Expect Intrinsics --------------------------===//
 //
 def int_expect : DefaultAttrsIntrinsic<[llvm_anyint_ty],
   [LLVMMatchType<0>, LLVMMatchType<0>], [IntrNoMem, IntrWillReturn]>;
 
 def int_expect_with_probability : DefaultAttrsIntrinsic<[llvm_anyint_ty],
   [LLVMMatchType<0>, LLVMMatchType<0>, llvm_double_ty],
-  [IntrNoMem, IntrWillReturn]>;
+  [IntrNoMem, IntrWillReturn, ImmArg<ArgIndex<2>>]>;
 
 //===-------------------- Bit Manipulation Intrinsics ---------------------===//
 //
 
 // None of these intrinsics accesses memory at all.
 let IntrProperties = [IntrNoMem, IntrSpeculatable, IntrWillReturn] in {
   def int_bswap: DefaultAttrsIntrinsic<[llvm_anyint_ty], [LLVMMatchType<0>]>;
   def int_ctpop: DefaultAttrsIntrinsic<[llvm_anyint_ty], [LLVMMatchType<0>]>;

llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

 /// edge of the phi is unlikely to be taken. Using that information,
 /// the branch probability info for the originating branch can be inferred.
 static void handlePhiDef(CallInst *Expect) {
   Value &Arg = *Expect->getArgOperand(0);
   ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Expect->getArgOperand(1));
   if (!ExpectedValue)
     return;
   const APInt &ExpectedPhiValue = ExpectedValue->getValue();
+  bool ExpectedValueIsLikely = true;
+  Function *Fn = Expect->getCalledFunction();
+  // If the function is expect_with_probability, then we need to take the
+  // probability into consideration. For example, in
+  // expect.with.probability.i64(i64 %a, i64 1, double 0.0), the
+  // "ExpectedValue" 1 is unlikely. This affects probability propagation later.
+  if (Fn->getIntrinsicID() == Intrinsic::expect_with_probability) {
+    auto *Confidence = cast<ConstantFP>(Expect->getArgOperand(2));
+    double TrueProb = Confidence->getValueAPF().convertToDouble();
+    ExpectedValueIsLikely = (TrueProb > 0.5);
+  }
 
   // Walk up in backward a list of instructions that
   // have 'copy' semantics by 'stripping' the copies
   // until a PHI node or an instruction of unknown kind
   // is reached. Negation via xor is also handled.
   //
   //       C = PHI(...);
   //       B = C;
static void handlePhiDef(CallInst *Expect) {
   for (unsigned i = 0, e = PhiDef->getNumIncomingValues(); i != e; ++i) {
 
     Value *PhiOpnd = PhiDef->getIncomingValue(i);
     ConstantInt *CI = dyn_cast<ConstantInt>(PhiOpnd);
     if (!CI)
       continue;
 
     // Not an interesting case when IsUnlikely is false -- we can not infer
-    // anything useful when the operand value matches the expected phi
-    // output.
-    if (ExpectedPhiValue == ApplyOperations(CI->getValue()))
+    // anything useful when:
+    // (1) We expect some phi output and the operand value matches it, or
+    // (2) We don't expect some phi output (i.e. the "ExpectedValue" has low
+    //     probability) and the operand value doesn't match that.
+    const APInt &CurrentPhiValue = ApplyOperations(CI->getValue());
+    if (ExpectedValueIsLikely == (ExpectedPhiValue == CurrentPhiValue))
       continue;
 
     BranchInst *BI = GetDomConditional(i);
     if (!BI)
       continue;
 
     MDBuilder MDB(PhiDef->getContext());
 
auto IsOpndComingFromSuccessor = [&](BasicBlock *Succ) {
         // to the Phi, this successor is the one feeding this
         // Phi operand.
         return true;
       return false;
     };
     uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
     std::tie(LikelyBranchWeightVal, UnlikelyBranchWeightVal) = getBranchWeight(
         Expect->getCalledFunction()->getIntrinsicID(), Expect, 2);
+    if (!ExpectedValueIsLikely)
+      std::swap(LikelyBranchWeightVal, UnlikelyBranchWeightVal);
 
     if (IsOpndComingFromSuccessor(BI->getSuccessor(1)))
       BI->setMetadata(LLVMContext::MD_prof,
                       MDB.createBranchWeights(LikelyBranchWeightVal,
                                               UnlikelyBranchWeightVal));
     else if (IsOpndComingFromSuccessor(BI->getSuccessor(0)))
       BI->setMetadata(LLVMContext::MD_prof,
                       MDB.createBranchWeights(UnlikelyBranchWeightVal,

llvm/test/Transforms/LowerExpectIntrinsic/phi_unexpect.ll

+; RUN: opt -S -passes='function(lower-expect),strip-dead-prototypes' -likely-branch-weight=2147483647 -unlikely-branch-weight=1 < %s | FileCheck %s
+
+; The C case
+; if (__builtin_expect_with_probability(((a0 == 1) || (a1 == 1) || (a2 == 1)), 1, 0))
+; For the above case, all 3 branches should be annotated
+; which should be equivalent to if (__builtin_expect(((a0 == 1) || (a1 == 1) || (a2 == 1)), 0))
+
+; The C case
+; if (__builtin_expect_with_probability(((a0 == 1) || (a1 == 1) || (a2 == 1)), 1, 1))
+; For the above case, we do not have enough information, so only the last branch could be annotated
+; which should be equivalent to if (__builtin_expect(((a0 == 1) || (a1 == 1) || (a2 == 1)), 1))
+
+declare void @foo()
+
+declare i64 @llvm.expect.i64(i64, i64) nounwind readnone
+declare i64 @llvm.expect.with.probability.i64(i64, i64, double) nounwind readnone
+
+; CHECK-LABEL: @test1_expect_1(
+; CHECK: block0:
+; CHECK-NOT: prof
+; CHECK: block1:
+; CHECK-NOT: prof
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !0
+define void @test1_expect_1(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.i64(i64 %cond2, i64 1)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; should have exactly the same behavior as test1
+; CHECK-LABEL: @test2_expect_with_prob_1_1(
+; CHECK: block0:
+; CHECK-NOT: prof
+; CHECK: block1:
+; CHECK-NOT: prof
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !0
+define void @test2_expect_with_prob_1_1(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.with.probability.i64(i64 %cond2, i64 1, double 1.0)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; should have exactly the same behavior as test1
+; CHECK-LABEL: @test3_expect_with_prob_0_0(
+; CHECK: block0:
+; CHECK-NOT: prof
+; CHECK: block1:
+; CHECK-NOT: prof
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !0
+define void @test3_expect_with_prob_0_0(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.with.probability.i64(i64 %cond2, i64 0, double 0.0)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; CHECK-LABEL: @test4_expect_0(
+; CHECK: block0:
+; CHECK: br i1 %c0, label %block3, label %block1, !prof !1
+; CHECK: block1:
+; CHECK: br i1 %c1, label %block3, label %block2, !prof !1
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !1
+define void @test4_expect_0(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.i64(i64 %cond2, i64 0)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; should have exactly the same behavior as test4
+; CHECK-LABEL: @test5_expect_with_prob_1_0(
+; CHECK: block0:
+; CHECK: br i1 %c0, label %block3, label %block1, !prof !1
+; CHECK: block1:
+; CHECK: br i1 %c1, label %block3, label %block2, !prof !1
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !1
+define void @test5_expect_with_prob_1_0(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.with.probability.i64(i64 %cond2, i64 1, double 0.0)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; should have exactly the same behavior as test4
+; CHECK-LABEL: @test6_expect_with_prob_0_1(
+; CHECK: block0:
+; CHECK: br i1 %c0, label %block3, label %block1, !prof !1
+; CHECK: block1:
+; CHECK: br i1 %c1, label %block3, label %block2, !prof !1
+; CHECK: block3:
+; CHECK: br i1 %tobool, label %block4, label %block5, !prof !1
+define void @test6_expect_with_prob_0_1(i8 %a0, i8 %a1, i8 %a2) {
+block0:
+  %c0 = icmp eq i8 %a0, 1
+  br i1 %c0, label %block3, label %block1
+
+block1:
+  %c1 = icmp eq i8 %a1, 1
+  br i1 %c1, label %block3, label %block2
+
+block2:
+  %c2 = icmp eq i8 %a2, 1
+  br label %block3
+
+block3:
+  %cond0 = phi i1 [ true, %block0 ], [ true, %block1 ], [ %c2, %block2 ]
+  %cond1 = zext i1 %cond0 to i32
+  %cond2 = sext i32 %cond1 to i64
+  %expval = call i64 @llvm.expect.with.probability.i64(i64 %cond2, i64 0, double 1.0)
+  %tobool = icmp ne i64 %expval, 0
+  br i1 %tobool, label %block4, label %block5
+
+block4:
+  call void @foo()
+  br label %block5
+
+block5:
+  ret void
+}
+
+; CHECK: !0 = !{!"branch_weights", i32 2147483647, i32 1}
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 2147483647}