Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add new intrinsic is_constant and optimize pow #114390

Closed
wants to merge 11 commits into from
14 changes: 14 additions & 0 deletions compiler/rustc_codegen_llvm/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
4 changes: 4 additions & 0 deletions compiler/rustc_codegen_llvm/src/intrinsic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 => {
Expand Down
5 changes: 5 additions & 0 deletions compiler/rustc_const_eval/src/const_eval/machine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?,
Centri3 marked this conversation as resolved.
Show resolved Hide resolved
_ => {
throw_unsup_format!(
"intrinsic `{intrinsic_name}` is not supported at compile-time"
Expand Down
2 changes: 2 additions & 0 deletions compiler/rustc_hir_analysis/src/check/intrinsic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 => {
Expand Down
1 change: 1 addition & 0 deletions compiler/rustc_span/src/symbol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -861,6 861,7 @@ symbols! {
intra_doc_pointers,
intrinsics,
irrefutable_let_patterns,
is_val_statically_known,
isa_attribute,
isize,
issue,
Expand Down
46 changes: 46 additions & 0 deletions library/core/src/intrinsics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2489,6 2489,52 @@ 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)))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't the black_box be the inner call?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Either way, the behavior is unspecified, even without black_box. imo I think it being outside the call makes more sense since it conveys better that optimizations may change the result, while calling is_val_statically_known on black_box may return false because it's a function call, or smth.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about entirely removing black_box from the example?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good

/// ```
Comment on lines 2506 to 2515
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both doc tests fail when ran with ./x test core --stage 1 --doc -- is_val_statically_known. The doc tests are never ran with --stage 0 because of #[cfg(not(bootstrap())].

You need to add unsafe, use, and feature statements to make them compile. If you think they should be omitted from the docs, then you can prepend them with #.

Here is an example that passes the tests:

  /// ```rust
  /// #![feature(is_val_statically_known)]
  /// #![feature(core_intrinsics)]
  /// use std::hint::unreachable_unchecked;
  /// use std::intrinsics::is_val_statically_known;
  ///
  /// unsafe {
  ///    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
  /// #![feature(is_val_statically_known)]
  /// #![feature(core_intrinsics)]
  /// use std::hint::black_box;
  /// use std::intrinsics::is_val_statically_known;
  ///
  /// unsafe {
  ///     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.
Centri3 marked this conversation as resolved.
Show resolved Hide resolved
///
/// 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*.
Comment on lines 2521 to 2524
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are they allowed to differ in panics? E.g. if an invalid input is handed to a safe function then it may either panic or return garbage and this differs between the branches?

Copy link
Member Author

@Centri3 Centri3 Sep 5, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding onto this, are panic messages/locations allowed to differ (if it always panics with the same inputs in CTFE/runtime)? They do in pow.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd say they must either both panic or neither panic. But the exact panic location is not relevant.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That should be documented then. Though I think it might be tricky to uphold. Would it be possible to have is_val_statically_known to always stick to one value in const eval? That would be less footgun-ish.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we could (in debug mode?) fork the const evaluator and run both paths and compare the result 🙃

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That should be documented then. Though I think it might be tricky to uphold. Would it be possible to have is_val_statically_known to always stick to one value in const eval? That would be less footgun-ish.

It's the same with const_eval_select. We already have an RFC up for resolving this properly; in the mean time I think we should keep the strict rules.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

They do in pow.

Could #[track_caller] and custom panic messages fix this? For that matter, is there a reason pow didn't have #[track_caller] to start out with?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't think there's any reason not to. Custom panic messages would be a nice addition also

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I want to add, for the benefit of those reading in the future, that the core library already intentionally gives different panic messages depending on const_eval_select. It appears this is for performance reasons because the runtime variant is static, while the compile time variant is not. As far as I know, that is the only current use of const_eval_select.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's used for a couple other things as well, like having certain assertions only in the run-time version of some low-level functions (assert_unsafe_precondition), and const-ifying functions whose runtime panic message needs the formatting machinery which is not supported at compile-time.

#[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>(t: T) -> bool {
mem::forget(t);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is forget okay here? Does the real intrinsic also forget its argument? Should we restrict it to Copy types to avoid the confusion that could cause?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy is fine. This is cfg(bootstrap) so it should be gone by the next beta release anyway.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yea that may be better actually ^^ We don't support anything else in the backend anyway

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be Copy on both of them then.

false
}

// Some functions are defined here because they accidentally got made
Expand Down
1 change: 1 addition & 0 deletions library/core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
57 changes: 40 additions & 17 deletions library/core/src/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is incorrect to use #[cfg(debug_assertions)] here. In general, we can't assume debug_assertions matches overflow_checks even though they go together by default. This is especially true because this function has the #[rustc_inherit_overflow_checks].

Without #[cfg(debug_assertions)], it is still incomplete because the greatest possible u32 is much greater than the number of bits in an integer.

To correctly check for overflow, test to see if the final result is less than 1. If that is the case use _ = <Self>::MAX * <Self>::MAX; to force a panic iff overflow_checks are enabled.

Although it isn't necessary, you can make the code less verbose by using saturating_mul instead of overflowing_mul and by using (1 as Self).checked_shl().unwrap_or_default() instead of (1 << num_shl) * fine as Self.

{
_ = 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`.
Expand Down
68 changes: 51 additions & 17 deletions library/core/src/num/uint_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
12 changes: 12 additions & 0 deletions src/tools/miri/src/shims/intrinsics/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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)?;
Expand Down
15 changes: 15 additions & 0 deletions src/tools/miri/tests/pass/intrinsics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Comment on lines 37 to 38
Copy link
Member

@RalfJung RalfJung Sep 5, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
let mut saw_false = false;
let mut saw_false = false;

just a nit to make it visually more clear that these are all together testing one thing

EDIT: hm I tried to remove the empty line, not sure if github supports that...

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(()));
Expand Down
50 changes: 50 additions & 0 deletions tests/codegen/is_val_statically_known.rs
Original file line number Diff line number Diff line change
@@ -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)
}
Loading
Loading