编写 std::copysign 的可移植 SSE/AVX 版本

2023-11-23

我目前正在使用 SSE 和 AVX 内在函数编写 QR 分解(线性系统求解器)的矢量化版本。其中一个子步骤需要选择与另一个值相反/等于的值的符号。在串行版本中,我为此使用了 std::copysign 。现在我想为 SSE/AVX 寄存器创建一个类似的函数。不幸的是,STL 使用内置函数来实现此目的,因此我不能仅复制代码并将其转换为 SSE/AVX 指令。

我还没有尝试过(所以我现在没有代码可以显示),但我的简单方法是创建一个寄存器,将所有值设置为 -0.0,以便仅设置有符号位。然后我会对源使用 AND 运算来查明其符号是否已设置。此操作的结果可能是 0.0 或 -0.0,具体取决于源的符号。由此,我将创建一个位掩码(使用逻辑运算),我可以将其与目标寄存器(使用另一个逻辑运算)结合起来以相应地设置符号。

但是,我不确定是否有更聪明的方法来解决这个问题。如果有一个用于基本数据类型(如浮点数和双精度数)的内置函数,也许还有一个我错过的内在函数。有什么建议么?

提前致谢

EDIT:

感谢“chtz”提供了这个有用的链接:

https://godbolt.org/z/oY0f7c

所以基本上 std::copysign 编译为一系列 2 AND 操作和后续的 OR 操作。我将为 SSE/AVX 复制此内容并将结果发布在这里,以防有一天其他人需要它:)

EDIT 2:

这是我的工作版本:

__m128 CopySign(__m128 srcSign, __m128 srcValue)
{
    // Extract the signed bit from srcSign
    const __m128 mask0 = _mm_set1_ps(-0.);
    __m128 tmp0 = _mm_and_ps(srcSign, mask0);

    // Extract the number without sign of srcValue (abs(srcValue))
    __m128 tmp1 = _mm_andnot_ps(mask0, srcValue);

    // Merge signed bit with number and return
    return _mm_or_ps(tmp0, tmp1);
}

测试它:

__m128 a = _mm_setr_ps(1, -1, -1, 1);
__m128 b = _mm_setr_ps(-5, -11, 3, 4);

__m128 c = CopySign(a, b);

for (U32 i = 0; i < 4; ++i)
    std::cout << simd::GetValue(c, i) << std::endl;

输出如预期:

5
-11
-3
4

但是,我也尝试了反汇编中的版本

__m128 tmp1 = _mm_andnot_ps(mask0, srcValue);

替换为:

const __m128 mask1 = _mm_set1_ps(NAN);
__m128 tmp1 = _mm_and_ps(srcValue, mask1);

结果很奇怪:

4
-8
-3
4

根据所选的数字,该数字有时可以,有时则不行。标志总是正确的。 由于某种原因,NaN 似乎不是 !(-0.0) 。我记得之前当我尝试将寄存器值设置为 NaN 或特定位模式时遇到了一些问题。也许有人知道问题的根源?

EDIT 3:

正如“Maxim Egorushkin”在他的答案的评论中澄清的那样,我对 NaN 的期望是 !(-0.0) 是错误的。 NaN 似乎不是一个独特的位模式(参见https://steve.hollasch.net/cgindex/coding/ieeefloat.html).

非常感谢大家!


AVX 版本float and double:

#include <immintrin.h>

__m256 copysign_ps(__m256 from, __m256 to) {
    constexpr float signbit = -0.f;
    auto const avx_signbit = _mm256_broadcast_ss(&signbit);
    return _mm256_or_ps(_mm256_and_ps(avx_signbit, from), _mm256_andnot_ps(avx_signbit, to)); // (avx_signbit & from) | (~avx_signbit & to)
}

__m256d copysign_pd(__m256d from, __m256d to) {
    constexpr double signbit = -0.;
    auto const avx_signbit = _mm256_broadcast_sd(&signbit);
    return _mm256_or_pd(_mm256_and_pd(avx_signbit, from), _mm256_andnot_pd(avx_signbit, to)); // (avx_signbit & from) | (~avx_signbit & to)
}

assembly

英特尔内联指南


与AVX2avx_signbit可以在没有常数的情况下生成:

__m256 copysign2_ps(__m256 from, __m256 to) {
    auto a = _mm256_castps_si256(from);
    auto avx_signbit = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cmpeq_epi32(a, a), 31));
    return _mm256_or_ps(_mm256_and_ps(avx_signbit, from), _mm256_andnot_ps(avx_signbit, to)); // (avx_signbit & from) | (~avx_signbit & to)
}

__m256d copysign2_pd(__m256d from, __m256d to) {
    auto a = _mm256_castpd_si256(from);
    auto avx_signbit = _mm256_castsi256_pd(_mm256_slli_epi64(_mm256_cmpeq_epi64(a, a), 63));
    return _mm256_or_pd(_mm256_and_pd(avx_signbit, from), _mm256_andnot_pd(avx_signbit, to)); // (avx_signbit & from) | (~avx_signbit & to)
}

尽管如此,两者clang and gcc计算avx_signbit在编译时并将其替换为从加载的常量.rodata部分,在我看来,这是次优的。

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

编写 std::copysign 的可移植 SSE/AVX 版本 的相关文章

随机推荐