Skip to content

Commit f23f1e5

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 14eb251 commit f23f1e5

3 files changed

Lines changed: 125 additions & 11 deletions

File tree

‎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
@@ -2280,11 +2280,23 @@ macro_rules! uint_impl {
22802280
without modifying the original"]
22812281
#[inline]
22822282
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
2283+
let mut base = self;
2284+
let mut acc: Self = 1;
2285+
2286+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
2287+
// change of base:
2288+
// if base == 2 ** k, then
2289+
// (2 ** k) ** n
2290+
// == 2 ** (k * n)
2291+
// == 1 << (k * n)
2292+
let k = base.ilog2();
2293+
let shift = try_opt!(k.checked_mul(exp));
2294+
return (1 as Self).checked_shl(shift);
2295+
}
2296+
22832297
if exp == 0 {
22842298
return Some(1);
22852299
}
2286-
let mut base = self;
2287-
let mut acc: Self = 1;
22882300

22892301
if intrinsics::is_val_statically_known(exp) {
22902302
while exp > 1 {
@@ -3479,14 +3491,28 @@ macro_rules! uint_impl {
34793491
without modifying the original"]
34803492
#[inline]
34813493
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
3482-
if exp == 0 {
3483-
return (1, false);
3484-
}
34853494
let mut base = self;
34863495
let mut acc: Self = 1;
34873496
let mut overflow = false;
34883497
let mut tmp_overflow;
34893498

3499+
if intrinsics::is_val_statically_known(base) && base.is_power_of_two() {
3500+
// change of base:
3501+
// if base == 2 ** k, then
3502+
// (2 ** k) ** n
3503+
// == 2 ** (k * n)
3504+
// == 1 << (k * n)
3505+
let k = base.ilog2();
3506+
let Some(shift) = k.checked_mul(exp) else {
3507+
return (0, true)
3508+
};
3509+
return ((1 as Self).unbounded_shl(shift), shift >= Self::BITS)
3510+
}
3511+
3512+
if exp == 0 {
3513+
return (1, false);
3514+
}
3515+
34903516
if intrinsics::is_val_statically_known(exp) {
34913517
while exp > 1 {
34923518
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)