[llvm] [SelectionDAG] Add generic implementation for @llvm.expect.with.probability when optimizations are disabled (PR #117459)

[via llvm-commits]

https://github.com/antangelo updated https://github.com/llvm/llvm-project/pull/117459

>From a8690d68cf77b56a3001831883c301b8b3f89637 Mon Sep 17 00:00:00 2001
From: antangelo <contact at antangelo.com>
Date: Thu, 21 Nov 2024 01:35:47 -0500
Subject: [PATCH] [SelectionDAG] Add generic implementation for
 @llvm.expect.with.probability when optimizations are disabled

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.
---
 llvm/lib/CodeGen/IntrinsicLowering.cpp                    | 6 ++++--
 llvm/lib/CodeGen/SelectionDAG/FastISel.cpp                | 3 ++-
 llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp     | 4 +++-
 .../CodeGen/Generic/builtin-expect-with-probability.ll    | 8 ++++++++
 4 files changed, 17 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/CodeGen/Generic/builtin-expect-with-probability.ll

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
+}