[SelectionDAG] Add generic implementation for @llvm.expect.with.probability when optimizations are disabled

[antangelo]

Handle @llvm.expect.with.probability in SelectionDAGBuilder, FastISel, and IntrinsicLowering in the same way @llvm.expect is handled, where the value is passed through as-is. This can be reached if the intrinsic is used without optimizations, where it would otherwise be properly transformed out.

Fixes #115411 for SelectionDAG. A similar patch is likely needed for GlobalISel.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

Author: None (antangelo)

Changes

Handle @llvm.expect.with.probability in SelectionDAGBuilder, FastISel, and IntrinsicLowering in the same way @llvm.expect is handled, where the value is passed through as-is. This can be reached if the intrinsic is used without optimizations, where it would otherwise be properly transformed out.

Fixes #115411 for SelectionDAG. A similar patch is likely needed for GlobalISel.

---

Full diff: https://github.com/llvm/llvm-project/pull/117459.diff

4 Files Affected:

• (modified) llvm/lib/CodeGen/IntrinsicLowering.cpp (+4-2)
• (modified) llvm/lib/CodeGen/SelectionDAG/FastISel.cpp (+2-1)
• (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp (+3-1)
• (added) llvm/test/CodeGen/Generic/builtin-expect-with-probability.ll (+8)

diff --git a/llvm/lib/CodeGen/IntrinsicLowering.cpp b/llvm/lib/CodeGen/IntrinsicLowering.cpp
index f799a8cfc1ba7e..1518ead7698bee 100644
--- a/llvm/lib/CodeGen/IntrinsicLowering.cpp
+++ b/llvm/lib/CodeGen/IntrinsicLowering.cpp
@@ -236,8 +236,10 @@ void IntrinsicLowering::LowerIntrinsicCall(CallInst *CI) {
     report_fatal_error("Code generator does not support intrinsic function '"+
                       Callee->getName()+"'!");
 
-  case Intrinsic::expect: {
-    // Just replace __builtin_expect(exp, c) with EXP.
+  case Intrinsic::expect:
+  case Intrinsic::expect_with_probability: {
+    // Just replace __builtin_expect(exp, c) and
+    // __builtin_expect_with_probability(exp, c, p) with EXP.
     Value *V = CI->getArgOperand(0);
     CI->replaceAllUsesWith(V);
     break;
diff --git a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
index b431ecc8472671..eede879e7e80d1 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
@@ -1456,7 +1456,8 @@ bool FastISel::selectIntrinsicCall(const IntrinsicInst *II) {
 
   case Intrinsic::launder_invariant_group:
   case Intrinsic::strip_invariant_group:
-  case Intrinsic::expect: {
+  case Intrinsic::expect:
+  case Intrinsic::expect_with_probability: {
     Register ResultReg = getRegForValue(II->getArgOperand(0));
     if (!ResultReg)
       return false;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 9d729d448502d8..a38a3e9b91052d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -7532,7 +7532,9 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     return;
 
   case Intrinsic::expect:
-    // Just replace __builtin_expect(exp, c) with EXP.
+  case Intrinsic::expect_with_probability:
+    // Just replace __builtin_expect(exp, c) and
+    // __builtin_expect_with_probability(exp, c, p) with EXP.
     setValue(&I, getValue(I.getArgOperand(0)));
     return;
 
diff --git a/llvm/test/CodeGen/Generic/builtin-expect-with-probability.ll b/llvm/test/CodeGen/Generic/builtin-expect-with-probability.ll
new file mode 100644
index 00000000000000..aef134b636d5a7
--- /dev/null
+++ b/llvm/test/CodeGen/Generic/builtin-expect-with-probability.ll
@@ -0,0 +1,8 @@
+; RUN: llc < %s
+
+declare i32 @llvm.expect.with.probability(i32, i32, double)
+
+define i32 @test1(i32 %val) nounwind {
+    %expected = call i32 @llvm.expect.with.probability(i32 %val, i32 1, double 0.5)
+    ret i32 %expected
+}