[llvm] b49ee01 - [LowerExpectIntrinsic] Propagate branch weights through phi values when ExpectedValue is unlikely in LowerExpectIntrinsic

Zhi Zhuang via llvm-commits llvm-commits at lists.llvm.org
Thu Dec 22 14:35:05 PST 2022


Author: Zhi Zhuang
Date: 2022-12-22T17:33:52-05:00
New Revision: b49ee01fe288abda9269114cc3aef6caade6015a

URL: https://github.com/llvm/llvm-project/commit/b49ee01fe288abda9269114cc3aef6caade6015a
DIFF: https://github.com/llvm/llvm-project/commit/b49ee01fe288abda9269114cc3aef6caade6015a.diff

LOG: [LowerExpectIntrinsic] Propagate branch weights through phi values when ExpectedValue is unlikely in LowerExpectIntrinsic

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

Added: 
    llvm/test/Transforms/LowerExpectIntrinsic/phi_unexpect.ll

Modified: 
    llvm/include/llvm/IR/Intrinsics.td
    llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index 7bd27674a80b6..41dd958e068d7 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -959,7 +959,7 @@ def int_expect : DefaultAttrsIntrinsic<[llvm_anyint_ty],
 
 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 ---------------------===//
 //

diff  --git a/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp b/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
index a1c2c3cd7275c..454aa56be5319 100644
--- a/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
@@ -125,6 +125,17 @@ static void handlePhiDef(CallInst *Expect) {
   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
@@ -213,9 +224,12 @@ static void handlePhiDef(CallInst *Expect) {
       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);
@@ -248,6 +262,8 @@ static void handlePhiDef(CallInst *Expect) {
     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,

diff  --git a/llvm/test/Transforms/LowerExpectIntrinsic/phi_unexpect.ll b/llvm/test/Transforms/LowerExpectIntrinsic/phi_unexpect.ll
new file mode 100644
index 0000000000000..2bad66343b761
--- /dev/null
+++ b/llvm/test/Transforms/LowerExpectIntrinsic/phi_unexpect.ll
@@ -0,0 +1,239 @@
+; 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}


        


More information about the llvm-commits mailing list