diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs index 24fd5bbf8c526..36502f2d2f0b4 100644 --- a/compiler/rustc_codegen_llvm/src/context.rs +++ b/compiler/rustc_codegen_llvm/src/context.rs @@ -892,6 +892,20 @@ impl<'ll> CodegenCx<'ll, '_> { ifn!("llvm.lifetime.start.p0i8", fn(t_i64, ptr) -> void); ifn!("llvm.lifetime.end.p0i8", fn(t_i64, ptr) -> void); + // FIXME: This is an infinitesimally small portion of the types you can + // pass to this intrinsic, if we can ever lazily register intrinsics we + // should register these when they're used, that way any type can be + // passed. + ifn!("llvm.is.constant.i1", fn(i1) -> i1); + ifn!("llvm.is.constant.i8", fn(t_i8) -> i1); + ifn!("llvm.is.constant.i16", fn(t_i16) -> i1); + ifn!("llvm.is.constant.i32", fn(t_i32) -> i1); + ifn!("llvm.is.constant.i64", fn(t_i64) -> i1); + ifn!("llvm.is.constant.i128", fn(t_i128) -> i1); + ifn!("llvm.is.constant.isize", fn(t_isize) -> i1); + ifn!("llvm.is.constant.f32", fn(t_f32) -> i1); + ifn!("llvm.is.constant.f64", fn(t_f64) -> i1); + ifn!("llvm.expect.i1", fn(i1, i1) -> i1); ifn!("llvm.eh.typeid.for", fn(ptr) -> t_i32); ifn!("llvm.localescape", fn(...) -> void); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index a9b06030e70a6..033680f9e140a 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -119,6 +119,10 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> { sym::likely => { self.call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(true)]) } + sym::is_val_statically_known => self.call_intrinsic( + &format!("llvm.is.constant.{:?}", args[0].layout.immediate_llvm_type(self.cx)), + &[args[0].immediate()], + ), sym::unlikely => self .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]), kw::Try => { diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs index b740b79d16241..00c95dae604ff 100644 --- a/compiler/rustc_const_eval/src/const_eval/machine.rs +++ b/compiler/rustc_const_eval/src/const_eval/machine.rs @@ -542,6 +542,11 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir, )?; } } + // The intrinsic represents whether the value is known to the optimizer (LLVM). + // We're not doing any optimizations here, so there is no optimizer that could know the value. + // (We know the value here in the machine of course, but this is the runtime of that code, + // not the optimization stage.) + sym::is_val_statically_known => ecx.write_scalar(Scalar::from_bool(false), dest)?, _ => { throw_unsup_format!( "intrinsic `{intrinsic_name}` is not supported at compile-time" diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index f89e2e5c25bf2..929c76f8a574c 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -477,6 +477,8 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) { sym::black_box => (1, vec![param(0)], param(0)), + sym::is_val_statically_known => (1, vec![param(0)], tcx.types.bool), + sym::const_eval_select => (4, vec![param(0), param(1), param(2)], param(3)), sym::vtable_size | sym::vtable_align => { diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 28a2dfebcfe96..5544572ff4b96 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -861,6 +861,7 @@ symbols! { intra_doc_pointers, intrinsics, irrefutable_let_patterns, + is_val_statically_known, isa_attribute, isize, issue, diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs index 676d4f2f38ca3..653eee91c124a 100644 --- a/library/core/src/intrinsics.rs +++ b/library/core/src/intrinsics.rs @@ -2489,6 +2489,52 @@ extern "rust-intrinsic" { /// constructing an empty slice) is returned. #[rustc_nounwind] pub fn option_payload_ptr(arg: *const Option) -> *const T; + + /// Returns whether the argument's value is statically known at + /// compile-time. + /// + /// This is useful when there is a way of writing the code that will + /// be *faster* when some variables have known values, but *slower* + /// in the general case: an `if is_val_statically_known(var)` can be used + /// to select between these two variants. The `if` will be optimized away + /// and only the desired branch remains. + /// + /// Formally speaking, this function non-deterministically returns `true` + /// or `false`, and the caller has to ensure sound behavior for both cases. + /// In other words, the following code has *Undefined Behavior*: + /// + /// ```rust + /// if !is_val_statically_known(0) { unreachable_unchecked(); } + /// ``` + /// + /// This also means that the following code's behavior is unspecified; it + /// may panic, or it may not: + /// + /// ```rust,no_run + /// assert_eq!(is_val_statically_known(0), black_box(is_val_statically_known(0))) + /// ``` + /// + /// Unsafe code may not rely on `is_val_statically_known` returning any + /// particular value, ever. However, the compiler will generally make it + /// return `true` only if the value of the argument is actually known. + /// + /// When calling this in a `const fn`, both paths must be semantically + /// equivalent, that is, the result of the `true` branch and the `false` + /// branch must return the same value and have the same side-effects *no + /// matter what*. + #[rustc_const_unstable(feature = "is_val_statically_known", issue = "none")] + #[rustc_nounwind] + #[cfg(not(bootstrap))] + pub fn is_val_statically_known(arg: T) -> bool; +} + +// FIXME: Seems using `unstable` here completely ignores `rustc_allow_const_fn_unstable` +// and thus compiling stage0 core doesn't work. +#[rustc_const_stable(feature = "is_val_statically_known", since = "never")] +#[cfg(bootstrap)] +pub const unsafe fn is_val_statically_known(t: T) -> bool { + mem::forget(t); + false } // Some functions are defined here because they accidentally got made diff --git a/library/core/src/lib.rs b/library/core/src/lib.rs index a2729b3743cc2..94831f15d1a81 100644 --- a/library/core/src/lib.rs +++ b/library/core/src/lib.rs @@ -187,6 +187,7 @@ // // Language features: // tidy-alphabetical-start +#![cfg_attr(not(bootstrap), feature(is_val_statically_known))] #![feature(abi_unadjusted)] #![feature(adt_const_params)] #![feature(allow_internal_unsafe)] diff --git a/library/core/src/num/int_macros.rs b/library/core/src/num/int_macros.rs index 1f43520e1b30a..b8672668862f3 100644 --- a/library/core/src/num/int_macros.rs +++ b/library/core/src/num/int_macros.rs @@ -2039,26 +2039,49 @@ macro_rules! int_impl { without modifying the original"] #[inline] #[rustc_inherit_overflow_checks] + #[rustc_allow_const_fn_unstable(is_val_statically_known)] pub const fn pow(self, mut exp: u32) -> Self { - if exp == 0 { - return 1; - } - let mut base = self; - let mut acc = 1; - - while exp > 1 { - if (exp & 1) == 1 { - acc = acc * base; + // SAFETY: This path has the same behavior as the other. + if unsafe { intrinsics::is_val_statically_known(self) } + && self > 0 + && (self & (self - 1) == 0) + { + let power_used = match self.checked_ilog2() { + Some(v) => v, + // SAFETY: We just checked this is a power of two. and above zero. + None => unsafe { core::hint::unreachable_unchecked() }, + }; + // So it panics. Have to use `overflowing_mul` to efficiently set the + // result to 0 if not. + #[cfg(debug_assertions)] + { + _ = power_used * exp; + } + let (num_shl, overflowed) = power_used.overflowing_mul(exp); + let fine = !overflowed + & (num_shl < (mem::size_of::() * 8) as u32); + (1 << num_shl) * fine as Self + } else { + if exp == 0 { + return 1; + } + let mut base = self; + let mut acc = 1; + + while exp > 1 { + if (exp & 1) == 1 { + acc = acc * base; + } + exp /= 2; + base = base * base; } - exp /= 2; - base = base * base; - } - // since exp!=0, finally the exp must be 1. - // Deal with the final bit of the exponent separately, since - // squaring the base afterwards is not necessary and may cause a - // needless overflow. - acc * base + // since exp!=0, finally the exp must be 1. + // Deal with the final bit of the exponent separately, since + // squaring the base afterwards is not necessary and may cause a + // needless overflow. + acc * base + } } /// Calculates the quotient of Euclidean division of `self` by `rhs`. diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs index 81148c7cc51f7..1e29d426ae54f 100644 --- a/library/core/src/num/uint_macros.rs +++ b/library/core/src/num/uint_macros.rs @@ -1957,26 +1957,60 @@ macro_rules! uint_impl { without modifying the original"] #[inline] #[rustc_inherit_overflow_checks] + #[rustc_allow_const_fn_unstable(is_val_statically_known)] pub const fn pow(self, mut exp: u32) -> Self { - if exp == 0 { - return 1; - } - let mut base = self; - let mut acc = 1; - - while exp > 1 { - if (exp & 1) == 1 { - acc = acc * base; + // LLVM now knows that `self` is a constant value, but not a + // constant in Rust. This allows us to compute the power used at + // compile-time. + // + // This will likely add a branch in debug builds, but this should + // be ok. + // + // This is a massive performance boost in release builds as you can + // get the power of a power of two and the exponent through a `shl` + // instruction, but we must add a couple more checks for parity with + // our own `pow`. + // SAFETY: This path has the same behavior as the other. + if unsafe { intrinsics::is_val_statically_known(self) } + && self.is_power_of_two() + { + let power_used = match self.checked_ilog2() { + Some(v) => v, + // SAFETY: We just checked this is a power of two. `0` is not a + // power of two. + None => unsafe { core::hint::unreachable_unchecked() }, + }; + // So it panics. Have to use `overflowing_mul` to efficiently set the + // result to 0 if not. + #[cfg(debug_assertions)] + { + _ = power_used * exp; + } + let (num_shl, overflowed) = power_used.overflowing_mul(exp); + let fine = !overflowed + & (num_shl < (mem::size_of::() * 8) as u32); + (1 << num_shl) * fine as Self + } else { + if exp == 0 { + return 1; + } + let mut base = self; + let mut acc = 1; + + while exp > 1 { + if (exp & 1) == 1 { + acc = acc * base; + } + exp /= 2; + base = base * base; } - exp /= 2; - base = base * base; - } - // since exp!=0, finally the exp must be 1. - // Deal with the final bit of the exponent separately, since - // squaring the base afterwards is not necessary and may cause a - // needless overflow. - acc * base + // since exp!=0, finally the exp must be 1. + // Deal with the final bit of the exponent separately, since + // squaring the base afterwards is not necessary and may cause a + // needless overflow. + acc * base + } } /// Performs Euclidean division. diff --git a/src/tools/miri/src/shims/intrinsics/mod.rs b/src/tools/miri/src/shims/intrinsics/mod.rs index b2c297fe542c5..0da6c22326a27 100644 --- a/src/tools/miri/src/shims/intrinsics/mod.rs +++ b/src/tools/miri/src/shims/intrinsics/mod.rs @@ -5,6 +5,7 @@ use std::iter; use log::trace; +use rand::Rng; use rustc_apfloat::{Float, Round}; use rustc_middle::ty::layout::LayoutOf; use rustc_middle::{ @@ -131,6 +132,17 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> { this.write_pointer(Pointer::new(ptr.provenance, masked_addr), dest)?; } + // We want to return either `true` or `false` at random, or else something like + // ``` + // if !is_val_statically_known(0) { unreachable_unchecked(); } + // ``` + // Would not be considered UB, or the other way around (`is_val_statically_known(0)`). + "is_val_statically_known" => { + let [_] = check_arg_count(args)?; + let branch: bool = this.machine.rng.get_mut().gen(); + this.write_scalar(Scalar::from_bool(branch), dest)?; + } + // Floating-point operations "fabsf32" => { let [f] = check_arg_count(args)?; diff --git a/src/tools/miri/tests/pass/intrinsics.rs b/src/tools/miri/tests/pass/intrinsics.rs index 8c6eeab22195c..8e46bd7ad48fb 100644 --- a/src/tools/miri/tests/pass/intrinsics.rs +++ b/src/tools/miri/tests/pass/intrinsics.rs @@ -33,6 +33,21 @@ fn main() { assert_eq!(intrinsics::likely(false), false); assert_eq!(intrinsics::unlikely(true), true); + let mut saw_true = false; + let mut saw_false = false; + + for _ in 0..50 { + if unsafe { intrinsics::is_val_statically_known(0) } { + saw_true = true; + } else { + saw_false = true; + } + } + assert!( + saw_true && saw_false, + "`is_val_statically_known` failed to return both true and false. Congrats, you won the lottery!" + ); + intrinsics::forget(Bomb); let _v = intrinsics::discriminant_value(&Some(())); diff --git a/tests/codegen/is_val_statically_known.rs b/tests/codegen/is_val_statically_known.rs new file mode 100644 index 0000000000000..4dcab7442356b --- /dev/null +++ b/tests/codegen/is_val_statically_known.rs @@ -0,0 +1,50 @@ +// #[cfg(bootstrap)] +// ignore-stage1 +// compile-flags: --crate-type=lib -Zmerge-functions=disabled + +#![feature(core_intrinsics)] + +use std::intrinsics::is_val_statically_known; + +pub struct A(u32); +pub enum B { + Ye(u32), +} + +#[inline] +pub fn _u32(a: u32) -> i32 { + if unsafe { is_val_statically_known(a) } { 1 } else { 0 } +} + +// CHECK-LABEL: @_u32_true( +#[no_mangle] +pub fn _u32_true() -> i32 { + // CHECK: ret i32 1 + _u32(1) +} + +// CHECK-LABEL: @_u32_false( +#[no_mangle] +pub fn _u32_false(a: u32) -> i32 { + // CHECK: ret i32 0 + _u32(a) +} + +#[inline] +pub fn _bool(b: bool) -> i32 { + if unsafe { is_val_statically_known(b) } { 3 } else { 2 } +} + +// CHECK-LABEL: @_bool_true( +#[no_mangle] +pub fn _bool_true() -> i32 { + // CHECK: ret i32 3 + _bool(true) +} + +// CHECK-LABEL: @_bool_false( +#[no_mangle] +pub fn _bool_false(b: bool) -> i32 { + // CHECK: ret i32 2 + _bool(b) +} diff --git a/tests/codegen/pow_of_two.rs b/tests/codegen/pow_of_two.rs new file mode 100644 index 0000000000000..3bce5535c66ec --- /dev/null +++ b/tests/codegen/pow_of_two.rs @@ -0,0 +1,68 @@ +// #[cfg(bootstrap)] +// ignore-stage1 +// compile-flags: --crate-type=lib -Zmerge-functions=disabled + +// CHECK-LABEL: @a( +#[no_mangle] +pub fn a(exp: u32) -> u64 { + // CHECK: %[[R:.+]] = and i32 %exp, 63 + // CHECK: %[[R:.+]] = zext i32 %[[R:.+]] to i64 + // CHECK: %[[R:.+]] = shl nuw i64 %[[R:.+]].i, %[[R:.+]] + // CHECK: ret i64 %[[R:.+]] + 2u64.pow(exp) +} + +#[no_mangle] +pub fn b(exp: u32) -> i64 { + // CHECK: %[[R:.+]] = and i32 %exp, 63 + // CHECK: %[[R:.+]] = zext i32 %[[R:.+]] to i64 + // CHECK: %[[R:.+]] = shl nuw i64 %[[R:.+]].i, %[[R:.+]] + // CHECK: ret i64 %[[R:.+]] + 2i64.pow(exp) +} + +// CHECK-LABEL: @c( +#[no_mangle] +pub fn c(exp: u32) -> u32 { + // CHECK: %[[R:.+]].0.i = shl i32 %exp, 1 + // CHECK: %[[R:.+]].1.i = icmp sgt i32 %exp, -1 + // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32 + // CHECK: %fine.i = and i1 %[[R:.+]].1.i, %[[R:.+]].i + // CHECK: %0 = and i32 %[[R:.+]].0.i, 30 + // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32 + // CHECK: %[[R:.+]] = shl nuw nsw i32 %[[R:.+]].i, %0 + // CHECK: ret i32 %[[R:.+]] + 4u32.pow(exp) +} + +// CHECK-LABEL: @d( +#[no_mangle] +pub fn d(exp: u32) -> u32 { + // CHECK: tail call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %exp, i32 5) + // CHECK: %[[R:.+]].0.i = extractvalue { i32, i1 } %[[R:.+]], 0 + // CHECK: %[[R:.+]].1.i = extractvalue { i32, i1 } %[[R:.+]], 1 + // CHECK: %[[R:.+]].i = xor i1 %[[R:.+]].1.i, true + // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32 + // CHECK: %fine.i = and i1 %[[R:.+]].i, %[[R:.+]].i + // CHECK: %[[R:.+]] = and i32 %[[R:.+]].0.i, 31 + // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32 + // CHECK: %[[R:.+]] = shl nuw i32 %[[R:.+]].i, %1 + // CHECK: ret i32 %[[R:.+]] + 32u32.pow(exp) +} + +// CHECK-LABEL: @e( +#[no_mangle] +pub fn e(exp: u32) -> i32 { + // CHECK: tail call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %exp, i32 5) + // CHECK: %[[R:.+]].0.i = extractvalue { i32, i1 } %[[R:.+]], 0 + // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32 + // CHECK: %[[R:.+]].1.i = extractvalue { i32, i1 } %[[R:.+]], 1 + // CHECK: %[[R:.+]].i = xor i1 %[[R:.+]].1.i, true + // CHECK: %fine.i = and i1 %[[R:.+]].i, %[[R:.+]].i + // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32 + // CHECK: %[[R:.+]] = and i32 %[[R:.+]].0.i, 31 + // CHECK: %[[R:.+]] = shl nuw i32 %[[R:.+]].i, %1 + // CHECK: ret i32 %[[R:.+]] + 32i32.pow(exp) +} diff --git a/tests/ui/consts/is_val_statically_known.rs b/tests/ui/consts/is_val_statically_known.rs new file mode 100644 index 0000000000000..b0565842eb4e2 --- /dev/null +++ b/tests/ui/consts/is_val_statically_known.rs @@ -0,0 +1,15 @@ +// run-pass + +#![feature(core_intrinsics)] +#![feature(is_val_statically_known)] + +use std::intrinsics::is_val_statically_known; + +const CONST_TEST: bool = unsafe { is_val_statically_known(0) }; + +fn main() { + if CONST_TEST { + unreachable!("currently expected to return false during const eval"); + // but note that this is not a guarantee! + } +}