Weekend project: Exploring the feasibility of replacing MUFU.RCP and MUFU. RSQ

Recently I was staring at some code that was making fairly extensive use of MUFU operations (the built-in special functions of the GPU like reciprocal and reciprocal square root) and started to ponder how such code would fare on some of the latest GPU architectures where MUFU operations have just 1/8 the throughput of simple FP32 operations.

How might one go about shifting some of the MUFU work back to the FP32/INT32 pipes? A good starting point for such exploration would be the simplest operations, MUFU.RCP and MUFU.RSQ. Obviously, any replacements would have to utilize an absolute minimum of instructions and registers, so table-based methods did not seem suitable. I quickly settled on using starting approximations based on bit manipulation which exploit the semi-logarithmic nature of IEEE-754 floating-point encodings. Such approximations can be generated for general roots, as I explained in this post on Stackoverflow:

The initial low-accuracy approximations can be refined by well-known Newton-Raphson iteration (quadratic convergence) or Halley iteration (cubic convergence). However, instead of using the classic textbook variants of these, one can skew them to achieve somewhat faster convergence by modifying the constants involved. This observation is neither new nor mine and goes back at least to a publication by 3D graphics guru Jim Blinn from almost three decades ago:

James F. Blinn, “Floating-point Tricks.” IEEE Computer Graphics and Applications, Vol. 17, No. 4, Jul.-Aug. 1997, pp. 80-84.

The iterations themselves can be arranged in various mathematically equivalent ways, which, however, result in different numerical properties and resource usage. In the code below I have focused on minimizing resource usage in the specific context of GPUs. All my experiments were performed with the CUDA compiler default setting, which include -ftz=false in particular. Adjustments may be necessary in environments that use -ftz=true or compiler flags that include this, such as -use_fast_math. Fairly recent publications that have examined the design space for these two operations in a more general fashion include:

Leonid Moroz and Volodymyr Samotyy, “Efficient floating-point division for digital signal processing application.” IEEE Signal Processing Magazine, Vol 36, No. 1, Jan. 2019, pp. 159-163

Leonid Moroz, Volodymyr Samotyy, and Oleh Horyachyy, “Modified fast inverse square root and square root approximation algorithms: The method of switching magic constants.” Computation, Vol. 9, No. 2, Feb. 2021, 22 pp.


Since its starting approximation is generated by exploiting the properties of normalized floating-point numbers, the reciprocal computation can at best only support arguments whose magnitude is in [2-126 , 2126 ]. Combination with two skewed Newton-Raphson iterations achieves results with about 18 good bits. The maximum errors occur for arguments whose magnitude is close to 2126, as this causes intermediate computation to veer into subnormal range. By reducing the supported domain slightly, one can achieve results accurate to 20 bits. Compiled for sm_90 using CUDA 13.0, the code generated for these two variants is as follows:

rcp_full_range(float):
 IADD3 R3, -R4, 0x7ef6ffd7, RZ 
 FFMA R0, R3, -R4, 1.0025414228439331055 
 FFMA R3, R3, R0, R3 
 FFMA R4, R3, -R4, 1.0000033378601074219 
 FFMA R4, R3, R4, R3 

rcp_reduced_range(float):
 IADD3 R3, -R4, 0x7ef33be2, RZ 
 FFMA R0, R3, -R4, 1.0012841224670410156 
 FFMA R3, R3, R0, R3 
 FFMA R4, R3, -R4, 1.0000008344650268555 
 FFMA R4, R3, R4, R3 

Clearly, at just five FP32/INT32 instructions, both of these seem to be alternatives worth trying in code characterized by a very high density of MUFU operations that may be (partially) bottlenecked on them on recent GPU architectures.

Since the reciprocal square root is a contracting operation, i.e. given a limited function domain the corresponding range is narrower, arguments in [2-126, 2128 ) can be supported by the bit-twiddling starting approximation. I started with the use of two Newton-Raphson iterations, which provides results with 20 good bits. However, this requires a sequence of nine instructions, which does not make for a performant alternative to MUFU.RSQ. Changing to a single Halley iteration delivers results accurate to 15 bits, but still requires eight instructions, which also doesn’t look suitable from a performance perspective. Only when reducing to a single Newton-Raphson iteration does the code reduce sufficiently (six instructions) that it might be interesting performance-wise, but of course accuracy drops to just 10 bits, which then is unlikely to be attractive from a functional perspective. The code generated by CUDA 13.0 for an sm_90 target for these three variants looks like so:

