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

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];
}

也就是说，我需要按元素选择 a来自u和v基于mask ，并进入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 在上面的评论中所建议的那样，还有其他方法，例如取绝对值，后跟最小值/最大值，然后重新插入符号位，但我找不到比这更短/更低延迟的方法这一系列指令。

<小时/>

实际上，还有另一种通过 _mm512_range_ps() 内部函数使用 AVX512DQ 的 VRANGEPS 的方法。 Intel's intrinsic guide描述如下:

Calculate the max, min, absolute max, or absolute min (depending on control in imm8) for packed single-precision (32-bit) floating-point elements in a and b, and store the results in dst. imm8[1:0] specifies the operation control: 00 = min, 01 = max, 10 = absolute max, 11 = absolute min. imm8[3:2] specifies the sign control: 00 = sign from a, 01 = sign from compare result, 10 = clear sign bit, 11 = set sign bit.

请注意，上面似乎有一个拼写错误；实际上，如果您查看详细信息，则 imm8[3:2] == 10 是“绝对最小值”，而 imm8[3:2] == 11 是“绝对最大值”每个元素的操作:

CASE opCtl[1:0] OF
0: tmp[31:0] := (src1[31:0] <= src2[31:0]) ? src1[31:0] : src2[31:0]
1: tmp[31:0] := (src1[31:0] <= src2[31:0]) ? src2[31:0] : src1[31:0]
2: tmp[31:0] := (ABS(src1[31:0]) <= ABS(src2[31:0])) ? src1[31:0] : src2[31:0]
3: tmp[31:0] := (ABS(src1[31:0]) <= ABS(src2[31:0])) ? src2[31:0] : src1[31:0]
ESAC

CASE signSelCtl[1:0] OF
0: dst[31:0] := (src1[31] << 31) OR (tmp[30:0])
1: dst[31:0] := tmp[63:0]
2: dst[31:0] := (0 << 31) OR (tmp[30:0])
3: dst[31:0] := (1 << 31) OR (tmp[30:0])
ESAC

RETURN dst

因此，只需两条指令即可获得相同的结果:

auto a = _mm512_range_ps(v, u, 0x7);  // 0b0111 = sign from compare result, absolute max
auto b = _mm512_range_ps(v, u, 0x6);  // 0b0110 = sign from compare result, absolute min

参数顺序 (v, u) 有点不直观，但为了在元素具有相等绝对值的情况下获得与您在 OP 中描述的相同行为，这是必需的(即，u 中的值传递到 a，v 传递到 b)。

在 Skylake 和 Ice Lake Xeon 平台上(可能是任何具有双 FMA 单元的 Xeon 平台？)，VRANGEPS 的吞吐量为 2，因此两个检查可以同时发出和执行，但有延迟4 个周期。这只是对原始方法的适度延迟改进，但吞吐量更好，并且需要更少的指令/微指令/指令缓存空间。