SSE/AVX：根据每个元素的最小和最大绝对值从两个 __m256 浮点向量中选择

Question

我正在寻找

的高效 AVX (AVX512) 实现

// Given
float u[8];
float v[8];

// Compute
float a[8];
float b[8];

//  Such that
for ( int i = 0; i < 8; ++i )
{
    a[i] = fabs(u[i]) >= fabs(v[i]) ? u[i] : v[i];
    b[i] = fabs(u[i]) <  fabs(v[i]) ? u[i] : v[i];
}

即，我需要从 u 和基于 mask 的 v 按元素 select 到 a，然后到 b 基于 !mask，其中 mask = (fabs(u) >= fabs(v)) 元素。

Answer 1

前几天我遇到了完全相同的问题。我提出的解决方案（仅使用 AVX）是：

// take the absolute value of u and v
__m256 sign_bit = _mm256_set1_ps(-0.0f);
__m256 u_abs = _mm256_andnot_ps(sign_bit, u);
__m256 v_abs = _mm256_andnot_ps(sign_bit, v);
// get a mask indicating the indices for which abs(u[i]) >= abs(v[i])
__m256 u_ge_v = _mm256_cmp_ps(u_abs, v_abs, _CMP_GE_OS);
// use the mask to select the appropriate elements into a and b, flipping the argument
// order for b to invert the sense of the mask
__m256 a = _mm256_blendv_ps(u, v, u_ge_v);
__m256 b = _mm256_blendv_ps(v, u, u_ge_v);

AVX512 等效项为：

// take the absolute value of u and v
__m512 sign_bit = _mm512_set1_ps(-0.0f);
__m512 u_abs = _mm512_andnot_ps(sign_bit, u);
__m512 v_abs = _mm512_andnot_ps(sign_bit, v);
// get a mask indicating the indices for which abs(u[i]) >= abs(v[i])
__mmask16 u_ge_v = _mm512_cmp_ps_mask(u_abs, v_abs, _CMP_GE_OS);
// use the mask to select the appropriate elements into a and b, flipping the argument
// order for b to invert the sense of the mask
__m512 a = _mm512_mask_blend_ps(u_ge_v, u, v);
__m512 b = _mm512_mask_blend_ps(u_ge_v, v, u);

正如 Peter Cordes 在上面的评论中所建议的那样，还有其他方法，比如取绝对值后跟 min/max 然后重新插入符号位，但我找不到任何 shorter/lower 比这个指令序列延迟。

Answer 2

clang 用 -ffast-math 和必要的 __restrict 限定符：https://godbolt.org/z/NMvN1u 自动向量化它做了一个相当合理的工作。和 ABS 的两个输入，比较一次，vblendvps 两次在原始输入上使用相同的掩码，但其他来源以相反的顺序获得最小值和最大值。

这几乎就是我在检查编译器做了什么之前的想法，并查看它们的输出以确认我还没有想清楚的细节。我看不出有什么比这更聪明的了。我不认为我们可以避免分别对 a 和 b 进行 abs() 操作；没有 cmpps 比较谓词来比较大小并忽略符号位。

// untested: I *might* have reversed min/max, but I think this is right.
#include <immintrin.h>
// returns min_abs
__m256 minmax_abs(__m256 u, __m256 v,  __m256 *max_result) {
    const __m256 signbits = _mm256_set1_ps(-0.0f);
    __m256 abs_u = _mm256_andnot_ps(signbits, u);
    __m256 abs_v = _mm256_andnot_ps(signbits, v);  // strip the sign bit

    __m256 maxabs_is_v = _mm256_cmp_ps(abs_u, abs_v, _CMP_LT_OS);  // u < v

    *max_result = _mm256_blendv_ps(v, u, maxabs_is_v);
    return        _mm256_blendv_ps(u, v, maxabs_is_v);
}

你会用 AVX512 做同样的事情，除了你比较到一个掩码而不是另一个矢量。

// returns min_abs
__m512 minmax_abs512(__m512 u, __m512 v,  __m512 *max_result) {
    const __m512 absmask = _mm512_castsi512_ps(_mm512_set1_epi32(0x7fffffff));
    __m512 abs_u = _mm512_and_ps(absmask, u);
    __m512 abs_v = _mm512_and_ps(absmask, v);  // strip the sign bit

    __mmask16 maxabs_is_v = _mm512_cmp_ps_mask(abs_u, abs_v, _CMP_LT_OS);  // u < v

    *max_result = _mm512_mask_blend_ps(maxabs_is_v, v, u);
    return        _mm512_mask_blend_ps(maxabs_is_v, u, v);
}

Clang 以一种有趣的方式编译 return 语句 (Godbolt):

.LCPI2_0:
    .long   2147483647              # 0x7fffffff
minmax_abs512(float __vector(16), float __vector(16), float __vector(16)*):           # @minmax_abs512(float __vector(16), float __vector(16), float __vector(16)*)
    vbroadcastss    zmm2, dword ptr [rip + .LCPI2_0]
    vandps  zmm3, zmm0, zmm2
    vandps  zmm2, zmm1, zmm2
    vcmpltps        k1, zmm3, zmm2
    vblendmps       zmm2 {k1}, zmm1, zmm0
    vmovaps zmmword ptr [rdi], zmm2   ## store the blend result
    vmovaps zmm0 {k1}, zmm1           ## interesting choice: blend merge-masking
    ret

而不是使用另一个 vblendmps，clang 注意到 zmm0 已经有一个混合输入，并使用带有规则向量 vmovaps 的合并掩码。对于 512 位 vblendmps（端口 0 或 5 的单 uop 指令），Skylake-AVX512 的优势为零，但如果 Agner Fog's instruction tables 是正确的，vblendmps x/y/zmm 只会运行s 在端口 0 或 5 上，但是屏蔽的 256 位或 128 位 vmovaps x/ymm{k}, x/ymm 可以运行在任何 p0/p1/p5.

上

两者都是单微指令/单周期延迟，不像 AVX2 vblendvps 基于掩码 vector 是 2 微指令。（所以 AVX512 即使对于 256 位向量也是一个优势）。不幸的是，gcc、clang 或 ICC 的 none 将 _mm256_cmp_ps 转换为 _mm256_cmp_ps_mask，并在使用 -march=skylake-avx512 编译时将 AVX2 内在函数优化为 AVX512 指令。）

s/512/256/ 制作一个使用 AVX512 处理 256 位向量的 minmax_abs512 版本。

Gcc 更进一步，做了

的可疑 "optimization"

    vmovaps zmm2, zmm1        # tmp118, v
    vmovaps zmm2{k1}, zmm0    # tmp118, tmp114, tmp118, u

而不是使用一个混合指令。（我一直在想我看到一个商店后面跟着一个蒙面商店，但是没有，编译器都没有那样混合）。

SSE/AVX：根据每个元素的最小和最大绝对值从两个 __m256 浮点向量中选择

SSE/AVX: Choose from two __m256 float vectors based on per-element min and max absolute value

sse

intrinsics

avx

avx512