Skip to content

Commit 1b0e46f

Browse files
matthiaskrgrmatthiaskrgr
authoredApr 2, 2024
Rollup merge of #123226 - scottmcm:u32-shifts, r=WaffleLapkin
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
2 parents d63ddef + 4626521 commit 1b0e46f

File tree

32 files changed

+408
-576
lines changed

32 files changed

+408
-576
lines changed
 

‎compiler/rustc_codegen_ssa/src/base.rs

+31-4
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::back::write::{
55
compute_per_cgu_lto_type, start_async_codegen, submit_codegened_module_to_llvm,
66
submit_post_lto_module_to_llvm, submit_pre_lto_module_to_llvm, ComputedLtoType, OngoingCodegen,
77
};
8-
use crate::common::{IntPredicate, RealPredicate, TypeKind};
8+
use crate::common::{self, IntPredicate, RealPredicate, TypeKind};
99
use crate::errors;
1010
use crate::meth;
1111
use crate::mir;
@@ -33,7 +33,7 @@ use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder, MonoItem};
3333
use rustc_middle::query::Providers;
3434
use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout};
3535
use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
36-
use rustc_session::config::{self, CrateType, EntryFnType, OutputType};
36+
use rustc_session::config::{self, CrateType, EntryFnType, OptLevel, OutputType};
3737
use rustc_session::Session;
3838
use rustc_span::symbol::sym;
3939
use rustc_span::Symbol;
@@ -300,14 +300,35 @@ pub fn coerce_unsized_into<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
300300
}
301301
}
302302

303-
pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
303+
/// Returns `rhs` sufficiently masked, truncated, and/or extended so that
304+
/// it can be used to shift `lhs`.
305+
///
306+
/// Shifts in MIR are all allowed to have mismatched LHS & RHS types.
307+
/// The shift methods in `BuilderMethods`, however, are fully homogeneous
308+
/// (both parameters and the return type are all the same type).
309+
///
310+
/// If `is_unchecked` is false, this masks the RHS to ensure it stays in-bounds,
311+
/// as the `BuilderMethods` shifts are UB for out-of-bounds shift amounts.
312+
/// For 32- and 64-bit types, this matches the semantics
313+
/// of Java. (See related discussion on #1877 and #10183.)
314+
///
315+
/// If `is_unchecked` is true, this does no masking, and adds sufficient `assume`
316+
/// calls or operation flags to preserve as much freedom to optimize as possible.
317+
pub fn build_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
304318
bx: &mut Bx,
305319
lhs: Bx::Value,
306-
rhs: Bx::Value,
320+
mut rhs: Bx::Value,
321+
is_unchecked: bool,
307322
) -> Bx::Value {
308323
// Shifts may have any size int on the rhs
309324
let mut rhs_llty = bx.cx().val_ty(rhs);
310325
let mut lhs_llty = bx.cx().val_ty(lhs);
326+
327+
let mask = common::shift_mask_val(bx, lhs_llty, rhs_llty, false);
328+
if !is_unchecked {
329+
rhs = bx.and(rhs, mask);
330+
}
331+
311332
if bx.cx().type_kind(rhs_llty) == TypeKind::Vector {
312333
rhs_llty = bx.cx().element_type(rhs_llty)
313334
}
@@ -317,6 +338,12 @@ pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
317338
let rhs_sz = bx.cx().int_width(rhs_llty);
318339
let lhs_sz = bx.cx().int_width(lhs_llty);
319340
if lhs_sz < rhs_sz {
341+
if is_unchecked && bx.sess().opts.optimize != OptLevel::No {
342+
// FIXME: Use `trunc nuw` once that's available
343+
let inrange = bx.icmp(IntPredicate::IntULE, rhs, mask);
344+
bx.assume(inrange);
345+
}
346+
320347
bx.trunc(rhs, lhs_llty)
321348
} else if lhs_sz > rhs_sz {
322349
// We zero-extend even if the RHS is signed. So e.g. `(x: i32) << -1i8` will zero-extend the

‎compiler/rustc_codegen_ssa/src/common.rs

+1-40
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,9 @@
33
use rustc_hir::LangItem;
44
use rustc_middle::mir;
55
use rustc_middle::ty::Instance;
6-
use rustc_middle::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
6+
use rustc_middle::ty::{self, layout::TyAndLayout, TyCtxt};
77
use rustc_span::Span;
88

9-
use crate::base;
109
use crate::traits::*;
1110

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

131-
// To avoid UB from LLVM, these two functions mask RHS with an
132-
// appropriate mask unconditionally (i.e., the fallback behavior for
133-
// all shifts). For 32- and 64-bit types, this matches the semantics
134-
// of Java. (See related discussion on #1877 and #10183.)
135-
136-
pub fn build_masked_lshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
137-
bx: &mut Bx,
138-
lhs: Bx::Value,
139-
rhs: Bx::Value,
140-
) -> Bx::Value {
141-
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
142-
// #1877, #10183: Ensure that input is always valid
143-
let rhs = shift_mask_rhs(bx, rhs);
144-
bx.shl(lhs, rhs)
145-
}
146-
147-
pub fn build_masked_rshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
148-
bx: &mut Bx,
149-
lhs_t: Ty<'tcx>,
150-
lhs: Bx::Value,
151-
rhs: Bx::Value,
152-
) -> Bx::Value {
153-
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
154-
// #1877, #10183: Ensure that input is always valid
155-
let rhs = shift_mask_rhs(bx, rhs);
156-
let is_signed = lhs_t.is_signed();
157-
if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
158-
}
159-
160-
fn shift_mask_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
161-
bx: &mut Bx,
162-
rhs: Bx::Value,
163-
) -> Bx::Value {
164-
let rhs_llty = bx.val_ty(rhs);
165-
let shift_val = shift_mask_val(bx, rhs_llty, rhs_llty, false);
166-
bx.and(rhs, shift_val)
167-
}
168-
169130
pub fn shift_mask_val<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
170131
bx: &mut Bx,
171132
llty: Bx::Type,

