Skip to content

Commit

Permalink
Optimize integer pow by removing exit branch
Browse files Browse the repository at this point in the history 
The branch at the end of the `pow` implementations is redundant
with multiplication code already present in the loop. By rotating
the exit check, this branch can be largely removed, improving code size
and instruction cache coherence.
  • Loading branch information
@mzabaluev
mzabaluev committed Mar 22, 2024
1 parent eb80be2 commit 76d2530
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 73 deletions.
61 changes: 26 additions & 35 deletions library/core/src/num/int_macros.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1416,18 +1416,17 @@ macro_rules! int_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = try_opt!(acc.checked_mul(base));
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return Some(acc);
}
}
exp /= 2;
base = try_opt!(base.checked_mul(base));
}
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.checked_mul(base)
}

/// Strict exponentiation. Computes `self.pow(exp)`, panicking if
Expand Down Expand Up @@ -1467,18 +1466,17 @@ macro_rules! int_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc.strict_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.strict_mul(base);
}
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.strict_mul(base)
}

/// Returns the square root of the number, rounded down.
Expand Down Expand Up @@ -2102,19 +2100,17 @@ macro_rules! int_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc.wrapping_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.wrapping_mul(base);
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.wrapping_mul(base)
}

/// Calculates `self` + `rhs`
Expand Down Expand Up @@ -2608,9 +2604,14 @@ macro_rules! int_impl {
// Scratch space for storing results of overflowing_mul.
let mut r;

while exp > 1 {
loop {
if (exp & 1) == 1 {
r = acc.overflowing_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
r.1 |= overflown;
return r;
}
acc = r.0;
overflown |= r.1;
}
Expand All @@ -2619,14 +2620,6 @@ macro_rules! int_impl {
base = r.0;
overflown |= r.1;
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
r = acc.overflowing_mul(base);
r.1 |= overflown;
r
}

/// Raises self to the power of `exp`, using exponentiation by squaring.
Expand All @@ -2653,19 +2646,17 @@ macro_rules! int_impl {
let mut base = self;
let mut acc = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc * base;
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base * base;
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc * base
}

/// Returns the square root of the number, rounded down.
Expand Down
64 changes: 26 additions & 38 deletions library/core/src/num/uint_macros.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1413,20 +1413,17 @@ macro_rules! uint_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = try_opt!(acc.checked_mul(base));
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return Some(acc);
}
}
exp /= 2;
base = try_opt!(base.checked_mul(base));
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.

acc.checked_mul(base)
}

/// Strict exponentiation. Computes `self.pow(exp)`, panicking if
Expand Down Expand Up @@ -1466,18 +1463,17 @@ macro_rules! uint_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc.strict_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.strict_mul(base);
}
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.strict_mul(base)
}

/// Saturating integer addition. Computes `self + rhs`, saturating at
Expand Down Expand Up @@ -1938,19 +1934,17 @@ macro_rules! uint_impl {
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc.wrapping_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.wrapping_mul(base);
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.wrapping_mul(base)
}

/// Calculates `self` + `rhs`
Expand Down Expand Up @@ -2395,9 +2389,14 @@ macro_rules! uint_impl {
// Scratch space for storing results of overflowing_mul.
let mut r;

while exp > 1 {
loop {
if (exp & 1) == 1 {
r = acc.overflowing_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
r.1 |= overflown;
return r;
}
acc = r.0;
overflown |= r.1;
}
Expand All @@ -2406,15 +2405,6 @@ macro_rules! uint_impl {
base = r.0;
overflown |= r.1;
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
r = acc.overflowing_mul(base);
r.1 |= overflown;

r
}

/// Raises self to the power of `exp`, using exponentiation by squaring.
Expand All @@ -2439,19 +2429,17 @@ macro_rules! uint_impl {
let mut base = self;
let mut acc = 1;

while exp > 1 {
loop {
if (exp & 1) == 1 {
acc = acc * base;
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base * base;
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc * base
}

/// Returns the square root of the number, rounded down.
Expand Down

0 comments on commit 76d2530

Please sign in to comment.