Add new intrinsic is_constant and optimize pow

Author: Centri3

compiler/rustc_codegen_llvm/src/context.rs

@@ -908,6 +908,8 @@ impl<'ll> CodegenCx<'ll, '_> {
         ifn!("llvm.lifetime.start.p0i8", fn(t_i64, ptr) -> void);
         ifn!("llvm.lifetime.end.p0i8", fn(t_i64, ptr) -> void);
 
+        ifn!("llvm.is.constant", fn(...) -> 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,7 @@ 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_constant => self.call_intrinsic("llvm.is.constant", &[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/interpret/intrinsics.rs

@@ -216,6 +216,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             sym::copy => {
                 self.copy_intrinsic(&args[0], &args[1], &args[2], /*nonoverlapping*/ false)?;
             }
+            sym::is_constant => self.write_scalar(Scalar::from_bool(true), dest)?,
             sym::write_bytes => {
                 self.write_bytes_intrinsic(&args[0], &args[1], &args[2])?;
             }

compiler/rustc_hir_analysis/messages.ftl

@@ -178,6 +178,8 @@ hir_analysis_invalid_union_field =
 hir_analysis_invalid_union_field_sugg =
     wrap the field type in `ManuallyDrop<...>`
 
+hir_analysis_is_constant_zst = parameter for `is_constant` cannot be zero-sized
+
 hir_analysis_late_bound_const_in_apit = `impl Trait` can only mention const parameters from an fn or impl
     .label = const parameter declared here
 

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -112,7 +112,8 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: DefId) -> hir
         | sym::forget
         | sym::black_box
         | sym::variant_count
-        | sym::ptr_mask => hir::Unsafety::Normal,
+        | sym::ptr_mask
+        | sym::is_constant => hir::Unsafety::Normal,
         _ => hir::Unsafety::Unsafe,
     };
 
@@ -453,6 +454,11 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
 
             sym::black_box => (1, vec![param(0)], param(0)),
 
+            sym::is_constant => {
+                // FIXME: ZSTs cause an ICE. We should check for this.
+                (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

@@ -907,6 +907,7 @@ symbols! {
         io_stderr,
         io_stdout,
         irrefutable_let_patterns,
+        is_constant,
         isa_attribute,
         isize,
         issue,

library/core/src/intrinsics.rs

@@ -2511,6 +2511,25 @@ extern "rust-intrinsic" {
     where
         G: FnOnce<ARG, Output = RET>,
         F: FnOnce<ARG, Output = RET>;
+
+    /// Returns whether the argument is known at compile-time. This opens the
+    /// door for additional optimizations, in that LLVM can then optimize
+    /// following checks to either `true` or `false`. This is often paired with
+    /// an `if-else` statement, as LLVM will only keep one branch (if
+    /// optimizations are on).
+    ///
+    /// "Constant" in this context is not the same as a constant in Rust. As
+    /// such, this should only be used for optimizations.
+    ///
+    /// 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.
+    #[rustc_const_stable(feature = "todo", since = "never")]
+    #[rustc_safe_intrinsic]
+    #[rustc_nounwind]
+    #[cfg(not(bootstrap))]
+    pub fn is_constant<T>(arg: T) -> bool;
 }
 
 // Some functions are defined here because they accidentally got made

library/core/src/num/int_macros.rs

@@ -2089,25 +2089,65 @@ macro_rules! int_impl {
         #[inline]
         #[rustc_inherit_overflow_checks]
         pub const fn pow(self, mut exp: u32) -> Self {
-            if exp == 0 {
-                return 1;
+            #[cfg(not(bootstrap))]
+            if intrinsics::is_constant(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
+                    & (power_used < (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;
+                }
+
+                // 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
             }
-            let mut base = self;
-            let mut acc = 1;
 
-            while exp > 1 {
-                if (exp & 1) == 1 {
-                    acc = acc * base;
+            #[cfg(bootstrap)]
+            {
+                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
+            }
         }
 
         /// Returns the square root of the number, rounded down.

library/core/src/num/uint_macros.rs

@@ -1974,25 +1974,77 @@ macro_rules! uint_impl {
         #[inline]
         #[rustc_inherit_overflow_checks]
         pub const fn pow(self, mut exp: u32) -> Self {
-            if exp == 0 {
-                return 1;
+            // 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`.
+            #[cfg(not(bootstrap))]
+            if intrinsics::is_constant(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
+                    & (power_used < (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;
+                }
+
+                // 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
             }
-            let mut base = self;
-            let mut acc = 1;
 
-            while exp > 1 {
-                if (exp & 1) == 1 {
-                    acc = acc * base;
+            #[cfg(bootstrap)]
+            {
+                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
+            }
         }
 
         /// Returns the square root of the number, rounded down.

tests/codegen/is_constant.rs

@@ -0,0 +1,47 @@
+// compile-flags: --crate-type=lib
+#![feature(core_intrinsics)]
+
+use std::intrinsics::is_constant;
+
+pub struct A(u32);
+pub enum B {
+    Ye(u32),
+}
+
+#[inline]
+pub fn tuple_struct(a: A) -> i32 {
+    if is_constant(a) { 1 } else { 0 }
+}
+
+// CHECK-LABEL: @tuple_struct_true(
+#[no_mangle]
+pub fn tuple_struct_true() -> i32 {
+    // CHECK: ret i32 1
+    tuple_struct(A(1))
+}
+
+// CHECK-LABEL: @tuple_struct_false(
+#[no_mangle]
+pub fn tuple_struct_false(a: A) -> i32 {
+    // CHECK: ret i32 0
+    tuple_struct(a)
+}
+
+#[inline]
+pub fn r#enum(b: B) -> i32 {
+    if is_constant(b) { 3 } else { 2 }
+}
+
+// CHECK-LABEL: @enum_true(
+#[no_mangle]
+pub fn enum_true() -> i32 {
+    // CHECK: ret i32 3
+    r#enum(B::Ye(2))
+}
+
+// CHECK-LABEL: @enum_false(
+#[no_mangle]
+pub fn enum_false(b: B) -> i32 {
+    // CHECK: ret i32 2
+    r#enum(b)
+}