Regarding isqrt performance

[leonardo-m]

This is experimental code, it could be compiled with one or the other implementation of the integer square root:

#![feature(cfg_overflow_checks, stmt_expr_attributes)]
#![allow(unused_macros)]

#[cfg(overflow_checks)]
macro_rules! fto { ($e:expr, $t:ident) => (($e) as $t) }

#[cfg(not(overflow_checks))]
macro_rules! fto { ($e:expr, $t:ident) => (#[allow(unused_unsafe)] unsafe { $e.to_int_unchecked::<$t>() }) }

macro_rules! isqrt { ($e:expr, $t:ident) => ( fto!{(($e) as f64).sqrt(), $t} ) }

fn main() {
    let mut tot: u64 = 0;
    for x in (1_u64 .. 1_000_000_000).step_by(3) {
        tot += isqrt!(x, u64); // Version A.  1.04 seconds.
        //tot += x.isqrt(); // Version B.       4.94 seconds.
    }
    println!("{tot}");
}

As you see the floating-point based isqrt implementation is almost five times faster on my PC (using rustc 1.87.0-nightly). The std lib isqrt has upside of being const, so I can use it for const generics calculations and similar situations. And when I need to compute only one or few isqrt, this performance difference isn't important, and I use the std one. But when I need to compute a lot of isqrt, I consider faster alternatives.

While I don't propose to replace the std library isqrt with the code I've shown here, I suggest std lib implementers to express an opinion regarding the usage of floating point sqrt+ int cast in some vetted and safe cases.

[bjorn3]

cc https://rust-lang.zulipchat.com/#narrow/channel/122651-general/topic/Altering.20a.20.60core.60.20method.20impl.20after.20.60core.60.20is.20compiled

[CatsAreFluffy]

This implementation of u32::isqrt is float-free and significantly faster than the stdlib version:

// SQRTS[i<256] = (i << 8).isqrt() as u8
const fn new_isqrt(x: u32) -> u32 {
    if x < 256 {
        return SQRTS[x as usize] as u32 >> 4;
    }
    let idx = x >> ((25 - x.leading_zeros()) & !1);
    // SAFETY: If x has y leading zeros, the shift count is either 24 - y or
    // 25 - y. Thus idx has either 24 or 25 leading zeroes, and in particular
    // it's less than 256.
    unsafe { std::hint::assert_unchecked(idx < 256) };
    let approx1 = SQRTS[idx as usize] as u32;
    // SAFETY: Every element of SQRTS is at least 16 except element 0, and
    // idx is positive so that cannot be selected.
    unsafe { std::hint::assert_unchecked(approx1 >= 16) };
    let mut approx = (approx1 + 1) << ((25 - x.leading_zeros()) / 2) >> 4;
    approx = (approx + x / approx) / 2;
    if approx * approx > x {
        approx -= 1;
    }
    approx
}

Godbolt

On my M3 Max laptop, this implementation takes ~6.7 seconds to evaluate for all 2^32 inputs, while u32::isqrt takes ~15.7 seconds, and |x| (x as f64).sqrt() as u32 takes ~3.0 seconds. (I've included the test harness I used to get those numbers in the Godbolt if you want to try it for yourself.)

[cuviper]

With my Ryzen 7 9800X3D on Linux, that's significantly slower at first:

New isqrt:
Time = 28.09259456s
f64 isqrt:
Time = 6.522072657s
Stdlib isqrt:
Time = 18.097873232s

That turns around with -Ctarget-cpu=native (zen5):

New isqrt:
Time = 6.881527916s
f64 isqrt:
Time = 6.457018049s
Stdlib isqrt:
Time = 17.119138915s

[ChaiTRex]

To avoid the SQRTS generation relying on the isqrt method (which may become circular if we replace more than the 32-bit isqrt methods), the following code can be used (checked to be identical with the code in the Godbolt above):

/// Fixed-point square roots of 0..=255, with 4 bits integer part and
/// 4 bits fractional part.
const FIXED_POINT_SQRTS: [u8; 256] = {
    let mut result = [0; 256];

    let mut sqrt = 0_u32;
    let mut i = 0_u32;
    while i < 256 {
        while sqrt * sqrt <= (i << 8) {
            sqrt += 1;
        }
        sqrt -= 1;

        result[i as usize] = sqrt as u8;
        i += 1;
    }

    result
};

[CatsAreFluffy]

With my Ryzen 7 9800X3D on Linux, that's significantly slower at first:

This is caused by an issue with how leading_zeros is compiled on x86 and x86-64 targets not supporting BMI1. (Specifically, the bsr instruction, which is used to implement leading_zeros, has a dependency on its output register, and for new_isqrt, the register LLVM picks for that is also used in the validation code, so iterations of the loop are almost completely serialized, while for u32::isqrt, LLVM picks a register which isn't used for anything else, so there's approximately no speed penalty.) You could work around this by using inline assembly that properly breaks the dependency instead of leading_zeros, although I'm not sure if it would matter much outside of simple benchmarks like this.

[tgross35]

This is caused by an issue with how leading_zeros is compiled on x86 and x86-64 targets not supporting BMI1 (specifically, the bsr instruction, which is used to implement leading_zeros, has a dependency on its output register, and for new_isqrt, the register LLVM picks for that is also used in the validation code, so iterations of the loop are almost completely serialized, while for u32::isqrt, LLVM picks a register which isn't used for anything else, so there's approximately no speed penalty).

Does LLVM have enough information to make a better register selection, in theory? If so, this would be worth opening an issue (if one doesn't already exist).

[CatsAreFluffy]

I'm not sure how feasible it is to solve that by picking a better output register, but it's also solvable by initializing the destination register immediately before the bsr so the output dependency can't get stuck on anything else, and I've opened llvm/llvm-project#129659 for that.

[CatsAreFluffy]

I've made a slightly faster isqrt implementation.

// SQRTS[i<256] = (i * 256).isqrt()
// RECIPS[i<192] = (1 << 38).div_ceil(SQRTS[i + 64] + 1)
fn newer_isqrt2(x: u32) -> u32 {
    if x < 256 {
        return SQRTS[x as usize] as u32 >> 4;
    }
    let idx = x >> ((25 - x.leading_zeros()) & !1);
    // SAFETY: If x has y leading zeros, the shift count is either 24 - y or
    // 25 - y. Thus idx has either 24 or 25 leading zeroes, and in particular
    // it's in 64..256.
    unsafe { std::hint::assert_unchecked(64 <= idx) };
    unsafe { std::hint::assert_unchecked(idx < 256) };
    let approx1 = SQRTS[idx as usize] as u32 + 1;
    let approx2 = approx1 << ((25 - x.leading_zeros()) / 2);
    let divmult = RECIPS[idx as usize - 64] as u64;
    // Approximately `x / approx2 * 16`.
    let approxr = (x as u64 * divmult >> (25 - x.leading_zeros()) / 2 + 30) as u32;
    let mut approx3 = approx2 + approxr >> 5;
    if approx3 * approx3 > x {
        approx3 -= 1;
    }
    approx3
}

Godbolt. (newer_isqrt2 is the faster one, newer_isqrt is similar but slower than new_isqrt.)

On my laptop, I get:

Newer isqrt 2:
Time = 6.589210709s
Newer isqrt:
Time = 7.007925959s
New isqrt:
Time = 6.910427791s
f64 isqrt:
Time = 4.150698459s
Stdlib isqrt:
Time = 14.867240041s