LukeZhuang created this revision.
LukeZhuang added reviewers: RKSimon, phosek, gribozavr, davidxl, chandlerc, 
pete, jyknight, dblaikie, kuba.
LukeZhuang added projects: LLVM, clang.
Herald added subscribers: llvm-commits, cfe-commits, jdoerfert, hiraditya.

Add a new builtin-function `__builtin_expect_with_probability` and intrinsic 
`llvm.expect.with.probability`.
The interface is `__builtin_expect_with_probability(long expr, long expected, 
double probability)`.
It is mainly the same as `__builtin_expect` besides one more argument 
indicating the probability of expression equal to expected value. The 
probability should be a constant floating-point expression and be in range 
[0.0, 1.0] inclusive.
It is similar to builtin-expect-with-probability function in GCC built-in 
functions <https://gcc.gnu.org/onlinedocs/gcc/Other-Builtins.html>.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D79830

Files:
  clang/include/clang/Basic/Builtins.def
  clang/lib/CodeGen/CGBuiltin.cpp
  clang/test/CodeGen/builtin-expect-with-probability-switch.c
  clang/test/CodeGen/builtin-expect-with-probability.c
  llvm/include/llvm/IR/Intrinsics.td
  llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

Index: clang/test/CodeGen/builtin-expect-with-probability.c
===================================================================
--- /dev/null
+++ clang/test/CodeGen/builtin-expect-with-probability.c
@@ -0,0 +1,11 @@
+// RUN: %clang_cc1 -emit-llvm -o - %s -O1 | FileCheck %s
+
+int expect_taken(int x) {
+// CHECK: !{{[0-9]+}} = !{!"branch_weights", i32 1932735283, i32 214748366}
+
+	if (__builtin_expect_with_probability (x == 100, 1, 0.9)) {
+		return 0;
+	}
+	return x;
+}
+
Index: clang/test/CodeGen/builtin-expect-with-probability-switch.c
===================================================================
--- /dev/null
+++ clang/test/CodeGen/builtin-expect-with-probability-switch.c
@@ -0,0 +1,19 @@
+// RUN: %clang_cc1 -emit-llvm -o - %s -O1 | FileCheck %s
+
+int expect_taken(int x) {
+// CHECK: !{{[0-9]+}} = !{!"branch_weights", i32 107374184, i32 107374184, i32 1717986918, i32 107374184, i32 107374184}
+	switch(__builtin_expect_with_probability(x, 1, 0.8)) {
+		case 0:
+			x = x + 0;
+		case 1:
+			x = x + 1;
+		case 2:
+			x = x + 2;
+		case 5:
+			x = x + 5;
+		default:
+			x = x + 6;
+	}
+	return x;
+}
+
Index: clang/lib/CodeGen/CGBuiltin.cpp
===================================================================
--- clang/lib/CodeGen/CGBuiltin.cpp
+++ clang/lib/CodeGen/CGBuiltin.cpp
@@ -2039,6 +2039,37 @@
         Builder.CreateCall(FnExpect, {ArgValue, ExpectedValue}, "expval");
     return RValue::get(Result);
   }
+  case Builtin::BI__builtin_expect_with_probability: {
+    Value *ArgValue = EmitScalarExpr(E->getArg(0));
+    llvm::Type *ArgType = ArgValue->getType();
+
+    Value *ExpectedValue = EmitScalarExpr(E->getArg(1));
+    Value *Confidence = EmitScalarExpr(E->getArg(2));
+    ConstantFP *Confidence_f = dyn_cast<ConstantFP>(Confidence);
+    if (!Confidence_f) {
+      CGM.Error(E->getArg(2)->getLocStart(),
+                "probability of __builtin_expect_with_probability must be "
+                "constant floating-point expression");
+    } else {
+      double prob = Confidence_f->getValueAPF().convertToDouble();
+      if (prob < 0.0 || prob > 1.0) {
+        CGM.Error(E->getArg(2)->getLocStart(),
+                  "probability of __builtin_expect_with_probability is "
+                  "outside the range [0.0, 1.0]");
+      }
+    }
+    // Don't generate llvm.expect.with.probability on -O0 as the backend
+    // won't use it for anything.
+    // Note, we still IRGen ExpectedValue because it could have side-effects.
+    if (CGM.getCodeGenOpts().OptimizationLevel == 0)
+      return RValue::get(ArgValue);
+
+    Value *FnExpect =
+        CGM.getIntrinsic(Intrinsic::expect_with_probability, ArgType);
+    Value *Result = Builder.CreateCall(
+        FnExpect, {ArgValue, ExpectedValue, Confidence}, "expval");
+    return RValue::get(Result);
+  }
   case Builtin::BI__builtin_assume_aligned: {
     const Expr *Ptr = E->getArg(0);
     Value *PtrValue = EmitScalarExpr(Ptr);
Index: clang/include/clang/Basic/Builtins.def
===================================================================
--- clang/include/clang/Basic/Builtins.def
+++ clang/include/clang/Basic/Builtins.def
@@ -562,6 +562,7 @@
 
 BUILTIN(__builtin_unpredictable, "LiLi"   , "nc")
 BUILTIN(__builtin_expect, "LiLiLi"   , "nc")
+BUILTIN(__builtin_expect_with_probability, "LiLiLid"   , "nc")
 BUILTIN(__builtin_prefetch, "vvC*.", "nc")
 BUILTIN(__builtin_readcyclecounter, "ULLi", "n")
 BUILTIN(__builtin_trap, "v", "nr")
Index: llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
+++ llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
@@ -54,13 +54,32 @@
     "unlikely-branch-weight", cl::Hidden, cl::init(1),
     cl::desc("Weight of the branch unlikely to be taken (default = 1)"));
 
