Skip to content

Commit

Permalink
Auto merge of #128166 - ChaiTRex:isqrt, r=<try>
Browse files Browse the repository at this point in the history
Improved `checked_isqrt` and `isqrt` methods

### Improved tests of `isqrt` and `checked_isqrt` implementations

* Inputs chosen more thoroughly and systematically.
* Checks that `isqrt` and `checked_isqrt` have equivalent results for signed types, either equivalent numerically or equivalent as a panic and a `None`.
* Checks that `isqrt` has numerically-equivalent results for unsigned types and their `NonZero` counterparts.

### Added benchmarks for `isqrt` implementations

### Greatly sped up `checked_isqrt` and `isqrt` methods

* Uses a lookup table for 8-bit integers and then the Karatsuba square root algorithm for larger integers.
* Includes optimization hints that give the compiler the exact numeric range of results.

### Feature tracking issue

`isqrt` is an unstable feature tracked at #116226.

<details><summary>Benchmarked improvements</summary>

### Command used to benchmark

    ./x bench library/core -- int_sqrt

### Before

    benchmarks:
        num::int_sqrt::i128::isqrt           439591.65/iter  +/- 6652.70
        num::int_sqrt::i16::isqrt              5302.97/iter   +/- 160.93
        num::int_sqrt::i32::isqrt             62999.11/iter  +/- 2022.05
        num::int_sqrt::i64::isqrt            125248.81/iter  +/- 1674.43
        num::int_sqrt::i8::isqrt                123.56/iter     +/- 1.87
        num::int_sqrt::isize::isqrt          125356.56/iter  +/- 1017.03
        num::int_sqrt::non_zero_u128::isqrt  437443.75/iter  +/- 3535.43
        num::int_sqrt::non_zero_u16::isqrt     8604.58/iter    +/- 94.76
        num::int_sqrt::non_zero_u32::isqrt    62933.33/iter   +/- 517.30
        num::int_sqrt::non_zero_u64::isqrt   125076.38/iter +/- 11340.61
        num::int_sqrt::non_zero_u8::isqrt       221.51/iter     +/- 1.58
        num::int_sqrt::non_zero_usize::isqrt 136005.21/iter  +/- 2020.35
        num::int_sqrt::u128::isqrt           439014.55/iter  +/- 3920.45
        num::int_sqrt::u16::isqrt              8575.08/iter   +/- 148.06
        num::int_sqrt::u32::isqrt             63008.89/iter   +/- 803.67
        num::int_sqrt::u64::isqrt            125088.09/iter   +/- 879.29
        num::int_sqrt::u8::isqrt                230.18/iter     +/- 2.04
        num::int_sqrt::usize::isqrt          125237.51/iter  +/- 4747.83
### After

    benchmarks:
        num::int_sqrt::i128::isqrt           105184.89/iter +/- 1171.38
        num::int_sqrt::i16::isqrt              1910.26/iter   +/- 78.50
        num::int_sqrt::i32::isqrt             34260.34/iter  +/- 960.84
        num::int_sqrt::i64::isqrt             45939.19/iter +/- 2525.65
        num::int_sqrt::i8::isqrt                 22.87/iter    +/- 0.45
        num::int_sqrt::isize::isqrt           45884.17/iter  +/- 595.49
        num::int_sqrt::non_zero_u128::isqrt  106344.27/iter  +/- 780.99
        num::int_sqrt::non_zero_u16::isqrt     2790.19/iter   +/- 53.43
        num::int_sqrt::non_zero_u32::isqrt    33613.99/iter  +/- 362.96
        num::int_sqrt::non_zero_u64::isqrt    46235.42/iter  +/- 429.69
        num::int_sqrt::non_zero_u8::isqrt        31.78/iter    +/- 0.75
        num::int_sqrt::non_zero_usize::isqrt  46208.75/iter  +/- 375.27
        num::int_sqrt::u128::isqrt           106385.94/iter +/- 1649.95
        num::int_sqrt::u16::isqrt              2747.69/iter   +/- 28.72
        num::int_sqrt::u32::isqrt             33627.09/iter  +/- 475.68
        num::int_sqrt::u64::isqrt             46182.29/iter  +/- 311.16
        num::int_sqrt::u8::isqrt                 33.10/iter    +/- 0.30
        num::int_sqrt::usize::isqrt           46165.00/iter  +/- 388.41

</details>

Tracking Issue for {u8,i8,...}::isqrt #116226

try-job: test-various
  • Loading branch information
bors committed Aug 29, 2024
2 parents acb4e8b + 7af8e21 commit daf8e6a
Show file tree
Hide file tree
Showing 11 changed files with 684 additions and 67 deletions.
1 change: 1 addition & 0 deletions library/core/benches/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#![feature(iter_array_chunks)]
#![feature(iter_next_chunk)]
#![feature(iter_advance_by)]
#![feature(isqrt)]

