Skip to content

Commit

Permalink
Rollup merge of #123226 - scottmcm:u32-shifts, r=WaffleLapkin
Browse files Browse the repository at this point in the history
De-LLVM the unchecked shifts [MCP#693]

This is just one part of the MCP (rust-lang/compiler-team#693), but it's the one that IMHO removes the most noise from the standard library code.

Seems net simpler this way, since MIR already supported heterogeneous shifts anyway, and thus it's not more work for backends than before.

r? WaffleLapkin
  • Loading branch information
matthiaskrgr authored Apr 2, 2024
2 parents d63ddef + 4626521 commit 1b0e46f
Show file tree
Hide file tree
Showing 32 changed files with 408 additions and 576 deletions.
35 changes: 31 additions & 4 deletions compiler/rustc_codegen_ssa/src/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use crate::back::write::{
compute_per_cgu_lto_type, start_async_codegen, submit_codegened_module_to_llvm,
submit_post_lto_module_to_llvm, submit_pre_lto_module_to_llvm, ComputedLtoType, OngoingCodegen,
};
use crate::common::{IntPredicate, RealPredicate, TypeKind};
use crate::common::{self, IntPredicate, RealPredicate, TypeKind};
use crate::errors;
use crate::meth;
use crate::mir;
Expand Down Expand Up @@ -33,7 +33,7 @@ use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder, MonoItem};
use rustc_middle::query::Providers;
use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout};
use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
use rustc_session::config::{self, CrateType, EntryFnType, OutputType};
use rustc_session::config::{self, CrateType, EntryFnType, OptLevel, OutputType};
use rustc_session::Session;
use rustc_span::symbol::sym;
use rustc_span::Symbol;
Expand Down Expand Up @@ -300,14 +300,35 @@ pub fn coerce_unsized_into<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
}
}

pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
/// Returns `rhs` sufficiently masked, truncated, and/or extended so that
/// it can be used to shift `lhs`.
///
/// Shifts in MIR are all allowed to have mismatched LHS & RHS types.
/// The shift methods in `BuilderMethods`, however, are fully homogeneous
/// (both parameters and the return type are all the same type).
///
/// If `is_unchecked` is false, this masks the RHS to ensure it stays in-bounds,
/// as the `BuilderMethods` shifts are UB for out-of-bounds shift amounts.
/// For 32- and 64-bit types, this matches the semantics
/// of Java. (See related discussion on #1877 and #10183.)
///
/// If `is_unchecked` is true, this does no masking, and adds sufficient `assume`
/// calls or operation flags to preserve as much freedom to optimize as possible.
pub fn build_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
bx: &mut Bx,
lhs: Bx::Value,
rhs: Bx::Value,
mut rhs: Bx::Value,
is_unchecked: bool,
) -> Bx::Value {
// Shifts may have any size int on the rhs
let mut rhs_llty = bx.cx().val_ty(rhs);
let mut lhs_llty = bx.cx().val_ty(lhs);

let mask = common::shift_mask_val(bx, lhs_llty, rhs_llty, false);
if !is_unchecked {
rhs = bx.and(rhs, mask);
}

if bx.cx().type_kind(rhs_llty) == TypeKind::Vector {
rhs_llty = bx.cx().element_type(rhs_llty)
}
Expand All @@ -317,6 +338,12 @@ pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
let rhs_sz = bx.cx().int_width(rhs_llty);
let lhs_sz = bx.cx().int_width(lhs_llty);
if lhs_sz < rhs_sz {
if is_unchecked && bx.sess().opts.optimize != OptLevel::No {
// FIXME: Use `trunc nuw` once that's available
let inrange = bx.icmp(IntPredicate::IntULE, rhs, mask);
bx.assume(inrange);
}

bx.trunc(rhs, lhs_llty)
} else if lhs_sz > rhs_sz {
// We zero-extend even if the RHS is signed. So e.g. `(x: i32) << -1i8` will zero-extend the
Expand Down
41 changes: 1 addition & 40 deletions compiler/rustc_codegen_ssa/src/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,9 @@
use rustc_hir::LangItem;
use rustc_middle::mir;
use rustc_middle::ty::Instance;
use rustc_middle::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
use rustc_middle::ty::{self, layout::TyAndLayout, TyCtxt};
use rustc_span::Span;

use crate::base;
use crate::traits::*;

#[derive(Copy, Clone)]
Expand Down Expand Up @@ -128,44 +127,6 @@ pub fn build_langcall<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
(bx.fn_abi_of_instance(instance, ty::List::empty()), bx.get_fn_addr(instance), instance)
}

// To avoid UB from LLVM, these two functions mask RHS with an
// appropriate mask unconditionally (i.e., the fallback behavior for
// all shifts). For 32- and 64-bit types, this matches the semantics
// of Java. (See related discussion on #1877 and #10183.)

pub fn build_masked_lshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
bx: &mut Bx,
lhs: Bx::Value,
rhs: Bx::Value,
) -> Bx::Value {
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
// #1877, #10183: Ensure that input is always valid
let rhs = shift_mask_rhs(bx, rhs);
bx.shl(lhs, rhs)
}