+std::pair<uint32_t, uint32_t> setBranchWeight(Intrinsic::ID IntrinsicID,
+                                              CallInst *CI, int BranchCount) {
+  if (IntrinsicID == Intrinsic::expect) {
+    // __builtin_expect
+    return { LikelyBranchWeight, UnlikelyBranchWeight };
+  } else {
+    // __builtin_expect_with_probability
+    assert(CI->getNumOperands() >= 3 &&
+           "expect with probability must have 3 arguments");
+    ConstantFP *Confidence = dyn_cast<ConstantFP>(CI->getArgOperand(2));
+    double TrueProb = Confidence->getValueAPF().convertToDouble();
+    double FalseProb = (1.0 - TrueProb) / (BranchCount - 1);
+    uint32_t LikelyBW = ceil((TrueProb * (double)(INT32_MAX - 1)) + 1.0);
+    uint32_t UnlikelyBW = ceil((FalseProb * (double)(INT32_MAX - 1)) + 1.0);
+    return { LikelyBW, UnlikelyBW };
+  }
+}
+
 static bool handleSwitchExpect(SwitchInst &SI) {
   CallInst *CI = dyn_cast<CallInst>(SI.getCondition());
   if (!CI)
     return false;
 
   Function *Fn = CI->getCalledFunction();
-  if (!Fn || Fn->getIntrinsicID() != Intrinsic::expect)
+  if (!Fn || (Fn->getIntrinsicID() != Intrinsic::expect &&
+              Fn->getIntrinsicID() != Intrinsic::expect_with_probability))
     return false;
 
   Value *ArgValue = CI->getArgOperand(0);
@@ -70,15 +89,19 @@
 
   SwitchInst::CaseHandle Case = *SI.findCaseValue(ExpectedValue);
   unsigned n = SI.getNumCases(); // +1 for default case.
-  SmallVector<uint32_t, 16> Weights(n + 1, UnlikelyBranchWeight);
+  std::pair<uint32_t, uint32_t> WeightNums =
+      setBranchWeight(Fn->getIntrinsicID(), CI, n + 1);
+  uint32_t LikelyBranchWeightVal = WeightNums.first;
+  uint32_t UnlikelyBranchWeightVal = WeightNums.second;
+
+  SmallVector<uint32_t, 16> Weights(n + 1, UnlikelyBranchWeightVal);
 
   uint64_t Index = (Case == *SI.case_default()) ? 0 : Case.getCaseIndex() + 1;
-  Weights[Index] = LikelyBranchWeight;
+  Weights[Index] = LikelyBranchWeightVal;
 
-  SI.setMetadata(
-      LLVMContext::MD_misexpect,
-      MDBuilder(CI->getContext())
-          .createMisExpect(Index, LikelyBranchWeight, UnlikelyBranchWeight));
+  SI.setMetadata(LLVMContext::MD_misexpect,
+                 MDBuilder(CI->getContext()).createMisExpect(
+                     Index, LikelyBranchWeightVal, UnlikelyBranchWeightVal));
 
   SI.setCondition(ArgValue);
   misexpect::checkFrontendInstrumentation(SI);
@@ -222,15 +245,19 @@
         return true;
       return false;
     };
