Skip to content

Commit f20dfa8

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 b44090a commit f20dfa8

3 files changed

Lines changed: 121 additions & 11 deletions

File tree

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

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1845,11 +1845,23 @@ macro_rules! int_impl {
18451845
without modifying the original"]
18461846
#[inline]
18471847
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
1848+
let mut base = self;
1849+
let mut acc: Self = 1;
1850+
1851+
if intrinsics::is_val_statically_known(base) && base.unsigned_abs().is_power_of_two() {
1852+
let k = base.unsigned_abs().ilog2();
1853+
let shift = try_opt!(k.checked_mul(exp));
1854+
let magnitude = try_opt!((1 as Self).checked_shl(shift));
1855+
return if base < 0 && (exp % 2) == 1 {
1856+
Some(magnitude.wrapping_neg())
1857+
} else {
1858+
Some(magnitude)
1859+
}
1860+
}
1861+
18481862
if exp == 0 {
18491863
return Some(1);
18501864
}
1851-
let mut base = self;
1852-
let mut acc: Self = 1;
18531865

18541866
if intrinsics::is_val_statically_known(exp) {
18551867
while exp > 1 {
@@ -3096,15 +3108,28 @@ macro_rules! int_impl {
30963108
without modifying the original"]
30973109
#[inline]
30983110
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3099-
if exp == 0 {
3100-
return (1, false);
3101-
}
3102-
31033111
let mut base = self;
31043112
let mut acc: Self = 1;
31053113
let mut overflow = false;
31063114
let mut tmp_overflow;
31073115

3116+
if intrinsics::is_val_statically_known(base) && base.unsigned_abs().is_power_of_two() {
3117+
let k = base.unsigned_abs().ilog2();
3118+
let Some(shift) = k.checked_mul(exp) else {
3119+
return (0, true)
3120+
};
3121+
let magnitude = (1 as Self).unbounded_shl(shift);
3122+
return if base < 0 && (exp % 2) == 1 {
3123+
(magnitude.wrapping_neg(), shift >= Self::BITS)
3124+
} else {
3125+
(magnitude, shift >= Self::BITS)
3126+
}
3127+
}
3128+
3129+
if exp == 0 {
3130+
return (1, false);
3131+
}
3132+
31083133
if intrinsics::is_val_statically_known(exp) {
31093134
while exp > 1 {
31103135
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
@@ -2367,11 +2367,23 @@ macro_rules! uint_impl {
23672367
without modifying the original"]
23682368
#[inline]
23692369
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
2370+
let mut base = self;
2371+
let mut acc: Self = 1;
2372+
2373+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
2374+
// change of base:
2375+
// if base == 2 ** k, then
2376+
// (2 ** k) ** n
2377+
// == 2 ** (k * n)
2378+
// == 1 << (k * n)
2379+
let k = base.ilog2();
2380+
let shift = try_opt!(k.checked_mul(exp));
2381+
return (1 as Self).checked_shl(shift);
2382+
}
2383+
23702384
if exp == 0 {
23712385
return Some(1);
23722386
}
2373-
let mut base = self;
2374-
let mut acc: Self = 1;
23752387

23762388
if intrinsics::is_val_statically_known(exp) {
23772389
while exp > 1 {
@@ -3566,14 +3578,28 @@ macro_rules! uint_impl {
35663578
without modifying the original"]
35673579
#[inline]
35683580
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3569-
if exp == 0 {
3570-
return (1, false);
3571-
}
35723581
let mut base = self;
35733582
let mut acc: Self = 1;
35743583
let mut overflow = false;
35753584
let mut tmp_overflow;
35763585

3586+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
3587+
// change of base:
3588+
// if base == 2 ** k, then
3589+
// (2 ** k) ** n
3590+
// == 2 ** (k * n)
3591+
// == 1 << (k * n)
3592+
let k = base.ilog2();
3593+
let Some(shift) = k.checked_mul(exp) else {
3594+
return (0, true)
3595+
};
3596+
return ((1 as Self).unbounded_shl(shift), shift >= Self::BITS)
3597+
}
3598+
3599+
if exp == 0 {
3600+
return (1, false);
3601+
}
3602+
35773603
if intrinsics::is_val_statically_known(exp) {
35783604
while exp > 1 {
35793605
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+
}

0 commit comments

Comments
 (0)