pub fn build_masked_rshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
bx: &mut Bx,
lhs_t: Ty<'tcx>,
lhs: Bx::Value,
rhs: Bx::Value,
) -> Bx::Value {
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
// #1877, #10183: Ensure that input is always valid
let rhs = shift_mask_rhs(bx, rhs);
let is_signed = lhs_t.is_signed();
if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
}

fn shift_mask_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
bx: &mut Bx,
rhs: Bx::Value,
) -> Bx::Value {
let rhs_llty = bx.val_ty(rhs);
let shift_val = shift_mask_val(bx, rhs_llty, rhs_llty, false);
bx.and(rhs, shift_val)
}

pub fn shift_mask_val<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
bx: &mut Bx,
llty: Bx::Type,
Expand Down
12 changes: 5 additions & 7 deletions compiler/rustc_codegen_ssa/src/mir/rvalue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use super::place::PlaceRef;
use super::{FunctionCx, LocalRef};

use crate::base;
use crate::common::{self, IntPredicate};
use crate::common::IntPredicate;
use crate::traits::*;
use crate::MemFlags;

Expand Down Expand Up @@ -860,14 +860,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
bx.inbounds_gep(llty, lhs, &[rhs])
}
}
mir::BinOp::Shl => common::build_masked_lshift(bx, lhs, rhs),
mir::BinOp::ShlUnchecked => {
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
mir::BinOp::Shl | mir::BinOp::ShlUnchecked => {
let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShlUnchecked);
bx.shl(lhs, rhs)
}
mir::BinOp::Shr => common::build_masked_rshift(bx, input_ty, lhs, rhs),
mir::BinOp::ShrUnchecked => {
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
mir::BinOp::Shr | mir::BinOp::ShrUnchecked => {
let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShrUnchecked);
if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
}
mir::BinOp::Ne
Expand Down
5 changes: 2 additions & 3 deletions compiler/rustc_hir_analysis/src/check/intrinsic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -454,9 +454,8 @@ pub fn check_intrinsic_type(
sym::unchecked_div | sym::unchecked_rem | sym::exact_div => {
(1, 0, vec![param(0), param(0)], param(0))
}
sym::unchecked_shl | sym::unchecked_shr | sym::rotate_left | sym::rotate_right => {
(1, 0, vec![param(0), param(0)], param(0))
}
sym::unchecked_shl | sym::unchecked_shr => (2, 0, vec![param(0), param(1)], param(0)),
sym::rotate_left | sym::rotate_right => (1, 0, vec![param(0), param(0)], param(0)),
sym::unchecked_add | sym::unchecked_sub | sym::unchecked_mul => {
(1, 0, vec![param(0), param(0)], param(0))
}
Expand Down
6 changes: 4 additions & 2 deletions library/core/src/intrinsics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2224,18 +2224,20 @@ extern "rust-intrinsic" {
/// Safe wrappers for this intrinsic are available on the integer
/// primitives via the `checked_shl` method. For example,
/// [`u32::checked_shl`]
#[cfg(not(bootstrap))]
#[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
#[rustc_nounwind]
pub fn unchecked_shl<T: Copy>(x: T, y: T) -> T;
pub fn unchecked_shl<T: Copy, U: Copy>(x: T, y: U) -> T;
/// Performs an unchecked right shift, resulting in undefined behavior when
/// `y < 0` or `y >= N`, where N is the width of T in bits.
///
/// Safe wrappers for this intrinsic are available on the integer
/// primitives via the `checked_shr` method. For example,
/// [`u32::checked_shr`]
#[cfg(not(bootstrap))]
#[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
#[rustc_nounwind]
pub fn unchecked_shr<T: Copy>(x: T, y: T) -> T;
pub fn unchecked_shr<T: Copy, U: Copy>(x: T, y: U) -> T;

/// Returns the result of an unchecked addition, resulting in
/// undefined behavior when `x + y > T::MAX` or `x + y < T::MIN`.
Expand Down
32 changes: 24 additions & 8 deletions library/core/src/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1227,10 +1227,18 @@ macro_rules! int_impl {
#[inline(always)]
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shl`.
// Any legal shift amount is losslessly representable in the self type.
unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
#[cfg(bootstrap)]
{
// For bootstrapping, just use built-in primitive shift.
// panicking is a legal manifestation of UB
self << rhs
}
#[cfg(not(bootstrap))]
{
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shl`.
unsafe { intrinsics::unchecked_shl(self, rhs) }
}
}

/// Checked shift right. Computes `self >> rhs`, returning `None` if `rhs` is
Expand Down Expand Up @@ -1310,10 +1318,18 @@ macro_rules! int_impl {
#[inline(always)]
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shr`.
// Any legal shift amount is losslessly representable in the self type.
unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
#[cfg(bootstrap)]
{
// For bootstrapping, just use built-in primitive shift.
// panicking is a legal manifestation of UB
self >> rhs
}
#[cfg(not(bootstrap))]
{
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shr`.
unsafe { intrinsics::unchecked_shr(self, rhs) }
}
}

