x86-64平台下如何让SIMD除零运算返回零而非NaN？

问题说明 #

在x86-64平台执行SIMD浮点数除法时，部分除数可能为零，希望除零运算返回零而非NaN。尝试设置MXCSR寄存器的FZ（Flush to Zero）和DAZ（Denormals Are Zero）标志后未生效，疑惑这两个标志的作用是否包含处理除零场景。测试代码如下：

#include <stdio.h>
#include <xmmintrin.h>

int main()
{
    #define CSR_FLUSH_TO_ZERO      (1 << 15)
    #define CSR_DENORMALS_ARE_ZERO (1 << 6)
    unsigned int csr = _mm_getcsr();
    csr |= CSR_FLUSH_TO_ZERO;
    csr |= CSR_DENORMALS_ARE_ZERO;
    _mm_setcsr(csr);

    __m128 a = { 0 };
    __m128 b = { 0 };
    a = _mm_div_ps(a, b);

    float f[4];
    _mm_store_ps(f, a);
    printf("%f\n", f[0]); // prints out 'nan'
}

关键误解澄清 #

FZ和DAZ标志不处理除零运算，它们的作用仅针对非规格化数：

• DAZ（Denormals Are Zero）：将输入的非规格化浮点数当作零参与运算
• FZ（Flush to Zero）：将运算产生的非规格化下溢结果强制置为零
  这两个标志对0/0（结果NaN）或非零数/0（结果±Inf）的场景完全没有影响。

正确实现方案 #

要实现除零返回零，需要手动检测除数为零的位置，然后将对应位置的除法结果替换为零。步骤如下：

1. 生成掩码：用_mm_cmpeq_ps比较除数和零，得到对应位置为全1（表示除数为零）的掩码
2. 执行除法：正常执行SIMD除法，得到包含NaN/Inf的结果
3. 结果替换：用掩码将除数为零位置的结果替换为零，可通过_mm_blendv_ps或位运算组合实现

修改后的代码示例 #

#include <stdio.h>
#include <xmmintrin.h>

int main()
{
    __m128 numerator = _mm_set_ps(1.0f, 2.0f, 0.0f, 3.0f); // 示例分子
    __m128 denominator = _mm_set_ps(0.0f, 4.0f, 0.0f, 0.0f); // 示例分母，包含零

    // 生成除数为零的掩码：对应位置为全1（0xffffffff）表示除数是0
    __m128 zero_mask = _mm_cmpeq_ps(denominator, _mm_setzero_ps());
    // 执行正常除法
    __m128 div_result = _mm_div_ps(numerator, denominator);
    // 将除数为零的位置替换为零：blendv会根据掩码选择第二个操作数的值（零）
    __m128 final_result = _mm_blendv_ps(div_result, _mm_setzero_ps(), zero_mask);

    float f[4];
    _mm_store_ps(f, final_result);
    for (int i = 0; i < 4; i++) {
        printf("f[%d] = %f\n", i, f[i]);
    }
    // 输出结果：
    // f[0] = 0.000000
    // f[1] = 0.500000
    // f[2] = 0.000000
    // f[3] = 0.000000
}

代码说明 #

• _mm_cmpeq_ps：逐元素比较分母和零，相等的位置返回全1的浮点值（对应整数0xffffffff），否则返回全0
• _mm_blendv_ps：根据掩码的每一位选择结果，掩码为1的位置取第二个操作数（零），掩码为0的位置取第一个操作数（正常除法结果）
• 也可以用位运算组合实现：final_result = _mm_or_ps(_mm_andnot_ps(zero_mask, div_result), _mm_and_ps(zero_mask, _mm_setzero_ps()))，效果和_mm_blendv_ps一致

内容的提问来源于stack exchange，提问作者aganm