+    std::pair<uint32_t, uint32_t> WeightNums = setBranchWeight(
+        Expect->getCalledFunction()->getIntrinsicID(), Expect, 2);
+    uint32_t LikelyBranchWeightVal = WeightNums.first;
+    uint32_t UnlikelyBranchWeightVal = WeightNums.second;
 
     if (IsOpndComingFromSuccessor(BI->getSuccessor(1)))
-      BI->setMetadata(
-          LLVMContext::MD_prof,
-          MDB.createBranchWeights(LikelyBranchWeight, UnlikelyBranchWeight));
+      BI->setMetadata(LLVMContext::MD_prof,
+                      MDB.createBranchWeights(LikelyBranchWeightVal,
+                                              UnlikelyBranchWeightVal));
     else if (IsOpndComingFromSuccessor(BI->getSuccessor(0)))
-      BI->setMetadata(
-          LLVMContext::MD_prof,
-          MDB.createBranchWeights(UnlikelyBranchWeight, LikelyBranchWeight));
+      BI->setMetadata(LLVMContext::MD_prof,
+                      MDB.createBranchWeights(UnlikelyBranchWeightVal,
+                                              LikelyBranchWeightVal));
   }
 }
 
@@ -276,7 +303,8 @@
   }
 
   Function *Fn = CI->getCalledFunction();
-  if (!Fn || Fn->getIntrinsicID() != Intrinsic::expect)
+  if (!Fn || (Fn->getIntrinsicID() != Intrinsic::expect &&
+              Fn->getIntrinsicID() != Intrinsic::expect_with_probability))
     return false;
 
   Value *ArgValue = CI->getArgOperand(0);
@@ -288,13 +316,22 @@
   MDNode *Node;
   MDNode *ExpNode;
 
+  std::pair<uint32_t, uint32_t> WeightNums =
+      setBranchWeight(Fn->getIntrinsicID(), CI, 2);
+  uint32_t LikelyBranchWeightVal = WeightNums.first;
+  uint32_t UnlikelyBranchWeightVal = WeightNums.second;
+
   if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
       (Predicate == CmpInst::ICMP_EQ)) {
-    Node = MDB.createBranchWeights(LikelyBranchWeight, UnlikelyBranchWeight);
-    ExpNode = MDB.createMisExpect(0, LikelyBranchWeight, UnlikelyBranchWeight);
+    Node =
+        MDB.createBranchWeights(LikelyBranchWeightVal, UnlikelyBranchWeightVal);
+    ExpNode =
+        MDB.createMisExpect(0, LikelyBranchWeightVal, UnlikelyBranchWeightVal);
   } else {
-    Node = MDB.createBranchWeights(UnlikelyBranchWeight, LikelyBranchWeight);
-    ExpNode = MDB.createMisExpect(1, LikelyBranchWeight, UnlikelyBranchWeight);
+    Node =
+        MDB.createBranchWeights(UnlikelyBranchWeightVal, LikelyBranchWeightVal);
+    ExpNode =
+        MDB.createMisExpect(1, LikelyBranchWeightVal, UnlikelyBranchWeightVal);
   }
 
   BSI.setMetadata(LLVMContext::MD_misexpect, ExpNode);
@@ -346,7 +383,8 @@
       }
 
       Function *Fn = CI->getCalledFunction();
-      if (Fn && Fn->getIntrinsicID() == Intrinsic::expect) {
+      if (Fn && (Fn->getIntrinsicID() == Intrinsic::expect ||
+                 Fn->getIntrinsicID() == Intrinsic::expect_with_probability)) {
         // Before erasing the llvm.expect, walk backward to find
         // phi that define llvm.expect's first arg, and
         // infer branch probability:
Index: llvm/include/llvm/IR/Intrinsics.td
===================================================================
--- llvm/include/llvm/IR/Intrinsics.td
+++ llvm/include/llvm/IR/Intrinsics.td
@@ -752,6 +752,10 @@
 def int_expect : Intrinsic<[llvm_anyint_ty],
   [LLVMMatchType<0>, LLVMMatchType<0>], [IntrNoMem, IntrWillReturn]>;
 
+def int_expect_with_probability : Intrinsic<[llvm_anyint_ty],
+  [LLVMMatchType<0>, LLVMMatchType<0>, llvm_double_ty],
+  [IntrNoMem, IntrWillReturn]>;
+
 //===-------------------- Bit Manipulation Intrinsics ---------------------===//
 //
 
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to