Skip to content

Commit 97675af

Browse files
KmeakinKmeakin
committed
optimize: pow when base is a power of two
if base == 2 ** k, then (2 ** k) ** n == 2 ** (k * n) == 1 << (k * n)
1 parent 4a52fd0 commit 97675af

File tree

5 files changed

+127
-13
lines changed

5 files changed

+127
-13
lines changed

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

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1761,11 +1761,25 @@ macro_rules! int_impl {
17611761
without modifying the original"]
17621762
#[inline]
17631763
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
1764+
let mut base = self;
1765+
let mut acc: Self = 1;
1766+
1767+
if intrinsics::is_val_statically_known(base) {
1768+
if base.unsigned_abs().is_power_of_two() {
1769+
let k = base.unsigned_abs().ilog2();
1770+
let shift = try_opt!(k.checked_mul(exp));
1771+
let magnitude = try_opt!((1 as Self).checked_shl(shift));
1772+
return if base < 0 && (exp % 2) == 1 {
1773+
Some(magnitude.wrapping_neg())
1774+
} else {
1775+
Some(magnitude)
1776+
}
1777+
}
1778+
}
1779+
17641780
if exp == 0 {
17651781
return Some(1);
17661782
}
1767-
let mut base = self;
1768-
let mut acc: Self = 1;
17691783

17701784
if intrinsics::is_val_statically_known(exp) {
17711785
while exp > 1 {
@@ -3008,15 +3022,30 @@ macro_rules! int_impl {
30083022
without modifying the original"]
30093023
#[inline]
30103024
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3011-
if exp == 0 {
3012-
return (1, false);
3013-
}
3014-
30153025
let mut base = self;
30163026
let mut acc: Self = 1;
30173027
let mut overflow = false;
30183028
let mut tmp_overflow;
30193029

3030+
if intrinsics::is_val_statically_known(base) {
3031+
if base.unsigned_abs().is_power_of_two() {
3032+
let k = base.unsigned_abs().ilog2();
3033+
let Some(shift) = k.checked_mul(exp) else {
3034+
return (0, true)
3035+
};
3036+
let magnitude = (1 as Self).unbounded_shl(shift);
3037+
return if base < 0 && (exp % 2) == 1 {
3038+
(magnitude.wrapping_neg(), shift >= Self::BITS)
3039+
} else {
3040+
(magnitude, shift >= Self::BITS)
3041+
}
3042+
}
3043+
}
3044+
3045+
if exp == 0 {
3046+
return (1, false);
3047+
}
3048+
30203049
if intrinsics::is_val_statically_known(exp) {
30213050
while exp > 1 {
30223051
if (exp & 1) == 1 {

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

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2226,11 +2226,23 @@ macro_rules! uint_impl {
22262226
without modifying the original"]
22272227
#[inline]
22282228
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
2229+
let mut base = self;
2230+
let mut acc: Self = 1;
2231+
2232+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
2233+
// change of base:
2234+
// if base == 2 ** k, then
2235+
// (2 ** k) ** n
2236+
// == 2 ** (k * n)
2237+
// == 1 << (k * n)
2238+
let k = base.ilog2();
2239+
let shift = try_opt!(k.checked_mul(exp));
2240+
return (1 as Self).checked_shl(shift);
2241+
}
2242+
22292243
if exp == 0 {
22302244
return Some(1);
22312245
}
2232-
let mut base = self;
2233-
let mut acc: Self = 1;
22342246

22352247
if intrinsics::is_val_statically_known(exp) {
22362248
while exp > 1 {
@@ -3425,14 +3437,28 @@ macro_rules! uint_impl {
34253437
without modifying the original"]
34263438
#[inline]
34273439
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3428-
if exp == 0 {
3429-
return (1, false);
3430-
}
34313440
let mut base = self;
34323441
let mut acc: Self = 1;
34333442
let mut overflow = false;
34343443
let mut tmp_overflow;
34353444

3445+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
3446+
// change of base:
3447+
// if base == 2 ** k, then
3448+
// (2 ** k) ** n
3449+
// == 2 ** (k * n)
3450+
// == 1 << (k * n)
3451+
let k = base.ilog2();
3452+
let Some(shift) = k.checked_mul(exp) else {
3453+
return (0, true)
3454+
};
3455+
return ((1 as Self).unbounded_shl(shift), shift >= Self::BITS)
3456+
}
3457+
3458+
if exp == 0 {
3459+
return (1, false);
3460+
}
3461+
34363462
if intrinsics::is_val_statically_known(exp) {
34373463
while exp > 1 {
34383464
if (exp & 1) == 1 {

‎tests/codegen-llvm/pow_known_base.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
//@ compile-flags: -Copt-level=3
2+
// Test that `pow` can use a faster implementation when `base` is a
3+
// known power of two
4+
5+
#![crate_type = "lib"]
6+
7+
// 2 ** n == 2 ** (1 * n) == 1 << (1 * n)
8+
// CHECK-LABEL: @pow2
9+
#[no_mangle]
10+
pub fn pow2(exp: u32) -> u32 {
11+
// CHECK: %[[OVERFLOW:.+]] = icmp ult i32 %exp, 32
12+
// CHECK: %[[POW:.+]] = shl nuw i32 1, %exp
13+
// CHECK: %[[RET:.+]] = select i1 %[[OVERFLOW]], i32 %[[POW]], i32 0
14+
// CHECK: ret i32 %[[RET]]
15+
2u32.pow(exp)
16+
}
17+
18+
// 4 ** n == 2 ** (2 * n) == 1 << (2 * n)
19+
// CHECK-LABEL: @pow4
20+
#[no_mangle]
21+
pub fn pow4(exp: u32) -> u32 {
22+
// CHECK: %[[ICMP1:.+]] = icmp slt i32 %exp, 0
23+
// CHECK: %[[SHIFT_AMOUNT:.+]] = shl i32 %exp, 1
24+
// CHECK: %[[ICMP2:.+]] = icmp ult i32 %[[SHIFT_AMOUNT]], 32
25+
// CHECK: %[[POW:.+]] = shl nuw i32 1, %[[SHIFT_AMOUNT]]
26+
// CHECK: %[[SEL:.+]] = select i1 %[[ICMP2]], i32 %[[POW]], i32 0
27+
// CHECK: %[[RET:.+]] = select i1 %[[ICMP1]], i32 0, i32 %[[SEL]]
28+
// CHECK: ret i32 %[[RET]]
29+
4u32.pow(exp)
30+
}
31+
32+
// 16 ** n == 2 ** (4 * n) == 1 << (4 * n)
33+
// CHECK-LABEL: @pow16
34+
#[no_mangle]
35+
pub fn pow16(exp: u32) -> u32 {
36+
// CHECK: %[[ICMP1:.+]] = icmp ugt i32 %exp, 1073741823
37+
// CHECK: %[[SHIFT_AMOUNT:.+]] = shl i32 %exp, 2
38+
// CHECK: %[[ICMP2:.+]] = icmp ult i32 %[[SHIFT_AMOUNT]], 32
39+
// CHECK: %[[POW:.+]] = shl nuw i32 1, %[[SHIFT_AMOUNT]]
40+
// CHECK: %[[SEL:.+]] = select i1 %[[ICMP2]], i32 %[[POW]], i32 0
41+
// CHECK: %[[RET:.+]] = select i1 %[[ICMP1]], i32 0, i32 %[[SEL]]
42+
// CHECK: ret i32 %[[RET]]
43+
16u32.pow(exp)
44+
}
45+
46+
// (-2) ** n == (-2) ** (1 * n) == 1 << (1 * n)
47+
// CHECK-LABEL: @pow_minus_2
48+
#[no_mangle]
49+
pub fn pow_minus_2(exp: u32) -> i32 {
50+
// CHECK: %[[OVERFLOW:.+]] = icmp ult i32 %exp, 32
51+
// CHECK: %[[POW:.+]] = shl nuw i32 1, %exp
52+
// CHECK: %[[MAGNITUDE:.+]] = select i1 %[[OVERFLOW]], i32 %[[POW]], i32 0
53+
// CHECK: %[[IS_ODD:.+]] = and i32 %exp, 1
54+
// CHECK: %[[IS_EVEN:.+]] = icmp eq i32 %[[IS_ODD]], 0
55+
// CHECK: %[[NEG:.+]] = sub i32 0, %[[MAGNITUDE]]
56+
// CHECK: %[[RET:.+]] = select i1 %[[IS_EVEN]], i32 %[[MAGNITUDE]], i32 %[[NEG]]
57+
// CHECK: ret i32 %[[RET]]
58+
(-2i32).pow(exp)
59+
}

‎tests/ui/numbers-arithmetic/overflowing-pow-signed.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
//@ run-fail
22
//@ regex-error-pattern: thread 'main'.*panicked
3-
//@ error-pattern: attempt to multiply with overflow
3+
//@ regex-error-pattern: attempt to calculate the power with overflow
44
//@ needs-subprocess
55
//@ compile-flags: -C debug-assertions
66

‎tests/ui/numbers-arithmetic/overflowing-pow-unsigned.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
//@ run-fail
22
//@ regex-error-pattern: thread 'main'.*panicked
3-
//@ error-pattern: attempt to multiply with overflow
3+
//@ regex-error-pattern: attempt to calculate the power with overflow
44
//@ needs-subprocess
55
//@ compile-flags: -C debug-assertions
66

0 commit comments

Comments
 (0)