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`.
This commit is contained in:
Amanieu d'Antras 2024-07-26 19:36:21 +01:00
parent 355efacf0d
commit 4f78f9fbb0
7 changed files with 107 additions and 1 deletions

View file

@ -1353,6 +1353,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) }
}

View file

@ -11,7 +11,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,

View file

@ -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,

View file

@ -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 => {

View file

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

View file

@ -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<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.
///

View file

@ -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)
}