extern crate test;

Expand Down
62 changes: 62 additions & 0 deletions library/core/benches/num/int_sqrt/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
use rand::Rng;
use test::{black_box, Bencher};

macro_rules! int_sqrt_bench {
($t:ty, $predictable:ident, $random:ident, $random_small:ident, $random_uniform:ident) => {
#[bench]
fn $predictable(bench: &mut Bencher) {
bench.iter(|| {
for n in 0..(<$t>::BITS / 8) {
for i in 1..=(100 as $t) {
let x = black_box(i << (n * 8));
black_box(x.isqrt());
}
}
});
}

#[bench]
fn $random(bench: &mut Bencher) {
let mut rng = crate::bench_rng();
/* Exponentially distributed random numbers from the whole range of the type. */
let numbers: Vec<$t> =
(0..256).map(|_| rng.gen::<$t>() >> rng.gen_range(0..<$t>::BITS)).collect();
bench.iter(|| {
for x in &numbers {
black_box(black_box(x).isqrt());
}
});
}

#[bench]
fn $random_small(bench: &mut Bencher) {
let mut rng = crate::bench_rng();
/* Exponentially distributed random numbers from the range 0..256. */
let numbers: Vec<$t> =
(0..256).map(|_| (rng.gen::<u8>() >> rng.gen_range(0..u8::BITS)) as $t).collect();
bench.iter(|| {
for x in &numbers {
black_box(black_box(x).isqrt());
}
});
}

#[bench]
fn $random_uniform(bench: &mut Bencher) {
let mut rng = crate::bench_rng();
/* Exponentially distributed random numbers from the whole range of the type. */
let numbers: Vec<$t> = (0..256).map(|_| rng.gen::<$t>()).collect();
bench.iter(|| {
for x in &numbers {
black_box(black_box(x).isqrt());
}
});
}
};
}

int_sqrt_bench! {u8, u8_sqrt_predictable, u8_sqrt_random, u8_sqrt_random_small, u8_sqrt_uniform}
int_sqrt_bench! {u16, u16_sqrt_predictable, u16_sqrt_random, u16_sqrt_random_small, u16_sqrt_uniform}
int_sqrt_bench! {u32, u32_sqrt_predictable, u32_sqrt_random, u32_sqrt_random_small, u32_sqrt_uniform}
int_sqrt_bench! {u64, u64_sqrt_predictable, u64_sqrt_random, u64_sqrt_random_small, u64_sqrt_uniform}
int_sqrt_bench! {u128, u128_sqrt_predictable, u128_sqrt_random, u128_sqrt_random_small, u128_sqrt_uniform}
1 change: 1 addition & 0 deletions library/core/benches/num/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ mod dec2flt;
mod flt2dec;
mod int_log;
mod int_pow;
mod int_sqrt;

use std::str::FromStr;

Expand Down
36 changes: 29 additions & 7 deletions library/core/src/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1641,7 +1641,33 @@ macro_rules! int_impl {
if self < 0 {
None
} else {
Some((self as $UnsignedT).isqrt() as Self)
// SAFETY: Input is nonnegative in this `else` branch.
let result = unsafe {
crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT
};

// Inform the optimizer what the range of outputs is. If
// testing `core` crashes with no panic message and a
// `num::int_sqrt::i*` test failed, it's because your edits
// caused these assertions to become false.
//
// SAFETY: Integer square root is a monotonically nondecreasing
// function, which means that increasing the input will never
// cause the output to decrease. Thus, since the input for
// nonnegative signed integers is bounded by
// `[0, <$ActualT>::MAX]`, sqrt(n) will be bounded by
// `[sqrt(0), sqrt(<$ActualT>::MAX)]`.
unsafe {
// SAFETY: `<$ActualT>::MAX` is nonnegative.
const MAX_RESULT: $SelfT = unsafe {
crate::num::int_sqrt::$ActualT(<$ActualT>::MAX) as $SelfT
};

crate::hint::assert_unchecked(result >= 0);
crate::hint::assert_unchecked(result <= MAX_RESULT);
}

Some(result)
}
}

Expand Down Expand Up @@ -2862,15 +2888,11 @@ macro_rules! int_impl {
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
#[track_caller]
pub const fn isqrt(self) -> Self {
// I would like to implement it as
// ```
// self.checked_isqrt().expect("argument of integer square root must be non-negative")
// ```
// but `expect` is not yet stable as a `const fn`.
match self.checked_isqrt() {
Some(sqrt) => sqrt,
None => panic!("argument of integer square root must be non-negative"),
None => crate::num::int_sqrt::panic_for_negative_argument(),
}
}

Expand Down
Loading

0 comments on commit daf8e6a

Please sign in to comment.