/// Checked absolute value. Computes `self.abs()`, returning `None` if
Expand Down
11 changes: 0 additions & 11 deletions library/core/src/num/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -285,17 +285,6 @@ macro_rules! widening_impl {
};
}

macro_rules! conv_rhs_for_unchecked_shift {
($SelfT:ty, $x:expr) => {{
// If the `as` cast will truncate, ensure we still tell the backend
// that the pre-truncation value was also small.
if <$SelfT>::BITS < 32 {
intrinsics::assume($x <= (<$SelfT>::MAX as u32));
}
$x as $SelfT
}};
}

impl i8 {
int_impl! {
Self = i8,
Expand Down
32 changes: 24 additions & 8 deletions library/core/src/num/uint_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1286,10 +1286,18 @@ macro_rules! uint_impl {
#[inline(always)]
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shl`.
// Any legal shift amount is losslessly representable in the self type.
unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
#[cfg(bootstrap)]
{
// For bootstrapping, just use built-in primitive shift.
// panicking is a legal manifestation of UB
self << rhs
}
#[cfg(not(bootstrap))]
{
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shl`.
unsafe { intrinsics::unchecked_shl(self, rhs) }
}
}

/// Checked shift right. Computes `self >> rhs`, returning `None`
Expand Down Expand Up @@ -1369,10 +1377,18 @@ macro_rules! uint_impl {
#[inline(always)]
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shr`.
// Any legal shift amount is losslessly representable in the self type.
unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
#[cfg(bootstrap)]
{
// For bootstrapping, just use built-in primitive shift.
// panicking is a legal manifestation of UB
self >> rhs
}
#[cfg(not(bootstrap))]
{
// SAFETY: the caller must uphold the safety contract for
// `unchecked_shr`.
unsafe { intrinsics::unchecked_shr(self, rhs) }
}
}

/// Checked exponentiation. Computes `self.pow(exp)`, returning `None` if
Expand Down
14 changes: 12 additions & 2 deletions library/core/src/ptr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1781,9 +1781,19 @@ pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usiz
// FIXME(#75598): Direct use of these intrinsics improves codegen significantly at opt-level <=
// 1, where the method versions of these operations are not inlined.
use intrinsics::{
assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_shl,
unchecked_shr, unchecked_sub, wrapping_add, wrapping_mul, wrapping_sub,
assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_sub,
wrapping_add, wrapping_mul, wrapping_sub,
};
#[cfg(bootstrap)]
const unsafe fn unchecked_shl(value: usize, shift: usize) -> usize {
value << shift
}
#[cfg(bootstrap)]
const unsafe fn unchecked_shr(value: usize, shift: usize) -> usize {
value >> shift
}
#[cfg(not(bootstrap))]
use intrinsics::{unchecked_shl, unchecked_shr};

/// Calculate multiplicative modular inverse of `x` modulo `m`.
///
Expand Down
49 changes: 47 additions & 2 deletions tests/codegen/unchecked_shifts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#![crate_type = "lib"]
#![feature(unchecked_shifts)]
#![feature(core_intrinsics)]

// CHECK-LABEL: @unchecked_shl_unsigned_same
#[no_mangle]
Expand All @@ -19,7 +20,7 @@ pub unsafe fn unchecked_shl_unsigned_smaller(a: u16, b: u32) -> u16 {
// This uses -DAG to avoid failing on irrelevant reorderings,
// like emitting the truncation earlier.

// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 65536
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 16
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
// CHECK-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
// CHECK-DAG: shl i16 %a, %[[TRUNC]]
Expand Down Expand Up @@ -51,7 +52,7 @@ pub unsafe fn unchecked_shr_signed_smaller(a: i16, b: u32) -> i16 {
// This uses -DAG to avoid failing on irrelevant reorderings,
// like emitting the truncation earlier.

// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 32768
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 16
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
// CHECK-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
// CHECK-DAG: ashr i16 %a, %[[TRUNC]]
Expand All @@ -66,3 +67,47 @@ pub unsafe fn unchecked_shr_signed_bigger(a: i64, b: u32) -> i64 {
// CHECK: ashr i64 %a, %[[EXT]]
a.unchecked_shr(b)
}

// CHECK-LABEL: @unchecked_shr_u128_i8
#[no_mangle]
pub unsafe fn unchecked_shr_u128_i8(a: u128, b: i8) -> u128 {
// CHECK-NOT: assume
// CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
// CHECK: lshr i128 %a, %[[EXT]]
std::intrinsics::unchecked_shr(a, b)
}

// CHECK-LABEL: @unchecked_shl_i128_u8
#[no_mangle]
pub unsafe fn unchecked_shl_i128_u8(a: i128, b: u8) -> i128 {
// CHECK-NOT: assume
// CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
// CHECK: shl i128 %a, %[[EXT]]
std::intrinsics::unchecked_shl(a, b)
}

// CHECK-LABEL: @unchecked_shl_u8_i128
#[no_mangle]
pub unsafe fn unchecked_shl_u8_i128(a: u8, b: i128) -> u8 {
// This uses -DAG to avoid failing on irrelevant reorderings,
// like emitting the truncation earlier.

// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
// CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
// CHECK-DAG: shl i8 %a, %[[TRUNC]]
std::intrinsics::unchecked_shl(a, b)
}

// CHECK-LABEL: @unchecked_shr_i8_u128
#[no_mangle]
pub unsafe fn unchecked_shr_i8_u128(a: i8, b: u128) -> i8 {
// This uses -DAG to avoid failing on irrelevant reorderings,
// like emitting the truncation earlier.

// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
// CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
// CHECK-DAG: ashr i8 %a, %[[TRUNC]]
std::intrinsics::unchecked_shr(a, b)
}
Loading

0 comments on commit 1b0e46f

Please sign in to comment.