Add new intrinsic is_constant and optimize pow

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);

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.llvm_type(self.cx)),
+                &[args[0].immediate()],
+            ),
             sym::unlikely => self
                 .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]),
             kw::Try => {

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"

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

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,

library/core/src/intrinsics.rs

@@ -2489,6 +2489,51 @@ extern "rust-intrinsic" {
     /// constructing an empty slice) is returned.
     #[rustc_nounwind]
     pub fn option_payload_ptr<T>(arg: *const Option<T>) -> *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<T>(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: ~const crate::marker::Destruct>(_: T) -> bool {
+    false
 }
 
 // Some functions are defined here because they accidentally got made

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)]

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::<Self>() * 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`.

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::<Self>() * 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.

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)?;

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(()));

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