Auto merge of #128250 - Amanieu:select_unpredictable, r=nikic

Add `select_unpredictable` to force LLVM to use CMOV

Since https://reviews.llvm.org/D118118, LLVM will no longer turn CMOVs into branches if it comes from a `select` marked with an `unpredictable` metadata attribute.

This PR introduces `core::intrinsics::select_unpredictable` which emits such a `select` and uses it in the implementation of `binary_search_by`.

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1355,6 +1355,16 @@ fn noundef_metadata(&mut self, load: &'ll Value) {
         }
     }
 
+    pub fn set_unpredictable(&mut self, inst: &'ll Value) {
+        unsafe {
+            llvm::LLVMSetMetadata(
+                inst,
+                llvm::MD_unpredictable as c_uint,
+                llvm::LLVMMDNodeInContext(self.cx.llcx, ptr::null(), 0),
+            );
+        }
+    }
+
     pub fn minnum(&mut self, lhs: &'ll Value, rhs: &'ll Value) -> &'ll Value {
         unsafe { llvm::LLVMRustBuildMinNum(self.llbuilder, lhs, rhs) }
     }

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -4,7 +4,7 @@
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
 use rustc_codegen_ssa::errors::{ExpectedPointerMutability, InvalidMonomorphization};
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
-use rustc_codegen_ssa::mir::place::PlaceRef;
+use rustc_codegen_ssa::mir::place::{PlaceRef, PlaceValue};
 use rustc_codegen_ssa::traits::*;
 use rustc_hir as hir;
 use rustc_middle::mir::BinOp;
@@ -203,6 +203,35 @@ fn codegen_intrinsic_call(
             }
             sym::unlikely => self
                 .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]),
+            sym::select_unpredictable => {
+                let cond = args[0].immediate();
+                assert_eq!(args[1].layout, args[2].layout);
+                let select = |bx: &mut Self, true_val, false_val| {
+                    let result = bx.select(cond, true_val, false_val);
+                    bx.set_unpredictable(&result);
+                    result
+                };
+                match (args[1].val, args[2].val) {
+                    (OperandValue::Ref(true_val), OperandValue::Ref(false_val)) => {
+                        assert!(true_val.llextra.is_none());
+                        assert!(false_val.llextra.is_none());
+                        assert_eq!(true_val.align, false_val.align);
+                        let ptr = select(self, true_val.llval, false_val.llval);
+                        let selected =
+                            OperandValue::Ref(PlaceValue::new_sized(ptr, true_val.align));
+                        selected.store(self, result);
+                        return Ok(());
+                    }
+                    (OperandValue::Immediate(_), OperandValue::Immediate(_))
+                    | (OperandValue::Pair(_, _), OperandValue::Pair(_, _)) => {
+                        let true_val = args[1].immediate_or_packed_pair(self);
+                        let false_val = args[2].immediate_or_packed_pair(self);
+                        select(self, true_val, false_val)
+                    }
+                    (OperandValue::ZeroSized, OperandValue::ZeroSized) => return Ok(()),
+                    _ => span_bug!(span, "Incompatible OperandValue for select_unpredictable"),
+                }
+            }
             sym::catch_unwind => {
                 catch_unwind_intrinsic(
                     self,

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -426,6 +426,7 @@ pub enum MetadataType {
     MD_nontemporal = 9,
     MD_mem_parallel_loop_access = 10,
     MD_nonnull = 11,
+    MD_unpredictable = 15,
     MD_align = 17,
     MD_type = 19,
     MD_vcall_visibility = 28,

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -120,6 +120,7 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -
         | sym::type_id
         | sym::likely
         | sym::unlikely
+        | sym::select_unpredictable
         | sym::ptr_guaranteed_cmp
         | sym::minnumf16
         | sym::minnumf32
@@ -488,6 +489,7 @@ pub fn check_intrinsic_type(
             sym::assume => (0, 0, vec![tcx.types.bool], tcx.types.unit),
             sym::likely => (0, 0, vec![tcx.types.bool], tcx.types.bool),
             sym::unlikely => (0, 0, vec![tcx.types.bool], tcx.types.bool),
+            sym::select_unpredictable => (1, 0, vec![tcx.types.bool, param(0), param(0)], param(0)),
 
             sym::read_via_copy => (1, 0, vec![Ty::new_imm_ptr(tcx, param(0))], param(0)),
             sym::write_via_move => {

compiler/rustc_span/src/symbol.rs

@@ -1701,6 +1701,7 @@
         saturating_add,
         saturating_div,
         saturating_sub,
+        select_unpredictable,
         self_in_typedefs,
         self_struct_ctor,
         semitransparent,

library/core/src/intrinsics.rs

@@ -1008,6 +1008,34 @@ pub const fn unlikely(b: bool) -> bool {
     b
 }
 
+/// Returns either `true_val` or `false_val` depending on condition `b` with a
+/// hint to the compiler that this condition is unlikely to be correctly
+/// predicted by a CPU's branch predictor (e.g. a binary search).
+///
+/// This is otherwise functionally equivalent to `if b { true_val } else { false_val }`.
+///
+/// Note that, unlike most intrinsics, this is safe to call;
+/// it does not require an `unsafe` block.
+/// Therefore, implementations must not require the user to uphold
+/// any safety invariants.
+///
+/// This intrinsic does not have a stable counterpart.
+#[cfg(not(bootstrap))]
+#[unstable(feature = "core_intrinsics", issue = "none")]
+#[rustc_intrinsic]
+#[rustc_nounwind]
+#[miri::intrinsic_fallback_is_spec]
+#[inline]
+pub fn select_unpredictable<T>(b: bool, true_val: T, false_val: T) -> T {
+    if b { true_val } else { false_val }
+}
+
+#[cfg(bootstrap)]
+#[inline]
+pub fn select_unpredictable<T>(b: bool, true_val: T, false_val: T) -> T {
+    if b { true_val } else { false_val }
+}
+
 extern "rust-intrinsic" {
     /// Executes a breakpoint trap, for inspection by a debugger.
     ///

tests/codegen/intrinsics/select_unpredictable.rs

@@ -0,0 +1,35 @@
+//@ compile-flags: -O
+
+#![feature(core_intrinsics)]
+#![crate_type = "lib"]
+
+#[no_mangle]
+pub fn test_int(p: bool, a: u64, b: u64) -> u64 {
+    // CHECK-LABEL: define{{.*}} @test_int
+    // CHECK: select i1 %p, i64 %a, i64 %b, !unpredictable
+    core::intrinsics::select_unpredictable(p, a, b)
+}
+
+#[no_mangle]
+pub fn test_pair(p: bool, a: (u64, u64), b: (u64, u64)) -> (u64, u64) {
+    // CHECK-LABEL: define{{.*}} @test_pair
+    // CHECK: select i1 %p, {{.*}}, !unpredictable
+    core::intrinsics::select_unpredictable(p, a, b)
+}
+
+struct Large {
+    e: [u64; 100],
+}
+
+#[no_mangle]
+pub fn test_struct(p: bool, a: Large, b: Large) -> Large {
+    // CHECK-LABEL: define{{.*}} @test_struct
+    // CHECK: select i1 %p, {{.*}}, !unpredictable
+    core::intrinsics::select_unpredictable(p, a, b)
+}
+
+#[no_mangle]
+pub fn test_zst(p: bool, a: (), b: ()) -> () {
+    // CHECK-LABEL: define{{.*}} @test_zst
+    core::intrinsics::select_unpredictable(p, a, b)
+}