平方根倒数算法对比

这是游戏雷神之锤中的一段代码，出自于 John Carmack 之手，在那个硬件受限的年代，这段代码彻底地改变了整个游戏界。

初见卡马克快速平方根倒数代码，我完全没看懂它在做什么，只听说它计算速度巨快，后来系统地学习了这个算法，我想对比一下它究竟多快？首先介绍卡马克平方根倒数算法，再与 math.h 数学库、CPU 指令集详尽对比。

一、目标

平方根倒数是为了计算如下公式：

雷神之锤游戏的渲染会经常遇到曲面物体，对于计算机而言，画直线轻而易举，但曲线就会涉及大量的运算，光照射到曲面上，会反射，而要获得到反射光的效果，就需要得知曲面上入射点的单位法向量【点击查看 DezemingFamily 介绍的原理】，这个过程速度越快，游戏的帧率就会越高。

尤其当构成曲面（曲线）但平面（直线）过多时，大量光线照射到墙体，CPU 每次计算单位法向量都慢一拍，累积下来，游戏帧率就会大幅度下降。虽然现状有 NVIDIA 显卡并行计算，但当时只能靠 CPU 算力。同样的道理，平方根倒数算法也可以运用在余弦相似度计算：

二、卡马克平方根倒数算法

简单来说，该算法利用了 IEEE 754 浮点型数存储形式和牛顿迭代定理，实现线性时间复杂度计算平方根倒数，并且有着不错的精度。

32 位浮点数在计算机中的存储方式如下：

最高位第 31 位是 sign 符号位，用 0 代表正数，用 1 代表负数；中间第 23 ~ 30 位是 exponent 指数位；后面 23 个位是 mantissa 尾数位。

任意一个在作用域内的浮点型数 $a$ 都可以用二进制科学记数法表示：

譬如：0 01111100 11000000000000000000000 这样的浮点型存储，对应的浮点数为：

其中，指数 8 位 01111100 对应十进制 124，尾数 23 位用小 $m$ 表示为：

同理，可以反推，譬如浮点数 -3.1415 转化成 IEEE 754 存储。符号位是 1，代表负数，整数部分是 3，尾数部分是 0.1415。整数反复除以二，小数反复乘以二，都转化成二进制：

由于 0.1415 始终乘不尽，所以我用省略号表示了，因此 3.1415 的二进制：

再转为二进制科学记数法：

又因为 E - 127 要等于 1，所以 E 等于 128，对应二进制：

综上，浮点数 -3.1415 对应 IEEE 754 计算机存储为（其中省略号一直凑齐 23 个）：

不妨设 $\vartheta$ 是浮点数 $a$ 的平方根倒数解，即：

易知解 $\vartheta$ 也是浮点型，因此也可以用 IEEE 754 表示法：

由于 $M\in[0,1)$ 之间，因此有如下性质：

根据切比雪夫最佳直线逼近，可得到 $\varepsilon$ 的最优值：

因此，可进一步得到：

则：

卡马克在这里，取的幻数是 0x5f3759df，与切比雪夫最优逼近的幻数 0x5f37bcb6 相差十进制 25303 个单位。

其中，$m_{\vartheta}+2^{23}E_{\vartheta}$ 正是浮点型数的 IEEE 754 二进制表示，譬如先前示例：

0.21875 对应着 0 01111100 11000000000000000000000，则：

从 C 语言编程角度而言，$\dfrac{1}{2}(m_a+2^{23}E_a)$ 等价于二进制位运算右移一个单位。

i = 0x5f3759df - (i >> 1);   // what the fuck?

至于浮点型和二进制互相转化，对应源码中：

i = * ( long * ) &y;     // evil floating point bit level hacking
y = * ( float * ) &i;

只不过需要注意的是，当年 long 型代表 32 位，而如今 long 普遍是 64 位，所以今天我们需要修改成：i = *(int *)&y 即可。

又因为平方根倒数的精度不够，所以卡马克在后面补充了牛顿迭代，进一步逼近更精确的数值解。

y = y * ( 1.5F - ( 0.5F * number * y * y ) );     // 1st iteration

二、性能对比

在 Windows 10/11 AMD Ryzen 5 5600H、Linux Ubuntu Intel(R) Xeon(R) Silver 4310 和 macOS Sonoma M2 三台机器测试自然指数 e 的平方根倒数，计算十亿次，其中 GCC 使用 O3 优化。

