Integer square root

Accelerated Computing CUDA CUDA Programming and Performance
1

Occasionally the square root of an integer is needed, that is ⌊√x⌋ . On many modern systems, the double-precision square root can be computed by the FPU in fewer than 20 cycles, so the square root of a 32-bit unsigned integer x can be computed simply and efficiently by

__device__ uint32_t isqrt (uint32_t x)
{
    return (uint32_t)sqrt ((double) x);
}

The throughput of double-precision operations on GPUs other than certain (semi-)professional models is very low, so the above implementation is not a good idea from a performance perspective. However, we can take advantage of the fact that all GPUs have a multifunction unit which can quickly compute a reciprocal square root to almost full single-precision accuracy, and multiplying that with the original argument yields the square root. The only thing left to do is to ensure that the correct floor is being computed, i.e. the largest integer less than or equal to the mathematical square root. The exhaustively tested implementation below puts this approach to work.

NOTE: It has been pointed out by @ AxelW that notwithstanding a successful exhaustive test, the line s = (uint32_t)fmaf (fr, fa, -0.5f) ; invokes undefined behavior per the C++ standard when the function argument a is zero.

/*
  Copyright (c) 2021, Norbert Juffa
  All rights reserved.

  Redistribution and use in source and binary forms, with or without 
  modification, are permitted provided that the following conditions
  are met:

  1. Redistributions of source code must retain the above copyright 
     notice, this list of conditions and the following disclaimer.

  2. Redistributions in binary form must reproduce the above copyright
     notice, this list of conditions and the following disclaimer in the
     documentation and/or other materials provided with the distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 
  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/* compute floor (sqrt (a)) */
__device__ uint32_t isqrt (uint32_t a)
{
    uint32_t s, rem;
    float fa, fr;

    /* Approximate square root accurately. Make sure it's an underestimate! */
    fa = (float)a;
    asm ("rsqrt.approx.ftz.f32 %0,%1; \n\t" : "=f"(fr) : "f"(fa));
    s = (uint32_t)fmaf (fr, fa, -0.5f) ;
    /* Make sure we got the floor correct */
    rem = a - s * s;
    if (rem >= (2 * s + 1)) s++;
    return (a == 0) ? a : s;
}
2 Likes
2

Here is an implementation of the integer square root for 64-bit integers. Using the reciprocal square root approximation from the multifunction unit as a starting point, it performs one Newton-Raphson iteration in fixed-point arithmetic to arrive at the final result. Since simply going through the double-precision sqrt() without additional checks only works for integers ≤ 253, this might be of interest even on GPUs with high throughput for double-precision operations.

Obviously an exhaustive test is not feasible for 64-bit integers. The code below has passed more than 300 billion random test vectors as well as transition point tests.

[Note: Code below updated to fix functional bugs 12/23/2021]

/*
  Copyright (c) 2021, Norbert Juffa
  All rights reserved.

  Redistribution and use in source and binary forms, with or without 
  modification, are permitted provided that the following conditions
  are met:

  1. Redistributions of source code must retain the above copyright 
     notice, this list of conditions and the following disclaimer.

  2. Redistributions in binary form must reproduce the above copyright
     notice, this list of conditions and the following disclaimer in the
     documentation and/or other materials provided with the distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 
  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

__device__ unsigned long long int umul_wide (unsigned int a, unsigned int b)
{
    unsigned long long int r;
    asm ("mul.wide.u32 %0,%1,%2;\n\t" : "=l"(r) : "r"(a), "r"(b));
    return r;
}

__device__ uint32_t isqrtll (uint64_t a)
{
    uint64_t rem, arg;
    uint32_t b, r, s, scal;

    arg = a;
    /* Normalize argument */
    scal = __clzll (a) & ~1;
    a = a << scal;
    b = a >> 32;
    /* Approximate rsqrt accurately. Make sure it's an underestimate! */
    float fb, fr;
    fb = (float)b;
    asm ("rsqrt.approx.ftz.f32 %0,%1; \n\t" : "=f"(fr) : "f"(fb));
    r = (uint32_t) fmaf (1.407374884e14f, fr, -438.0f);
    /* Compute sqrt(a) as a * rsqrt(a) */
    s = __umulhi (r, b);
    /* NR iteration combined with back multiply */
    s = s * 2;
    rem = a - umul_wide (s, s);
    r = __umulhi ((uint32_t)(rem >> 32) + 1, r);
    s = s + r;
    /* Denormalize result */
    s = s >> (scal >> 1);
    /* Make sure we get the floor correct; can be off by one to either side */
    rem = arg - umul_wide (s, s);
    if ((int64_t)rem < 0) s--;
    else if (rem >= ((uint64_t)s * 2 + 1)) s++;
    return (arg == 0) ? 0 : s;
}
1 Like
3