rsq_bit_trick_2nr(float):
 SHF.R.U32.HI R0, RZ, 0x1, R4 
 FMUL R4, R4, -0.49788725376129150391 
 IADD3 R3, -R0, 0x5f37c185, RZ 
 FMUL R0, R3, R4 
 FFMA R0, R3, R0, 1.5001899003982543945 
 FMUL R3, R3, R0 
 FMUL R4, R4, R3 
 FFMA R4, R3, R4, 1.4978848695755004883 
 FMUL R4, R3, R4 

rsq_bit_trick_halley(float):
 SHF.R.U32.HI R3, RZ, 0x1, R4 
 HFMA2.MMA R5, -RZ, RZ, 1.5361328125, 0.10125732421875 
 IADD3 R3, -R3, 0x5f54401d, RZ 
 FMUL R4, R3, R4 
 FFMA R0, R3, -R4, 0.98139178752899169922 
 FFMA R4, R0, R5, 0.43742907047271728516 
 FMUL R0, R3, R0 
 FFMA R4, R4, R0, R3 

rsq_bit_trick_1nr(float):
 SHF.R.U32.HI R0, RZ, 0x1, R4 
 FMUL R3, R4, -0.53146761655807495117 
 IADD3 R0, -R0, 0x5f331ed1, RZ 
 FMUL R3, R0, R3 
 FFMA R3, R0, R3, 1.5316251516342163086 
 FMUL R4, R0, R3 

Below is the source code for all variants of reciprocal and reciprocal square root computation discussed above.

/*
  Copyright (c) 2025, Norbert Juffa

  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 ERFCRESS 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 reciprocal for |x| in [0x1.000000p-126, 0x1.000000p+126].
   maximum ulp error = 48.7394085;  maximum relative error = 3.39609828e-6
 */
__device__ float rcp_full_range (float x)
{
    float t, y;
    /* initial guess exploiting semi-logarithmic floating-point encoding */
    y = __int_as_float (0x7ef6ffd7 - __float_as_int (x));
    /* refine approximation using two skewed Newton-Raphson iterations */
    t = fmaf (-x, y, 1.00254142f);
    y = fmaf (y, t, y);
    t = fmaf (-x, y, 1.00000334f);
    y = fmaf (y, t, y);
    return y;
}

/* Compute reciprocal for |x| in [0x1.000000p-126, 0x1.e69b7ap+125].
   maximum ulp error = 14.0434610;  maximum relative error = 8.90010789e-7
 */
__device__ float rcp_reduced_range (float x)
{
    float t, y;
    /* initial guess exploiting semi-logarithmic floating-point encoding */
    y = __int_as_float (0x7ef33be2 - __float_as_int (x));
    /* refine approximation using two skewed Newton-Raphson iterations */
    t = fmaf (-x, y, 1.00128412f);
    y = fmaf (y, t, y);
    t = fmaf (-x, y, 1.00000083f);
    y = fmaf (y, t, y);
    return y;
}

// Compute rsqrt for x in [2**(-126), 2**128). One Newton-Raphson iteration.
// maximum ulp error = 12146.4007; maximum relative error = 7.79748218e-4
__device__ float rsq_bit_trick_1nr (float x)
{
    float t, y;
    y = __int_as_float (0x5f331ed1 - ((unsigned int)__float_as_int (x) >> 1));
    t = -5.31467617e-1f * x;
    y = y * fmaf (y * t, y, 1.53162515f);
    return y;
}

// Compute rsqrt for x in [2**(-126), 2**128). One Halley iteration.
// maximum ulp error = 277.529361;  maximum relative error = 1.84172627e-5
__device__ float rsq_bit_trick_halley (float x)
{
    float s, t, y;
    y = __int_as_float (0x5f54401d - ((unsigned int)__float_as_int (x) >> 1));
    s = fmaf (-x * y, y, 9.81391788e-1f);
    t = fmaf (1.61310121e-1f, s, 4.37429070e-1f);
    y = fmaf (t, s * y, y);
    return y;

}

// Compute rsqrt for x in [2**(-126), 2**128). Two Newton-Raphson iterations.
// maximum ulp error = 12.0396649  maximum relative error = 7.33878383e-7
__device__ float rsq_bit_trick_2nr (float x)
{
    float t, y;
    y = __int_as_float (0x5f37c185 - ((unsigned int)__float_as_int (x) >> 1));
    t = -0.497887254f * x;
    y = y * fmaf (y * t, y, 1.50018990f);
    y = y * fmaf (y * t, y, 1.49788487f);
    return y;
}
6 Likes