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

Question

我目前正在使用 SSE 和 AVX 内在函数编写 QR 分解（线性系统求解器）的矢量化版本。其中一个子步骤需要 select 值的符号 opposite/equal 到另一个值。在串行版本中，我为此使用了 std::copysign。现在我想为 SSE/AVX 寄存器创建一个类似的函数。不幸的是，STL 为此使用了一个内置函数，所以我不能只复制代码并将其转换为 SSE/AVX 指令。

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

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

提前致谢

编辑：

感谢 "chtz" 这个有用 link:

https://godbolt.org/z/oY0f7c

所以基本上 std::copysign 编译为一系列 2 个 AND 操作和一个后续的 OR。我将在此处为 SSE/AVX 和 post 重现此结果，以防有一天其他人需要它:)

编辑 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;

输出符合预期：

但是，我也尝试了

反汇编的版本

__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 或特定位模式时遇到了一些问题。也许有人知道问题的根源？

编辑 3：

正如 'Maxim Egorushkin' 在他的回答的评论中澄清的那样，我对 NaN 的期望是 !(-0.0) 是错误的。 NaN 似乎不是唯一的位模式（参见 https://steve.hollasch.net/cgindex/coding/ieeefloat.html）。

非常感谢大家！

Answer 1

float 和 double 的 AVX 版本：

#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

The Intel Intrinsics Guide

使用 AVX2 avx_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 和 gcc 都在编译时计算 avx_signbit 并将其替换为从 .rodata 部分加载的常量，即 IMO，sub-最佳。

Answer 2

如果您以 icc 为目标，我认为这是一个比公认的稍微好一点的版本：

__m256d copysign_pd(__m256d from, __m256d to) {
    __m256d const avx_sigbit = _mm256_set1_pd(-0.);
    return _mm256_or_pd(_mm256_and_pd(avx_sigbit, from), _mm256_andnot_pd(avx_sigbit, to));
}

它使用 _mm256_set1_pd 而不是广播内在函数。在 clang 和 gcc 上，这主要是一次清洗，但在 icc 上，广播版本实际上将一个常量写入堆栈，然后从中广播，这很糟糕。

Godbolt 显示 AVX-512 代码，将 -march= 调整为 -march=skylake 以查看 AVX2 代码。

这是一个未经测试的 AVX-512 版本，它直接使用 vpterlogdq，在 icc 和 clang 上编译为单个 vpterlogd 指令（gcc 包括一个单独的广播）：

__m512d copysign_pd_alt(__m512d from, __m512d to) {
    const __m512i sigbit = _mm512_castpd_si512(_mm512_set1_pd(-0.));
    return _mm512_castsi512_pd(_mm512_ternarylogic_epi64(_mm512_castpd_si512(from), _mm512_castpd_si512(to), sigbit, 0xE4));
}

当启用 AVX-512 但您正在处理 __m256* 向量时，您可以制作一个 256 位版本。

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

Writing a portable SSE/AVX version of std::copysign

c++

sse

x86-64

simd

avx