Skip to content

Commit 3926960

Browse files
Rollup merge of rust-lang#128166 - ChaiTRex:isqrt, r=tgross35
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 rust-lang#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>
2 parents acb4e8b + 8e54e87 commit 3926960

File tree

11 files changed

+675
-67
lines changed

11 files changed

+675
-67
lines changed

library/core/benches/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#![feature(iter_array_chunks)]
99
#![feature(iter_next_chunk)]
1010
#![feature(iter_advance_by)]
11+
#![feature(isqrt)]
1112

1213
extern crate test;
1314

+62
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
use rand::Rng;
2+
use test::{black_box, Bencher};
3+
4+
macro_rules! int_sqrt_bench {
5+
($t:ty, $predictable:ident, $random:ident, $random_small:ident, $random_uniform:ident) => {
6+
#[bench]
7+
fn $predictable(bench: &mut Bencher) {
8+
bench.iter(|| {
9+
for n in 0..(<$t>::BITS / 8) {
10+
for i in 1..=(100 as $t) {
11+
let x = black_box(i << (n * 8));
12+
black_box(x.isqrt());
13+
}
14+
}
15+
});
16+
}
17+
18+
#[bench]
19+
fn $random(bench: &mut Bencher) {
20+
let mut rng = crate::bench_rng();
21+
/* Exponentially distributed random numbers from the whole range of the type. */
22+
let numbers: Vec<$t> =
23+
(0..256).map(|_| rng.gen::<$t>() >> rng.gen_range(0..<$t>::BITS)).collect();
24+
bench.iter(|| {
25+
for x in &numbers {
26+
black_box(black_box(x).isqrt());
27+
}
28+
});
29+
}
30+
31+
#[bench]
32+
fn $random_small(bench: &mut Bencher) {
33+
let mut rng = crate::bench_rng();
34+
/* Exponentially distributed random numbers from the range 0..256. */
35+
let numbers: Vec<$t> =
36+
(0..256).map(|_| (rng.gen::<u8>() >> rng.gen_range(0..u8::BITS)) as $t).collect();
37+
bench.iter(|| {
38+
for x in &numbers {
39+
black_box(black_box(x).isqrt());
40+
}
41+
});
42+
}
43+
44+
#[bench]
45+
fn $random_uniform(bench: &mut Bencher) {
46+
let mut rng = crate::bench_rng();
47+
/* Exponentially distributed random numbers from the whole range of the type. */
48+
let numbers: Vec<$t> = (0..256).map(|_| rng.gen::<$t>()).collect();
49+
bench.iter(|| {
50+
for x in &numbers {
51+
black_box(black_box(x).isqrt());
52+
}
53+
});
54+
}
55+
};
56+
}
57+
58+
int_sqrt_bench! {u8, u8_sqrt_predictable, u8_sqrt_random, u8_sqrt_random_small, u8_sqrt_uniform}
59+
int_sqrt_bench! {u16, u16_sqrt_predictable, u16_sqrt_random, u16_sqrt_random_small, u16_sqrt_uniform}
60+
int_sqrt_bench! {u32, u32_sqrt_predictable, u32_sqrt_random, u32_sqrt_random_small, u32_sqrt_uniform}
61+
int_sqrt_bench! {u64, u64_sqrt_predictable, u64_sqrt_random, u64_sqrt_random_small, u64_sqrt_uniform}
62+
int_sqrt_bench! {u128, u128_sqrt_predictable, u128_sqrt_random, u128_sqrt_random_small, u128_sqrt_uniform}

library/core/benches/num/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ mod dec2flt;
22
mod flt2dec;
33
mod int_log;
44
mod int_pow;
5+
mod int_sqrt;
56

67
use std::str::FromStr;
78

library/core/src/num/int_macros.rs

+29-7
Original file line numberDiff line numberDiff line change
@@ -1641,7 +1641,33 @@ macro_rules! int_impl {
16411641
if self < 0 {
16421642
None
16431643
} else {
1644-
Some((self as $UnsignedT).isqrt() as Self)
1644+
// SAFETY: Input is nonnegative in this `else` branch.
1645+
let result = unsafe {
1646+
crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT
1647+
};
1648+
1649+
// Inform the optimizer what the range of outputs is. If
1650+
// testing `core` crashes with no panic message and a
1651+
// `num::int_sqrt::i*` test failed, it's because your edits
1652+
// caused these assertions to become false.
1653+
//
1654+
// SAFETY: Integer square root is a monotonically nondecreasing
1655+
// function, which means that increasing the input will never
1656+
// cause the output to decrease. Thus, since the input for
1657+
// nonnegative signed integers is bounded by
1658+
// `[0, <$ActualT>::MAX]`, sqrt(n) will be bounded by
1659+
// `[sqrt(0), sqrt(<$ActualT>::MAX)]`.
1660+
unsafe {
1661+
// SAFETY: `<$ActualT>::MAX` is nonnegative.
1662+
const MAX_RESULT: $SelfT = unsafe {
1663+
crate::num::int_sqrt::$ActualT(<$ActualT>::MAX) as $SelfT
1664+
};
1665+
1666+
crate::hint::assert_unchecked(result >= 0);
1667+
crate::hint::assert_unchecked(result <= MAX_RESULT);
1668+
}
1669+
1670+
Some(result)
16451671
}
16461672
}
16471673

@@ -2862,15 +2888,11 @@ macro_rules! int_impl {
28622888
#[must_use = "this returns the result of the operation, \
28632889
without modifying the original"]
28642890
#[inline]
2891+
#[track_caller]
28652892
pub const fn isqrt(self) -> Self {
2866-
// I would like to implement it as
2867-
// ```
2868-
// self.checked_isqrt().expect("argument of integer square root must be non-negative")
2869-
// ```
2870-
// but `expect` is not yet stable as a `const fn`.
28712893
match self.checked_isqrt() {
28722894
Some(sqrt) => sqrt,
2873-
None => panic!("argument of integer square root must be non-negative"),
2895+
None => crate::num::int_sqrt::panic_for_negative_argument(),
28742896
}
28752897
}
28762898

0 commit comments

Comments
 (0)