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

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

本文关键字:AVX SSE 版本 std copysign 编写 可移植      更新时间:2023-10-16

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

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

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

提前致谢

编辑:

感谢"chtz"这个有用的链接:

https://godbolt.org/z/oY0f7c

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

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

编辑3:

正如"马克西姆·叶戈鲁什金"在他的回答评论中澄清的那样,我对 NaN 的期望是 !(-0.0( 是错误的。NaN 似乎不是一个独特的位模式(见 https://steve.hollasch.net/cgindex/coding/ieeefloat.html(。

非常感谢大家!

AVX 版本用于floatdouble

#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)
}

集会

英特尔内部函数指南


使用 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)
}

尽管如此,clanggcc在编译时计算avx_signbit,并将其替换为从.rodata部分加载的常量,这是IMO的次优。

如果你针对 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 时,您可以制作一个 256 位版本,但您正在处理__m256*向量。