From 69942f01328629764558afc0e4754cdf6f663e5a Mon Sep 17 00:00:00 2001 From: joboet Date: Fri, 6 Dec 2024 15:05:14 +0100 Subject: [PATCH 1/4] core: implement `bool::select_unpredictable` --- library/core/src/bool.rs | 46 ++++++++++++++++++++++++++++++ library/core/src/intrinsics/mod.rs | 2 +- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/library/core/src/bool.rs b/library/core/src/bool.rs index 58a870d2e0725..b9590dcc1e7de 100644 --- a/library/core/src/bool.rs +++ b/library/core/src/bool.rs @@ -60,4 +60,50 @@ impl bool { pub fn then T>(self, f: F) -> Option { if self { Some(f()) } else { None } } + + /// Returns either `true_val` or `false_val` depending on the value of + /// `condition`, with a hint to the compiler that `condition` is unlikely + /// to be correctly predicted by a CPU’s branch predictor. + /// + /// This method is functionally equivalent to writing + /// ```ignore (this is just for illustrative purposes) + /// if b { true_val } else { false_val } + /// ``` + /// but might generate different assembly. In particular, on platforms with + /// a conditional move or select instruction (like `cmov` on x86 or `csel` + /// on ARM) the optimizer might use these instructions to avoid branches, + /// which can benefit performance if the branch predictor is struggling + /// with predicting `condition`, such as in an implementation of binary + /// search. + /// + /// Note however that this lowering is not guaranteed (on any platform) and + /// should not be relied upon when trying to write constant-time code. Also + /// be aware that this lowering might *decrease* performance if `condition` + /// is well-predictable. It is advisable to perform benchmarks to tell if + /// this function is useful. + /// + /// # Examples + /// + /// Distribute values evenly between two buckets: + /// ``` + /// #![feature(select_unpredictable)] + /// + /// use std::hash::BuildHasher; + /// + /// fn append(hasher: &H, v: i32, bucket_one: &mut Vec, bucket_two: &mut Vec) { + /// let hash = hasher.hash_one(&v); + /// let bucket = (hash % 2 == 0).select_unpredictable(bucket_one, bucket_two); + /// bucket.push(v); + /// } + /// # let hasher = std::collections::hash_map::RandomState::new(); + /// # let mut bucket_one = Vec::new(); + /// # let mut bucket_two = Vec::new(); + /// # append(&hasher, 42, &mut bucket_one, &mut bucket_two); + /// # assert_eq!(bucket_one.len() + bucket_two.len(), 1); + /// ``` + #[inline(always)] + #[unstable(feature = "select_unpredictable", issue = "133962")] + pub fn select_unpredictable(self, true_val: T, false_val: T) -> T { + crate::intrinsics::select_unpredictable(self, true_val, false_val) + } } diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 802b571c51067..f21e881641819 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -1544,7 +1544,7 @@ pub const fn unlikely(b: bool) -> bool { /// Therefore, implementations must not require the user to uphold /// any safety invariants. /// -/// This intrinsic does not have a stable counterpart. +/// The public form of this instrinsic is [`bool::select_unpredictable`]. #[unstable(feature = "core_intrinsics", issue = "none")] #[rustc_intrinsic] #[rustc_nounwind] From 13c77ba34d16255d9f5a57758140c73a86f1895c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20B=C3=B6ttiger?= Date: Mon, 9 Dec 2024 13:22:39 +0100 Subject: [PATCH 2/4] core: improve comments Co-authored-by: Yotam Ofek Co-authored-by: Hanna Kruppe --- library/core/src/bool.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/library/core/src/bool.rs b/library/core/src/bool.rs index b9590dcc1e7de..5c9b0e2810631 100644 --- a/library/core/src/bool.rs +++ b/library/core/src/bool.rs @@ -62,12 +62,14 @@ impl bool { } /// Returns either `true_val` or `false_val` depending on the value of - /// `condition`, with a hint to the compiler that `condition` is unlikely + /// `self`, with a hint to the compiler that `self` is unlikely /// to be correctly predicted by a CPU’s branch predictor. /// - /// This method is functionally equivalent to writing + /// This method is functionally equivalent to /// ```ignore (this is just for illustrative purposes) - /// if b { true_val } else { false_val } + /// fn select_unpredictable(b: bool, true_val: T, false_val: T) -> T { + /// if b { true_val } else { false_val } + /// } /// ``` /// but might generate different assembly. In particular, on platforms with /// a conditional move or select instruction (like `cmov` on x86 or `csel` From 49d76b8d0e9d3647bbc26dce9fd50aab66ea0617 Mon Sep 17 00:00:00 2001 From: joboet Date: Mon, 9 Dec 2024 13:27:06 +0100 Subject: [PATCH 3/4] core: use public method instead of instrinsic --- library/core/src/slice/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index 191eaccff9899..2bcea43d58843 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -7,7 +7,7 @@ #![stable(feature = "rust1", since = "1.0.0")] use crate::cmp::Ordering::{self, Equal, Greater, Less}; -use crate::intrinsics::{exact_div, select_unpredictable, unchecked_sub}; +use crate::intrinsics::{exact_div, unchecked_sub}; use crate::mem::{self, SizedTypeProperties}; use crate::num::NonZero; use crate::ops::{Bound, OneSidedRange, Range, RangeBounds, RangeInclusive}; @@ -2835,7 +2835,7 @@ impl [T] { // Binary search interacts poorly with branch prediction, so force // the compiler to use conditional moves if supported by the target // architecture. - base = select_unpredictable(cmp == Greater, base, mid); + base = (cmp == Greater).select_unpredictable(base, mid); // This is imprecise in the case where `size` is odd and the // comparison returns Greater: the mid element still gets included From 8f3aa358bf4c5507eccaca2f315d94eaef3b48d9 Mon Sep 17 00:00:00 2001 From: joboet Date: Fri, 3 Jan 2025 19:44:08 +0100 Subject: [PATCH 4/4] add codegen test for `bool::select_unpredictable` --- tests/codegen/bool-select-unpredictable.rs | 35 ++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 tests/codegen/bool-select-unpredictable.rs diff --git a/tests/codegen/bool-select-unpredictable.rs b/tests/codegen/bool-select-unpredictable.rs new file mode 100644 index 0000000000000..1562b17754203 --- /dev/null +++ b/tests/codegen/bool-select-unpredictable.rs @@ -0,0 +1,35 @@ +//@ compile-flags: -O + +#![feature(select_unpredictable)] +#![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 + p.select_unpredictable(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 + p.select_unpredictable(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 + p.select_unpredictable(a, b) +} + +#[no_mangle] +pub fn test_zst(p: bool, a: (), b: ()) -> () { + // CHECK-LABEL: define{{.*}} @test_zst + p.select_unpredictable(a, b) +}