Fast integer square-root

Question

I'm looking for the highest-performance method of calculating integer square roots in Mathematica of very big arbitrary-precision numbers.

As an example testcase, I use:

n = 10^1000000 - 3 ^ 2095903

On my Core i7 linux machine, calculating the integer (floor) square root using the straightforward method takes 2.68 seconds:

In[1]:= n = 10^1000000 - 3^2095903;
In[2]:= First@Timing@Floor[Sqrt[n]]
Out[2]= 2.68984

However, the same machine can calculate the integer square root much faster using the GMP library. Here's an example python program doing just that:

from gmpy2 import mpz, isqrt
from time import time

n = mpz(10**1000000 - 3**2095903)

t0 = time()
s = isqrt(n)
t1 = time()

print(t1 - t0)

On the same machine, this is approximately 160 times faster:

$ python3 benchmark.py
0.01685619354248047

The frustrating thing is that Mathematica already depends on the GMP library for its arbitrary-precision integers, but I cannot seem to make Mathematica use the GMP implementation of the integer square root.

Instead I suspect that its Floor[Sqrt[...]] function falls back on a generic algorithm that must calculate its argument to sufficient precision using arbitrary-precision floating-point approximation.

Is there any way to speed up the calculation of integer square roots, preferably approaching GMP's raw performance?

Answer 1

Sqrt is using exact methods in an effort to pull out "small" squares. This is going to take time. A direct approach, as already noted, would do the square root numerically. For purposes of Floor extraction, it suffices to use as precision half the digit size.

floorSqrt[n_Integer] := 
 Module[{prec = RealExponent[1.001*n]/2}, 
  Floor[Sqrt[SetPrecision[n, prec]]]]

Test:

n = 10^1000000 - 3^2095903;
AbsoluteTiming[sn1 = Floor[Sqrt[n]];]
AbsoluteTiming[sn2 = floorSqrt[n];]
sn1 === sn2

(* Out[410]= {2.774283, Null}

Out[411]= {0.016654, Null}

Out[412]= True *)

For variety, here is a top-level implementation of an integer-based method. My guess is it is similar to Zmmermann's, but he may well have had some extra efficiencies. The idea is to split the number into an upper and lower part, using a power of 4 for the split size so that we can shift back by a factor of 2. That is, write a = 4^n*b+c. Recursively compute the integer sqrt of b, multiply by 2^n, and use the usual Taylor approximation to get a correction that estimates sqrt(a). Last step is to iteratively repair that estimate. We use integer multiplication, squaring, and the integer Quotient function. The recursion has to bottom out so we provide some base size below which we use the numeric approximation method noted already.

iSqrt[a_, baselen_] := Catch[Module[
   {quarterscale = Ceiling[RealExponent[a, 2.]/4], aUpper, aLower, 
    sqrt, diff},
   If[quarterscale < baselen, Throw[Floor[Sqrt[N[a, 2*baselen + 4]]]]];
   aUpper = BitShiftRight[a, 2*quarterscale];
   aLower = a - BitShiftLeft[aUpper, 2*quarterscale];
   sqrt = BitShiftLeft[iSqrt[aUpper, baselen], quarterscale];
   sqrt += Quotient[aLower, (2*sqrt)];
   diff = a - sqrt^2;
   While[Abs[diff] >= 2 sqrt + 1 || diff < 0,
    sqrt += Quotient[diff, (2*sqrt)];
    diff = a - sqrt^2;
    ];
   sqrt
   ]]

This is not as fast as the numerical approximation method but it is not bad either. Below is on my slow laptop.

n = 10^1000000 - 3^2095903;
AbsoluteTiming[sn1 = Floor[Sqrt[n]];]
AbsoluteTiming[sn2 = floorSqrt[n];]
AbsoluteTiming[sn3 = iSqrt[n, 10];]
sn1 == sn2 == sn3

(* Out[642]= {7.20206, Null}

Out[643]= {0.0479314, Null}

Out[644]= {0.128822, Null}

Out[645]= True *)

So it's around 2.5 times slower than the numeric code. An internal implementation might do better in terms of avoiding some amount of overhead, and might possibly make up a good chunk from that factor of 2.5.

--- edit ---

I finally added this to the Wolfram Function Repository, as FloorSqrt.

First[Timing[isqrt = 
    ResourceFunction["FloorSqrt"][n = 10^1000000 - 3^2095903];]]

(* Out[200]= 0.02877 *)

Check the result:

IntegerQ[isqrt] && isqrt^2 <= n < (isqrt + 1)^2

(* Out[201]= True *)

Answer 2

Not directly addressing the question about GMP library, but you can get a Sqrt much faster starting with an extended precision float.

n = 10^1000000 - 3^2095903;
(g = Floor[Sqrt[n]]) // AbsoluteTiming // First

6.22126

(h = Floor[Sqrt[N[n, Log[10, n]]]]) // AbsoluteTiming // First

0.086576

g == h

True

using this you maybe want a validity check, (h + 1)^2 > n > h^2 takes very little time.