AVX2中_mm256_cmp_pd返回NaN而非1.0的问题及解决方法

问题概述 #

希望对比较操作进行向量化处理用作掩码，但使用AVX2的_mm256_cmp_pd指令时，返回结果为NaN而非预期的1.0。相关代码如下：

__m256d _numberToCompare = _mm256_set1_pd(1.0);
__m256d _compareConditions = _mm256_set_pd(0.0, 1.0, 2.0, 3.0);

__m256d _result = _mm256_cmp_pd(_numberToCompare, _compareConditions, _CMP_LT_OQ); //a < b ordered (non-signalling) 
alignas(8) double res[4];
_mm256_store_pd(&res[0], _result);
for (auto i : res) {
    std::cout << i << '\t';
}
   
__m256d _result2 = _mm256_cmp_pd(_numberToCompare, _compareConditions, _CMP_LE_OQ); //a <= b ordered (non-signalling)   
alignas(8) double res2[4];
_mm256_store_pd(&res2[0], _result2);
for (auto i : res2) {
    std::cout << i << '\t';
}
std::cout << '\n';

预期结果 #

0 0 1 1
0 1 1 1

实际结果 #

-nan    -nan    0       0
-nan    -nan    -nan    0

问题解答 #

1. 为何比较结果是NaN？ #

_mm256_cmp_pd的返回值并非0.0/1.0的浮点值，而是全0或全1的位模式：

• 比较结果为真时，对应双精度元素的所有64位都会被设为1，这个位模式解释为双精度浮点数就是NaN；
• 比较结果为假时，所有位设为0，对应浮点值0.0。

另外你看到的结果顺序和预期相反，是因为_mm256_set_pd是逆序填充寄存器的：你传入的0.0,1.0,2.0,3.0会被存储为寄存器中的[3.0,2.0,1.0,0.0]，导致输出顺序颠倒。

2. 如何得到预期的结果？ #

分两种场景处理：

场景一：需要浮点格式的0/1结果

如果确实需要得到0.0和1.0的浮点值，可将比较掩码转换为整数后再转回浮点：

// 将比较得到的掩码从浮点向量转换为整数向量
__m256i mask = _mm256_castpd_si256(_result);
// 右移63位，将全1的整数位模式转为1，全0保持0
mask = _mm256_srli_epi64(mask, 63);
// 转换回双精度浮点向量
__m256d float_result = _mm256_cvtepi64_pd(mask);

场景二：用作掩码（更高效）

如果只是为了后续作为掩码使用（比如_mm256_blendv_pd这类指令），直接使用_mm256_cmp_pd返回的掩码即可——这类指令会识别全0/全1的位模式作为掩码，无需转换为浮点格式的0/1。

同时修正_mm256_set_pd的参数顺序，让结果顺序符合预期：

__m256d _compareConditions = _mm256_set_pd(3.0, 2.0, 1.0, 0.0); // 逆序传入，寄存器中顺序为0.0,1.0,2.0,3.0

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