‎compiler/rustc_codegen_ssa/src/mir/rvalue.rs

+5-7
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use super::place::PlaceRef;
33
use super::{FunctionCx, LocalRef};
44

55
use crate::base;
6-
use crate::common::{self, IntPredicate};
6+
use crate::common::IntPredicate;
77
use crate::traits::*;
88
use crate::MemFlags;
99

@@ -860,14 +860,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
860860
bx.inbounds_gep(llty, lhs, &[rhs])
861861
}
862862
}
863-
mir::BinOp::Shl => common::build_masked_lshift(bx, lhs, rhs),
864-
mir::BinOp::ShlUnchecked => {
865-
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
863+
mir::BinOp::Shl | mir::BinOp::ShlUnchecked => {
864+
let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShlUnchecked);
866865
bx.shl(lhs, rhs)
867866
}
868-
mir::BinOp::Shr => common::build_masked_rshift(bx, input_ty, lhs, rhs),
869-
mir::BinOp::ShrUnchecked => {
870-
let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
867+
mir::BinOp::Shr | mir::BinOp::ShrUnchecked => {
868+
let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShrUnchecked);
871869
if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
872870
}
873871
mir::BinOp::Ne

‎compiler/rustc_hir_analysis/src/check/intrinsic.rs

+2-3
Original file line numberDiff line numberDiff line change
@@ -454,9 +454,8 @@ pub fn check_intrinsic_type(
454454
sym::unchecked_div | sym::unchecked_rem | sym::exact_div => {
455455
(1, 0, vec![param(0), param(0)], param(0))
456456
}
457-
sym::unchecked_shl | sym::unchecked_shr | sym::rotate_left | sym::rotate_right => {
458-
(1, 0, vec![param(0), param(0)], param(0))
459-
}
457+
sym::unchecked_shl | sym::unchecked_shr => (2, 0, vec![param(0), param(1)], param(0)),
458+
sym::rotate_left | sym::rotate_right => (1, 0, vec![param(0), param(0)], param(0)),
460459
sym::unchecked_add | sym::unchecked_sub | sym::unchecked_mul => {
461460
(1, 0, vec![param(0), param(0)], param(0))
462461
}

‎library/core/src/intrinsics.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -2224,18 +2224,20 @@ extern "rust-intrinsic" {
22242224
/// Safe wrappers for this intrinsic are available on the integer
22252225
/// primitives via the `checked_shl` method. For example,
22262226
/// [`u32::checked_shl`]
2227+
#[cfg(not(bootstrap))]
22272228
#[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
22282229
#[rustc_nounwind]
2229-
pub fn unchecked_shl<T: Copy>(x: T, y: T) -> T;
2230+
pub fn unchecked_shl<T: Copy, U: Copy>(x: T, y: U) -> T;
22302231
/// Performs an unchecked right shift, resulting in undefined behavior when
22312232
/// `y < 0` or `y >= N`, where N is the width of T in bits.
22322233
///
22332234
/// Safe wrappers for this intrinsic are available on the integer
22342235
/// primitives via the `checked_shr` method. For example,
22352236
/// [`u32::checked_shr`]
2237+
#[cfg(not(bootstrap))]
22362238
#[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
22372239
#[rustc_nounwind]
2238-
pub fn unchecked_shr<T: Copy>(x: T, y: T) -> T;
2240+
pub fn unchecked_shr<T: Copy, U: Copy>(x: T, y: U) -> T;
22392241

22402242
/// Returns the result of an unchecked addition, resulting in
22412243
/// undefined behavior when `x + y > T::MAX` or `x + y < T::MIN`.

‎library/core/src/num/int_macros.rs

+24-8
Original file line numberDiff line numberDiff line change
@@ -1227,10 +1227,18 @@ macro_rules! int_impl {
12271227
#[inline(always)]
12281228
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
12291229
pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
1230-
// SAFETY: the caller must uphold the safety contract for
1231-
// `unchecked_shl`.
1232-
// Any legal shift amount is losslessly representable in the self type.
1233-
unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
1230+
#[cfg(bootstrap)]
1231+
{
1232+
// For bootstrapping, just use built-in primitive shift.
1233+
// panicking is a legal manifestation of UB
1234+
self << rhs
1235+
}
1236+
#[cfg(not(bootstrap))]
1237+
{
1238+
// SAFETY: the caller must uphold the safety contract for
1239+
// `unchecked_shl`.
1240+
unsafe { intrinsics::unchecked_shl(self, rhs) }
1241+
}
12341242
}
12351243

12361244
/// Checked shift right. Computes `self >> rhs`, returning `None` if `rhs` is
@@ -1310,10 +1318,18 @@ macro_rules! int_impl {
13101318
#[inline(always)]
13111319
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
13121320
pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
1313-
// SAFETY: the caller must uphold the safety contract for
1314-
// `unchecked_shr`.
1315-
// Any legal shift amount is losslessly representable in the self type.
1316-
unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
1321+
#[cfg(bootstrap)]
1322+
{
1323+
// For bootstrapping, just use built-in primitive shift.
1324+
// panicking is a legal manifestation of UB
1325+
self >> rhs
1326+
}
1327+
#[cfg(not(bootstrap))]
1328+
{
1329+
// SAFETY: the caller must uphold the safety contract for
1330+
// `unchecked_shr`.
1331+
unsafe { intrinsics::unchecked_shr(self, rhs) }
1332+
}
13171333
}
13181334

13191335
/// Checked absolute value. Computes `self.abs()`, returning `None` if

‎library/core/src/num/mod.rs

-11
Original file line numberDiff line numberDiff line change
@@ -285,17 +285,6 @@ macro_rules! widening_impl {
285285
};
286286
}
287287

288-
macro_rules! conv_rhs_for_unchecked_shift {
289-
($SelfT:ty, $x:expr) => {{
290-
// If the `as` cast will truncate, ensure we still tell the backend
291-
// that the pre-truncation value was also small.
292-
if <$SelfT>::BITS < 32 {
293-
intrinsics::assume($x <= (<$SelfT>::MAX as u32));
294-
}
295-
$x as $SelfT
296-
}};
297-
}
298-
299288
impl i8 {
300289
int_impl! {
301290
Self = i8,

‎library/core/src/num/uint_macros.rs

+24-8
Original file line numberDiff line numberDiff line change
@@ -1286,10 +1286,18 @@ macro_rules! uint_impl {
12861286
#[inline(always)]
12871287
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
12881288
pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
1289-
// SAFETY: the caller must uphold the safety contract for
1290-
// `unchecked_shl`.
1291-
// Any legal shift amount is losslessly representable in the self type.
1292-
unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
1289+
#[cfg(bootstrap)]
1290+
{
1291+
// For bootstrapping, just use built-in primitive shift.
1292+
// panicking is a legal manifestation of UB
1293+
self << rhs
1294+
}
1295+
#[cfg(not(bootstrap))]
1296+
{
1297+
// SAFETY: the caller must uphold the safety contract for
1298+
// `unchecked_shl`.
1299+
unsafe { intrinsics::unchecked_shl(self, rhs) }
1300+
}
12931301
}
12941302

12951303
/// Checked shift right. Computes `self >> rhs`, returning `None`
@@ -1369,10 +1377,18 @@ macro_rules! uint_impl {
13691377
#[inline(always)]
13701378
#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
13711379
pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
1372-
// SAFETY: the caller must uphold the safety contract for
1373-
// `unchecked_shr`.
1374-
// Any legal shift amount is losslessly representable in the self type.
1375-
unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
1380+
#[cfg(bootstrap)]
1381+
{
1382+
// For bootstrapping, just use built-in primitive shift.
1383+
// panicking is a legal manifestation of UB
1384+
self >> rhs
1385+
}
1386+
#[cfg(not(bootstrap))]
1387+
{
1388+
// SAFETY: the caller must uphold the safety contract for
1389+
// `unchecked_shr`.
1390+
unsafe { intrinsics::unchecked_shr(self, rhs) }
1391+
}
13761392
}
13771393

13781394
/// Checked exponentiation. Computes `self.pow(exp)`, returning `None` if

‎library/core/src/ptr/mod.rs

+12-2
Original file line numberDiff line numberDiff line change
@@ -1781,9 +1781,19 @@ pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usiz
17811781
// FIXME(#75598): Direct use of these intrinsics improves codegen significantly at opt-level <=
17821782
// 1, where the method versions of these operations are not inlined.
17831783
use intrinsics::{
1784-
assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_shl,
1785-
unchecked_shr, unchecked_sub, wrapping_add, wrapping_mul, wrapping_sub,
1784+
assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_sub,
1785+
wrapping_add, wrapping_mul, wrapping_sub,
17861786
};
1787+
#[cfg(bootstrap)]
1788+
const unsafe fn unchecked_shl(value: usize, shift: usize) -> usize {
1789+
value << shift
1790+
}
1791+
#[cfg(bootstrap)]
1792+
const unsafe fn unchecked_shr(value: usize, shift: usize) -> usize {
1793+
value >> shift
1794+
}
1795+
#[cfg(not(bootstrap))]
1796+
use intrinsics::{unchecked_shl, unchecked_shr};
17871797

17881798
/// Calculate multiplicative modular inverse of `x` modulo `m`.
17891799
///

‎tests/codegen/unchecked_shifts.rs

+47-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#![crate_type = "lib"]
44
#![feature(unchecked_shifts)]
5+
#![feature(core_intrinsics)]
56

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

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

54-
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 32768
55+
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 16
5556
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
5657
// CHECK-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
5758
// CHECK-DAG: ashr i16 %a, %[[TRUNC]]
@@ -66,3 +67,47 @@ pub unsafe fn unchecked_shr_signed_bigger(a: i64, b: u32) -> i64 {
6667
// CHECK: ashr i64 %a, %[[EXT]]
6768
a.unchecked_shr(b)
6869
}
70+
71+
// CHECK-LABEL: @unchecked_shr_u128_i8
72+
#[no_mangle]
73+
pub unsafe fn unchecked_shr_u128_i8(a: u128, b: i8) -> u128 {
74+
// CHECK-NOT: assume
75+
// CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
76+
// CHECK: lshr i128 %a, %[[EXT]]
77+
std::intrinsics::unchecked_shr(a, b)
78+
}
79+
80+
// CHECK-LABEL: @unchecked_shl_i128_u8
81+
#[no_mangle]
82+
pub unsafe fn unchecked_shl_i128_u8(a: i128, b: u8) -> i128 {
83+
// CHECK-NOT: assume
84+
// CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
85+
// CHECK: shl i128 %a, %[[EXT]]
86+
std::intrinsics::unchecked_shl(a, b)
87+
}
88+
89+
// CHECK-LABEL: @unchecked_shl_u8_i128
90+
#[no_mangle]
91+
pub unsafe fn unchecked_shl_u8_i128(a: u8, b: i128) -> u8 {
92+
// This uses -DAG to avoid failing on irrelevant reorderings,
93+
// like emitting the truncation earlier.
94+
95+
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
96+
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
97+
// CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
98+
// CHECK-DAG: shl i8 %a, %[[TRUNC]]
99+
std::intrinsics::unchecked_shl(a, b)
100+
}
101+
102+
// CHECK-LABEL: @unchecked_shr_i8_u128
103+
#[no_mangle]
104+
pub unsafe fn unchecked_shr_i8_u128(a: i8, b: u128) -> i8 {
105+
// This uses -DAG to avoid failing on irrelevant reorderings,
106+
// like emitting the truncation earlier.
107+
108+
// CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
109+
// CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
110+
// CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
111+
// CHECK-DAG: ashr i8 %a, %[[TRUNC]]
112+
std::intrinsics::unchecked_shr(a, b)
113+
}

0 commit comments

Comments
 (0)