结果表明：在 Windows 操作系统下，无论是否开启 GCC 编译优化，卡马克平方根倒数快速幂算法优于数学库的标准函数，与牛顿迭代一次的 CPU 指令集策略速度几乎持平（但牛顿迭代一次的 CPU 指令集精度更高），卡马克策略虽然慢于直接使用 CPU 指令集策略，但精度比直接使用 CPU 指令集更高。在 Linux 操作系统中，无论是否开启 GCC 编译优化，卡马克策略均最慢，开启 O3 优化后，牛顿迭代一次的 CPU 指令集速度和精度都是最好的。在 macOS 系统下，不开启 GCC 编译优化的卡马克策略慢于数学标准库，但快于 CPU 指令集及其一次牛顿迭代这两种策略，不过从精度上而言，卡马克策略优于数学标准库，在开启 O3 优化后，四种策略速度得到大幅度提升，并且均不相上下地持平。

/*
Windows or macOS:
    >>> gcc rsqrt.c
    >>> gcc -O3 rsqrt.c
Linux:
    >>> gcc rsqrt.c -lm
    >>> gcc -O3 rsqrt.c -lm
*/

# include <math.h>
# include <time.h>
# include <stdio.h>

# ifdef __x86_64__
    # include <xmmintrin.h>
# elif __arm64__ || __aarch64__
    # include <arm_neon.h>
# endif

void Timer(const char *name, float (*func)(float), float x, int epoch) {
    clock_t start = clock();
    volatile float y;   // Prevents the gcc compiler from optimizing the result.
    for (int i = 0; i < epoch; i++) y = func(x);
    clock_t end = clock();
    double elapse = (double)(end - start) / CLOCKS_PER_SEC;
    printf("time = %.5lfs (epoch = %d)  -->  %s(%f) = %f\n", elapse, epoch, name, x, y);
}

float Q_rsqrt(float x) {
    float x_half = 0.5F * x;
    int i = *(int *)&x;     // A 32-bit (4-byte) int must be used instead of long (8-byte).
    i = 0x5f3759df - (i >> 1);
    x = *(float *)&i;
    x = x * (1.5F - x_half * x * x);
    x = x * (1.5F - x_half * x * x);
    return x;
}

float Math_rsqrt(float x) {
    return 1 / sqrtf(x);
}

# ifdef __x86_64__
    float SSE_rsqrt(float x) {
        __m128 vector_x = _mm_set_ss(x);
        __m128 vector_y = _mm_rsqrt_ss(vector_x);
        return _mm_cvtss_f32(vector_y);
    }

    float SSE_Newton_rsqrt(float x) {
        __m128 vector_x = _mm_set_ss(x);
        __m128 vector_y = _mm_rsqrt_ss(vector_x);
        float y = vector_y[0];
        y = y * (1.5F - 0.5F * x * y * y);
        return y;
    }
# endif

# if defined(__arm64__) || defined(__aarch64__)
    float NEON_rsqrt(float x) {
        float32x4_t vector_x = {x, 1.0F, 1.0F, 1.0F};
        float32x4_t vector_y = vrsqrteq_f32(vector_x);
        return vgetq_lane_f32(vector_y, 0);
    }

    float NEON_Newton_rsqrt(float x) {
        float32x4_t vector_x = {x, 1.0F, 1.0F, 1.0F};
        float32x4_t vector_y = vrsqrteq_f32(vector_x);
        float y = vgetq_lane_f32(vector_y, 0);
        y = y * (1.5F - 0.5F * x * y * y);
        return y;
    }
# endif

float cos_similarity(int dim, float *vector_a, float *vector_b) {
    float numerator = 0.0F;
    float module_a = 0.0F;
    float module_b = 0.0F;
    for (int i = 0; i < dim; i++) {
        numerator = numerator + vector_a[i] * vector_b[i];
        module_a = module_a + vector_a[i] * vector_a[i];
        module_b = module_b + vector_b[i] * vector_b[i];
    }
    return numerator * Q_rsqrt(module_a * module_b);
}

int main(int argc, char *argv[], char *env[]) {
    Timer("Q_rsqrt", Q_rsqrt, 2.71828F, 1e9);
    Timer("Math_rsqrt", Math_rsqrt, 2.71828F, 1e9);

    # ifdef __x86_64__
        Timer("SSE_rsqrt", SSE_rsqrt, 2.71828F, 1e9);
        Timer("SSE_Newton_rsqrt", SSE_Newton_rsqrt, 2.71828F, 1e9);
    # elif __arm64__ || __aarch64__
        Timer("NEON_rsqrt", NEON_rsqrt, 2.71828F, 1e9);
        Timer("NEON_Newton_rsqrt", NEON_Newton_rsqrt, 2.71828F, 1e9);
    # endif

    float vector_a[4] = {1.2, 2.4, 3.6, 4.8};
    float vector_b[4] = {4.8, 3.6, 2.4, 1.2};
    float y = cos_similarity(4, vector_a, vector_b);
    printf("cos similarity = %f\n", y);

    return 0;
}