Skip to content

Commit b1b3348

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 9988dcd commit b1b3348

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
@@ -1740,11 +1740,25 @@ macro_rules! int_impl {
17401740
without modifying the original"]
17411741
#[inline]
17421742
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
1743+
let mut base = self;
1744+
let mut acc: Self = 1;
1745+
1746+
if intrinsics::is_val_statically_known(base) {
1747+
if base.unsigned_abs().is_power_of_two() {
1748+
let k = base.unsigned_abs().ilog2();
1749+
let shift = try_opt!(k.checked_mul(exp));
1750+
let magnitude = try_opt!((1 as Self).checked_shl(shift));
1751+
return if base < 0 && (exp % 2) == 1 {
1752+
Some(magnitude.wrapping_neg())
1753+
} else {
1754+
Some(magnitude)
1755+
}
1756+
}
1757+
}
1758+
17431759
if exp == 0 {
17441760
return Some(1);
17451761
}
1746-
let mut base = self;
1747-
let mut acc: Self = 1;
17481762

17491763
if intrinsics::is_val_statically_known(exp) {
17501764
while exp > 1 {
@@ -2959,15 +2973,30 @@ macro_rules! int_impl {
29592973
without modifying the original"]
29602974
#[inline]
29612975
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
2962-
if exp == 0 {
2963-
return (1, false);
2964-
}
2965-
29662976
let mut base = self;
29672977
let mut acc: Self = 1;
29682978
let mut overflow = false;
29692979
let mut tmp_overflow;
29702980

2981+
if intrinsics::is_val_statically_known(base) {
2982+
if base.unsigned_abs().is_power_of_two() {
2983+
let k = base.unsigned_abs().ilog2();
2984+
let Some(shift) = k.checked_mul(exp) else {
2985+
return (0, true)
2986+
};
2987+
let magnitude = (1 as Self).unbounded_shl(shift);
2988+
return if base < 0 && (exp % 2) == 1 {
2989+
(magnitude.wrapping_neg(), shift >= Self::BITS)
2990+
} else {
2991+
(magnitude, shift >= Self::BITS)
2992+
}
2993+
}
2994+
}
2995+
2996+
if exp == 0 {
2997+
return (1, false);
2998+
}
2999+
29713000
if intrinsics::is_val_statically_known(exp) {
29723001
while exp > 1 {
29733002
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
@@ -2085,11 +2085,23 @@ macro_rules! uint_impl {
20852085
without modifying the original"]
20862086
#[inline]
20872087
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
2088+
let mut base = self;
2089+
let mut acc: Self = 1;
2090+
2091+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
2092+
// change of base:
2093+
// if base == 2 ** k, then
2094+
// (2 ** k) ** n
2095+
// == 2 ** (k * n)
2096+
// == 1 << (k * n)
2097+
let k = base.ilog2();
2098+
let shift = try_opt!(k.checked_mul(exp));
2099+
return (1 as Self).checked_shl(shift);
2100+
}
2101+
20882102
if exp == 0 {
20892103
return Some(1);
20902104
}
2091-
let mut base = self;
2092-
let mut acc: Self = 1;
20932105

20942106
if intrinsics::is_val_statically_known(exp) {
20952107
while exp > 1 {
@@ -3255,14 +3267,28 @@ macro_rules! uint_impl {
32553267
without modifying the original"]
32563268
#[inline]
32573269
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3258-
if exp == 0 {
3259-
return (1, false);
3260-
}
32613270
let mut base = self;
32623271
let mut acc: Self = 1;
32633272
let mut overflow = false;
32643273
let mut tmp_overflow;
32653274

3275+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
3276+
// change of base:
3277+
// if base == 2 ** k, then
3278+
// (2 ** k) ** n
3279+
// == 2 ** (k * n)
3280+
// == 1 << (k * n)
3281+
let k = base.ilog2();
3282+
let Some(shift) = k.checked_mul(exp) else {
3283+
return (0, true)
3284+
};
3285+
return ((1 as Self).unbounded_shl(shift), shift >= Self::BITS)
3286+
}
3287+
3288+
if exp == 0 {
3289+
return (1, false);
3290+
}
3291+
32663292
if intrinsics::is_val_statically_known(exp) {
32673293
while exp > 1 {
32683294
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)