Skip to content

Make checked ops emit *unchecked* LLVM operations where feasible #124114

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Apr 20, 2024
Merged

Make checked ops emit *unchecked* LLVM operations where feasible #124114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions library/core/src/num/int_macros.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1163,12 +1163,19 @@ macro_rules! int_impl {
/// ```
#[stable(feature = "wrapping", since = "1.7.0")]
#[rustc_const_stable(feature = "const_checked_int_methods", since = "1.47.0")]
// We could always go back to wrapping
#[rustc_allow_const_fn_unstable(unchecked_shifts)]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
pub const fn checked_shl(self, rhs: u32) -> Option<Self> {
let (a, b) = self.overflowing_shl(rhs);
if unlikely!(b) { None } else { Some(a) }
// Not using overflowing_shl as that's a wrapping shift
if rhs < Self::BITS {
// SAFETY: just checked the RHS is in-range
Some(unsafe { self.unchecked_shl(rhs) })
} else {
None
}
}

/// Strict shift left. Computes `self << rhs`, panicking if `rhs` is larger
Expand Down Expand Up @@ -1254,12 +1261,19 @@ macro_rules! int_impl {
/// ```
#[stable(feature = "wrapping", since = "1.7.0")]
#[rustc_const_stable(feature = "const_checked_int_methods", since = "1.47.0")]
// We could always go back to wrapping
#[rustc_allow_const_fn_unstable(unchecked_shifts)]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
pub const fn checked_shr(self, rhs: u32) -> Option<Self> {
let (a, b) = self.overflowing_shr(rhs);
if unlikely!(b) { None } else { Some(a) }
// Not using overflowing_shr as that's a wrapping shift
if rhs < Self::BITS {
// SAFETY: just checked the RHS is in-range
Some(unsafe { self.unchecked_shr(rhs) })
} else {
None
}
}

/// Strict shift right. Computes `self >> rhs`, panicking `rhs` is
Expand Down
35 changes: 29 additions & 6 deletions library/core/src/num/uint_macros.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,8 +579,17 @@ macro_rules! uint_impl {
without modifying the original"]
#[inline]
pub const fn checked_sub(self, rhs: Self) -> Option<Self> {
let (a, b) = self.overflowing_sub(rhs);
if unlikely!(b) { None } else { Some(a) }
// Per PR#103299, there's no advantage to the `overflowing` intrinsic
// for *unsigned* subtraction and we just emit the manual check anyway.
// Thus, rather than using `overflowing_sub` that produces a wrapping
// subtraction, check it ourself so we can use an unchecked one.

if self >= rhs {
// SAFETY: just checked this can't overflow
Some(unsafe { intrinsics::unchecked_sub(self, rhs) })
} else {
None
}
}

/// Strict integer subtraction. Computes `self - rhs`, panicking if
Expand Down Expand Up @@ -1222,12 +1231,19 @@ macro_rules! uint_impl {
/// ```
#[stable(feature = "wrapping", since = "1.7.0")]
#[rustc_const_stable(feature = "const_checked_int_methods", since = "1.47.0")]
// We could always go back to wrapping
#[rustc_allow_const_fn_unstable(unchecked_shifts)]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
pub const fn checked_shl(self, rhs: u32) -> Option<Self> {
let (a, b) = self.overflowing_shl(rhs);
if unlikely!(b) { None } else { Some(a) }
// Not using overflowing_shl as that's a wrapping shift
if rhs < Self::BITS {
// SAFETY: just checked the RHS is in-range
Some(unsafe { self.unchecked_shl(rhs) })
} else {
None
}
}

/// Strict shift left. Computes `self << rhs`, panicking if `rhs` is larger
Expand Down Expand Up @@ -1313,12 +1329,19 @@ macro_rules! uint_impl {
/// ```
#[stable(feature = "wrapping", since = "1.7.0")]
#[rustc_const_stable(feature = "const_checked_int_methods", since = "1.47.0")]
// We could always go back to wrapping
#[rustc_allow_const_fn_unstable(unchecked_shifts)]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
pub const fn checked_shr(self, rhs: u32) -> Option<Self> {
let (a, b) = self.overflowing_shr(rhs);
if unlikely!(b) { None } else { Some(a) }
// Not using overflowing_shr as that's a wrapping shift
if rhs < Self::BITS {
// SAFETY: just checked the RHS is in-range
Some(unsafe { self.unchecked_shr(rhs) })
} else {
None
}
}

/// Strict shift right. Computes `self >> rhs`, panicking `rhs` is
Expand Down
86 changes: 86 additions & 0 deletions tests/codegen/checked_math.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
//@ compile-flags: -O -Z merge-functions=disabled

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

// Because the result of something like `u32::checked_sub` can only be used if it
// didn't overflow, make sure that LLVM actually knows that in optimized builds.
// Thanks to poison semantics, this doesn't even need branches.

// CHECK-LABEL: @checked_sub_unsigned
// CHECK-SAME: (i16 noundef %a, i16 noundef %b)
#[no_mangle]
pub fn checked_sub_unsigned(a: u16, b: u16) -> Option<u16> {
// CHECK-DAG: %[[IS_SOME:.+]] = icmp uge i16 %a, %b
// CHECK-DAG: %[[DIFF_P:.+]] = sub nuw i16 %a, %b
// CHECK-DAG: %[[DISCR:.+]] = zext i1 %[[IS_SOME]] to i16
// CHECK-DAG: %[[DIFF_U:.+]] = select i1 %[[IS_SOME]], i16 %[[DIFF_P]], i16 undef

// CHECK: %[[R0:.+]] = insertvalue { i16, i16 } poison, i16 %[[DISCR]], 0
// CHECK: %[[R1:.+]] = insertvalue { i16, i16 } %[[R0]], i16 %[[DIFF_U]], 1
// CHECK: ret { i16, i16 } %[[R1]]
a.checked_sub(b)
}

// Note that `shl` and `shr` in LLVM are already unchecked. So rather than
// looking for no-wrap flags, we just need there to not be any masking.
Comment on lines +25 to +26
Copy link
Member Author

Choose a reason for hiding this comment

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

You can see the masking in optimized LLVM IR and optimized MIR today: https://rust.godbolt.org/z/x5Podh9z5

Copy link
Member

Choose a reason for hiding this comment

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

looks like the same assembly emitted, tho'.

Copy link
Member Author

@scottmcm scottmcm Apr 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

Well, on x86 the shift instructions are wrapping. To see a difference you'd need to make an example where LLVM can optimize based on knowing that the shift is in-range.

(Just like how sub and sub nuw have the same assembly too.)

Copy link
Member

@workingjubilee workingjubilee Apr 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

That's a good point! I tried an architecture that actually has a slightly more annoying shift instruction in this regard, PowerPC! (specifically --target=powerpc64le-unknown-linux-gnu)

old_checked_sub:
        clrldi  5, 4, 32
        clrlwi  4, 4, 27
        addi 5, 5, -32
        srw 4, 3, 4
        rldicl 5, 5, 1, 63
        mr      3, 5
        blr
        .long   0
        .quad   0

new_checked_sub:
        clrldi  5, 4, 32
        slw 4, 3, 4
        addi 5, 5, -32
        rldicl 5, 5, 1, 63
        mr      3, 5
        blr
        .long   0
        .quad   0

One less instruction!


// CHECK-LABEL: @checked_shl_unsigned
// CHECK-SAME: (i32 noundef %a, i32 noundef %b)
#[no_mangle]
pub fn checked_shl_unsigned(a: u32, b: u32) -> Option<u32> {
// CHECK-DAG: %[[IS_SOME:.+]] = icmp ult i32 %b, 32
// CHECK-DAG: %[[SHIFTED_P:.+]] = shl i32 %a, %b
// CHECK-DAG: %[[DISCR:.+]] = zext i1 %[[IS_SOME]] to i32
// CHECK-DAG: %[[SHIFTED_U:.+]] = select i1 %[[IS_SOME]], i32 %[[SHIFTED_P]], i32 undef

// CHECK: %[[R0:.+]] = insertvalue { i32, i32 } poison, i32 %[[DISCR]], 0
// CHECK: %[[R1:.+]] = insertvalue { i32, i32 } %[[R0]], i32 %[[SHIFTED_U]], 1
// CHECK: ret { i32, i32 } %[[R1]]
a.checked_shl(b)
}

// CHECK-LABEL: @checked_shr_unsigned
// CHECK-SAME: (i32 noundef %a, i32 noundef %b)
#[no_mangle]
pub fn checked_shr_unsigned(a: u32, b: u32) -> Option<u32> {
// CHECK-DAG: %[[IS_SOME:.+]] = icmp ult i32 %b, 32
// CHECK-DAG: %[[SHIFTED_P:.+]] = lshr i32 %a, %b
// CHECK-DAG: %[[DISCR:.+]] = zext i1 %[[IS_SOME]] to i32
// CHECK-DAG: %[[SHIFTED_U:.+]] = select i1 %[[IS_SOME]], i32 %[[SHIFTED_P]], i32 undef

// CHECK: %[[R0:.+]] = insertvalue { i32, i32 } poison, i32 %[[DISCR]], 0
// CHECK: %[[R1:.+]] = insertvalue { i32, i32 } %[[R0]], i32 %[[SHIFTED_U]], 1
// CHECK: ret { i32, i32 } %[[R1]]
a.checked_shr(b)
}

// CHECK-LABEL: @checked_shl_signed
// CHECK-SAME: (i32 noundef %a, i32 noundef %b)
#[no_mangle]
pub fn checked_shl_signed(a: i32, b: u32) -> Option<i32> {
// CHECK-DAG: %[[IS_SOME:.+]] = icmp ult i32 %b, 32
// CHECK-DAG: %[[SHIFTED_P:.+]] = shl i32 %a, %b
// CHECK-DAG: %[[DISCR:.+]] = zext i1 %[[IS_SOME]] to i32
// CHECK-DAG: %[[SHIFTED_U:.+]] = select i1 %[[IS_SOME]], i32 %[[SHIFTED_P]], i32 undef

// CHECK: %[[R0:.+]] = insertvalue { i32, i32 } poison, i32 %[[DISCR]], 0
// CHECK: %[[R1:.+]] = insertvalue { i32, i32 } %[[R0]], i32 %[[SHIFTED_U]], 1
// CHECK: ret { i32, i32 } %[[R1]]
a.checked_shl(b)
}

// CHECK-LABEL: @checked_shr_signed
// CHECK-SAME: (i32 noundef %a, i32 noundef %b)
#[no_mangle]
pub fn checked_shr_signed(a: i32, b: u32) -> Option<i32> {
// CHECK-DAG: %[[IS_SOME:.+]] = icmp ult i32 %b, 32
// CHECK-DAG: %[[SHIFTED_P:.+]] = ashr i32 %a, %b
// CHECK-DAG: %[[DISCR:.+]] = zext i1 %[[IS_SOME]] to i32
// CHECK-DAG: %[[SHIFTED_U:.+]] = select i1 %[[IS_SOME]], i32 %[[SHIFTED_P]], i32 undef

// CHECK: %[[R0:.+]] = insertvalue { i32, i32 } poison, i32 %[[DISCR]], 0
// CHECK: %[[R1:.+]] = insertvalue { i32, i32 } %[[R0]], i32 %[[SHIFTED_U]], 1
// CHECK: ret { i32, i32 } %[[R1]]
a.checked_shr(b)
}
53 changes: 13 additions & 40 deletions tests/mir-opt/pre-codegen/checked_ops.checked_shl.PreCodegen.after.panic-abort.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,60 +5,33 @@ fn checked_shl(_1: u32, _2: u32) -> Option<u32> {
debug rhs => _2;
let mut _0: std::option::Option<u32>;
scope 1 (inlined core::num::<impl u32>::checked_shl) {
debug self => _1;
debug rhs => _2;
let mut _6: bool;
scope 2 {
debug a => _4;
debug b => _5;
}
scope 3 (inlined core::num::<impl u32>::overflowing_shl) {
debug self => _1;
debug rhs => _2;
let mut _4: u32;
let mut _5: bool;
scope 4 (inlined core::num::<impl u32>::wrapping_shl) {
debug self => _1;
debug rhs => _2;
let mut _3: u32;
scope 5 (inlined core::num::<impl u32>::unchecked_shl) {
debug self => _1;
debug rhs => _3;
}
}
let mut _3: bool;
let mut _4: u32;
scope 2 (inlined core::num::<impl u32>::unchecked_shl) {
}
}

bb0: {
StorageLive(_4);
StorageLive(_5);
StorageLive(_3);
_3 = BitAnd(_2, const 31_u32);
_4 = ShlUnchecked(_1, _3);
StorageDead(_3);
_5 = Ge(_2, const core::num::<impl u32>::BITS);
StorageLive(_6);
_6 = unlikely(move _5) -> [return: bb1, unwind unreachable];
_3 = Lt(_2, const core::num::<impl u32>::BITS);
switchInt(move _3) -> [0: bb1, otherwise: bb2];
}

bb1: {
switchInt(move _6) -> [0: bb2, otherwise: bb3];
_0 = const Option::<u32>::None;
goto -> bb3;
}

bb2: {
_0 = Option::<u32>::Some(_4);
goto -> bb4;
StorageLive(_4);
_4 = ShlUnchecked(_1, _2);
_0 = Option::<u32>::Some(move _4);
StorageDead(_4);
goto -> bb3;
}

bb3: {
_0 = const Option::<u32>::None;
goto -> bb4;
}

bb4: {
StorageDead(_6);
StorageDead(_5);
StorageDead(_4);
StorageDead(_3);
return;
}
}
Expand Down
53 changes: 13 additions & 40 deletions tests/mir-opt/pre-codegen/checked_ops.checked_shl.PreCodegen.after.panic-unwind.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,60 +5,33 @@ fn checked_shl(_1: u32, _2: u32) -> Option<u32> {
debug rhs => _2;
let mut _0: std::option::Option<u32>;
scope 1 (inlined core::num::<impl u32>::checked_shl) {
debug self => _1;
debug rhs => _2;
let mut _6: bool;
scope 2 {
debug a => _4;
debug b => _5;
}
scope 3 (inlined core::num::<impl u32>::overflowing_shl) {
debug self => _1;
debug rhs => _2;
let mut _4: u32;
let mut _5: bool;
scope 4 (inlined core::num::<impl u32>::wrapping_shl) {
debug self => _1;
debug rhs => _2;
let mut _3: u32;
scope 5 (inlined core::num::<impl u32>::unchecked_shl) {
debug self => _1;
debug rhs => _3;
}
}
let mut _3: bool;
let mut _4: u32;
scope 2 (inlined core::num::<impl u32>::unchecked_shl) {
}
}

bb0: {
StorageLive(_4);
StorageLive(_5);
StorageLive(_3);
_3 = BitAnd(_2, const 31_u32);
_4 = ShlUnchecked(_1, _3);
StorageDead(_3);
_5 = Ge(_2, const core::num::<impl u32>::BITS);
StorageLive(_6);
_6 = unlikely(move _5) -> [return: bb1, unwind unreachable];
_3 = Lt(_2, const core::num::<impl u32>::BITS);
switchInt(move _3) -> [0: bb1, otherwise: bb2];
}

bb1: {
switchInt(move _6) -> [0: bb2, otherwise: bb3];
_0 = const Option::<u32>::None;
goto -> bb3;
}

bb2: {
_0 = Option::<u32>::Some(_4);
goto -> bb4;
StorageLive(_4);
_4 = ShlUnchecked(_1, _2);
_0 = Option::<u32>::Some(move _4);
StorageDead(_4);
goto -> bb3;
}

bb3: {
_0 = const Option::<u32>::None;
goto -> bb4;
}

bb4: {
StorageDead(_6);
StorageDead(_5);
StorageDead(_4);
StorageDead(_3);
return;
}
}
Expand Down
2 changes: 1 addition & 1 deletion tests/mir-opt/pre-codegen/checked_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
// skip-filecheck
//@ compile-flags: -O -Zmir-opt-level=2 -Cdebuginfo=2
Copy link
Member

Choose a reason for hiding this comment

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

why this change?

Copy link
Member Author

Choose a reason for hiding this comment

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

Well, first I was going to say for consistency with the other pre-codegen tests, but it looks like they're not as consistent as I thought. The mem-replace and range-iter tests don't ask for debuginfo, but then the slice-filter and chained-comparison ones do :/

I guess it's because I think it's more common for -O to be without debuginfo, and in general I don't think that the pre-codegen tests are about looking at debuginfo. Certainly in a normal --debug build the MIR inliner doesn't run so it'd look nothing like this.

I could put it back as it was if you'd like. It just adds 4 more debug lines; nothing interesting.

Copy link
Member

Choose a reason for hiding this comment

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

I don't mind it, either way, tbh, I just wanted to know the reasoning so that I know it wasn't Perfectly Random.

//@ compile-flags: -O -Zmir-opt-level=2
// EMIT_MIR_FOR_EACH_PANIC_STRATEGY

#![crate_type = "lib"]
Expand Down
Loading