From 4f78f9fbb05145d437540181fda9bcc83d5a53e4 Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Fri, 26 Jul 2024 19:36:21 +0100 Subject: [PATCH] Force LLVM to use CMOV for binary search 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 | 10 ++++++ compiler/rustc_codegen_llvm/src/intrinsic.rs | 31 +++++++++++++++- compiler/rustc_codegen_llvm/src/llvm/ffi.rs | 1 + .../rustc_hir_analysis/src/check/intrinsic.rs | 2 ++ compiler/rustc_span/src/symbol.rs | 1 + library/core/src/intrinsics.rs | 28 +++++++++++++++ .../intrinsics/select_unpredictable.rs | 35 +++++++++++++++++++ 7 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 tests/codegen/intrinsics/select_unpredictable.rs diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 72ff9ea118e..bcb6260862d 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -1353,6 +1353,16 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> { } } + 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) } } diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 68c3d47e826..cd977eb41df 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -11,7 +11,7 @@ use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh} 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 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> { } 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, diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs index 3beda28ac1f..1e2b0292d36 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs @@ -428,6 +428,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, diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index 6282499883b..dfe6a2675c3 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -119,6 +119,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 @@ -487,6 +488,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 => { diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 6d58c4877cb..99a68613324 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1701,6 +1701,7 @@ symbols! { saturating_add, saturating_div, saturating_sub, + select_unpredictable, self_in_typedefs, self_struct_ctor, semitransparent, diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs index c4c63883389..240cc8a21ff 100644 --- a/library/core/src/intrinsics.rs +++ b/library/core/src/intrinsics.rs @@ -1010,6 +1010,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(b: bool, true_val: T, false_val: T) -> T { + if b { true_val } else { false_val } +} + +#[cfg(bootstrap)] +#[inline] +pub fn select_unpredictable(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. /// diff --git a/tests/codegen/intrinsics/select_unpredictable.rs b/tests/codegen/intrinsics/select_unpredictable.rs new file mode 100644 index 00000000000..2054838dd79 --- /dev/null +++ b/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) +}