Thank you for sharing this approach. I’ve tried it out with some other ideas, using it as a case study for learning CUDA.
First of all I learned that profiling performance of device code is tricky, but it seems your approach is among the fastest.
Overestimating the result can also be fast, but the test equations seem to agree less with CUDA; dealing with its potential numeric overflow kills some performance.
Furthermore in both cases the special handling for an input of zero is killing performance, as soon as there is any “if” or “?:”. Not doing the test did speed up execution by about 20%, and — surprisingly — gave correct results on my GPU.
Investigating a little deeper into documentation, I found that casting NaN to int is actually undefined behavior in C++. Even if the result is not used, it theoretically could stop execution. (Ok, well, I know, … but please don’t laugh too hard, corporate safety culture has created stranger requirements.)
For CUDA intrinsic __float2uint_rd() casting from NaN is unfortunately also documented as undefined behavior.
Looking into the PTX, I saw that cvt.rzi.u32.f32 was doing the magic, and this time I was lucky and could find a guarantee in the PTX ISA 8.7 documentation: “In float-to-integer conversion, NaN inputs are converted to 0.” So replacing the cast (uint32_t) by a second asm() command gives a slightly faster and slightly more legal solution.

Here is my code:

__device__ inline
uint32_t uint32_sqrt_asm2_under(uint32_t value)
{
	float value_float = float(value);
	float rsqrt_float; // = rsqrt(value_float); // CUDA intrinsic __frsqrt_rn is slower than PTX rsqrt
	asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(rsqrt_float) : "f"(value_float)); // PTX ISA 8.7: Input +0.0 => Result +Inf
	float s_float = value_float * rsqrt_float - 0.5f; // Always underestimate, NaN for value = 0, because IEEE 754: 0 * Inf = NaN
	uint32_t s; // = isnan(s_float) ? 0 : uint32_t(s_float); // Cast of NaN to int is undefined behavior for C++ and CUDA intrinsic __float2uint_rd.
	asm("cvt.rzi.u32.f32 %0, %1;" : "=r"(s) : "f"(s_float)); // PTX ISA 8.7: "In float-to-integer conversion, NaN inputs are converted to 0."
	// Result s is never too big.
	// Result s is too small if:
	//     (s + 1) * (s + 1) <= value  // but the left side can overflow numerically
	// <=> s * s + 2 * s + 1 <= value
	// <=> value - s * s >= 2 * s + 1  // the left side does not underflow numerically
	// <=> value - s * s > 2 * s
	return value - s * s > 2 * s ? s + 1 : s;
}

Note that the eliminated “+ 1” in the (un)equation of the test does not seem to improve performance.

4

I think you meant to say Infinity, not NaN? For an argument of zero, rsqrt returns infinity, which means that in the case of a zero argument, the line

r = (uint32_t) fmaf (1.407374884e14f, fr, -438.0f);

effectively turns intor = (uint32_t)(INFINITY). I concur with the language lawyer-ly assessment that this results in undefined behavior per the C++ standard:

An rvalue of a floating point type can be converted to an rvalue of an integer type. The conversion truncates; that is, the fractional part is discarded. The behavior is undefined if the truncated value cannot be represented in the destination
type.

The standard quip is that undefined behavior can result in nasal demons, i.e. anything could happen. Historically at least, the behavior of the CUDA compiler has been a bit more sane than that, and this knowledge has probably caused me to subconsciously ignore certain “benign” instances of UB in common code paths, as long as any results emanating from such computation are annulled by special-case handling at the end, with such a code pattern resulting in high performance implementations.

I am curious: Did this code actually “blow up” with the latest CUDA compilers on account of the UB, or did you find the issue by inspection?

I appreciate your efforts in finding an alternate implementation that holds up to language lawyer-ly scrutiny without sacrificing performance. If your intention in sharing this code was to allow others to incorporate it into their projects, note that the lack of an explicitly stated license will preclude that, in particular where corporate lawyers get involved.

[Later: ] Sorry, I noticed belatedly that your analysis pertains to the 32-bit integer variant at the start of the thread, where the line

s = (uint32_t)fmaf (fr, fa, -0.5f) ;

effectively turns intos=(uint32_t)(NAN) when the function argument is zero, because 0 * INFINITY = NAN.

1 Like