有效 load/compute/pack 64 双重比较结果在 uint64_t 位掩码中

Question

我想 load/compare/pack 尽可能高效地将 64 次双重比较的结果放入 uint64_t 位掩码中。

我目前的方法是通过 AVX2 使用 _mm256_cmp_pd 比较 2*2 对。两次（=8）比较的结果使用 _mm256_movemask_pd 转换为位图，并通过 a|b<<4 分配为字节到联合（1x uint64_t / 8 uint8_t）以保存几个 shifts/or。

这个例子可能有助于形象化

union ui64 {
    uint64_t i64;
    struct { uint8_t i0,i1,i2,i3,i4,i5,i6,i7; } i8;
};
static inline uint64_t cmp64 (double* in0, double* in1) {
    ui64 result;
    // +0
    result.i8.i0 =
        _mm256_movemask_pd(
            _mm256_cmp_pd(
            _mm256_load_pd(in0 + 0), 
            _mm256_load_pd(in1 + 0), 
            _CMP_LT_OQ)) |
        _mm256_movemask_pd(
            _mm256_cmp_pd(
                _mm256_load_pd(in0 + 4), 
                _mm256_load_pd(in1 + 4),
                _CMP_LT_OQ)) << 4;

    // +8
    // load, compare, pack n+8 each 'iteration' using result.i8.i1,...i7
    // ...

    return result.i64;
}

compress&set 的变体看起来很直接，但使用较慢的指令：1x _mm256_set_m128 和 2x_mm256_cvtpd_ps 与标量 << 和 |像这样

_mm256_movemask_ps(_mm256_set_m128(
    _mm256_cvtpd_ps(v0),
    _mm256_cvtpd_ps(v1)));

使用的 CPU 是 Zen 1（最大 AVX2），不确定 GPU 使用（Nvidia）是否是一个选项。

请分享您的想法。

Answer 1

在 CPU 和 full-width AVX2（如 Zen2 或 Haswell / Skylake）上，你可能会用 vpackssdw / vpacksswb 水平压缩从 qwords 到字节每次缩小一半。因此，总共 8 个输入向量将成为您在 (_mm256_movemask_epi8) 上执行 vpmovmskb 的一个向量。 VCMPPD 结果是 all-ones (-1) 保持 -1，或 all-zeros 保持 0，在 qword 的两半中，即使您使用更窄的包元素大小。但是那个打包是 in-lane （在向量的 128 位一半内），所以在最终打包成字节之后你需要一个 vpshufb + vpermd 来在 [=12 之前按顺序获取字节=]. （AMD 在 Zen3 之前没有快速 pdep，否则如果你没有进行 lane-crossing 修复随机播放，你可以使用它来交错位对。）
4:1 包见； 8:1 使最终洗牌更加复杂，除非我们更早地进行更多洗牌，而双字块足够小。

（我正在使用 asm 助记符名称，因为它们比内在函数更短，更易于阅读，而且您需要在指令表中查找任何内容以找出花费多少 uops；https://uops.info/ or https://agner.org/optimize/)

但是每个 256 位 SIMD 操作花费 2 微秒，你可能在 Zen 1 上只用 vmovmskpd 和标量 bit-shift / OR 就做得很好。如果周围的代码都是向量，让这些 uops 使用标量整数 ALU 是很好的。 front-end 的宽度为 6 微指令，或 5 条指令，以较少者为准，但每个整数和 SIMD ALU 管道只有 4 条，因此理想情况下，较早和较晚的代码可以很好地重叠执行。（并且某些特定的 ALU 单元的吞吐量甚至更有限，例如，这些洗牌仅在 4 个端口中的 2 个上进行。）

或者一步向量打包然后 _mm256_movemask_ps？ Lane-crossing 洗牌在 Zen 1 上相对昂贵。但不是太坏：vpermq（或 vpermpd）只有 3 微指令和 2 个周期吞吐量，对比 vpackssdw 吞吐量为 1c 的 2 微指令。（以及 vpermd 的 4c 吞吐量的 3 微指令。）

假设 vpacksswd ymm 使用与 XMM 版本相同的端口，即 FP1 / FP2。因此它可以与 vcmppd 部分重叠，而在 FP01 上可以运行。（如果不与其他指令混合，YMM 版本也是 2 微指令，1c 吞吐量。）

https://uops.info/ 在某些 AMD CPU 上无法获得 multi-uop 指令的详细程度，但我们可以假设非-lane-crossing 版本只是与 XMM 版本相同的 uop 中的两个，它确实有该数据。

您很可能不想使用 _mm256_cvtpd_ps，它会花费 shuffle uops 和 FP->FP 转换。这需要 2 微指令，但只有一个输入向量，而不是两个。将比较结果解释为 -NaN double，您很可能会得到一个 float -NaN，因此它实际上可能有助于正确性。在大多数 CPU 上，这种方式肯定更慢。
在 Zen1 上它有 2 个周期的吞吐量，这是每个单个输入向量而不是一对向量。

使用 4x vpackssdw 我们可以将 8 个向量减少到 4 个。
然后 2x vpackssdw ymm 减少到 2 个向量。
然后 1x vpacksswb ymm 减少为 1 个向量，字节对的顺序错误。

对于 Zen 1，可能从 4 个输入向量开始，在减少到一个 YMM 之后，用 vextracti128 将其分成两半，这在 Zen 1 上只有一个 uop，对于任何端口（因为 YMM 寄存器的两半已经分别存储在物理寄存器中）。然后 vpacksswb 将两半放在一起（1 uop），设置 vpshufb xmm （1 uop）以按正确的顺序放置字节对。这为 vpmovmskb 做好了准备。所以唯一的lane-crossing洗牌只是一个摘录。

或者不是获取 16 位位图块，您可以执行上述两次，然后 vinserti128 ymm, xmm, 1（2 微指令，0.67c 吞吐量）/vpmovmskb ymm（1 微指令）以获得一个 32 位的位图块。这 3 个 uops 替换了 2x vpmovmskb xmm / shl / or，所以你节省了一个 uop，并且对它们可以运行的矢量 ALU 端口具有很好的灵活性。虽然是vector ALU压力更大。

Answer 2

举个例子。它使用最有效的指令将比较结果打包成字节，每 32 个数字洗牌一次，并使用 _mm256_movemask_epi8 一次产生 32 位。

// Compare 4 numbers, return 32 bytes with results of 4 comparisons:
// 00000000 11111111 22222222 33333333
inline __m256d compare4( const double* a, const double* b )
{
    return _mm256_cmp_pd( _mm256_load_pd( a ), _mm256_load_pd( b ), _CMP_LT_OQ );
}

// Compare 8 numbers, combine 8 results into the following bytes:
// 0000 4444 1111 5555  2222 6666 3333 7777
inline __m256i compare8( const double* a, const double* b )
{
    __m256 c0 = _mm256_castpd_ps( compare4( a, b ) );
    __m256 c1 = _mm256_castpd_ps( compare4( a + 4, b + 4 ) );
    return _mm256_castps_si256( _mm256_blend_ps( c0, c1, 0b10101010 ) );
}

// Compare 16 numbers, combine 16 results into following bytes:
// 00 44 11 55  88 CC 99 DD  22 66 33 77  AA EE BB FF
inline __m256i compare16( const double* a, const double* b )
{
    __m256i c0 = compare8( a, b );
    __m256i c1 = compare8( a + 8, b + 8 );
    return _mm256_packs_epi32( c0, c1 );
}

inline uint32_t compare32( const double* a, const double* b )
{
    // Compare 32 numbers and merge them into a single vector
    __m256i c0 = compare16( a, b );
    __m256i c1 = compare16( a + 16, b + 16 );
    __m256i src = _mm256_packs_epi16( c0, c1 );

    // We got the 32 bytes, but the byte order is screwed up in that vector:
    // 0   4   1   5   8   12  9   13  16  20  17  21  24  28  25  29
    // 2   6   3   7   10  14  11  15  18  22  19  23  26  30  27  31
    // The following 2 instructions are fixing the order

    // Shuffle 8-byte pieces across the complete vector
    // That instruction is relatively expensive on most CPUs, but we only doing it once per 32 numbers
    src = _mm256_permute4x64_epi64( src, _MM_SHUFFLE( 3, 1, 2, 0 ) );

    // The order of bytes in the vector is still wrong:
    // 0    4   1   5   8  12   9  13    2   6   3   7  10  14  11  15
    // 13  16  20  17  21  24  28  25   18  22  19  23  26  30  27  31
    // Need one more shuffle instruction

    const __m128i perm16 = _mm_setr_epi8(
        0, 2, 8, 10,  1, 3, 9, 11,   4, 6, 12, 14,  5, 7, 13, 15 );
    // If you calling this in a loop and everything is inlined,
    // the shuffling vector should be pre-loaded outside of the loop.
    const __m256i perm = _mm256_broadcastsi128_si256( perm16 );

    // Shuffle the bytes
    src = _mm256_shuffle_epi8( src, perm );

    // The order of bytes is now correct, can use _mm256_movemask_epi8 to make 32 bits of the result
    return (uint32_t)_mm256_movemask_epi8( src );
}


uint64_t compareAndPack64( const double* a, const double* b )
{
    uint64_t low = compare32( a, b );
    uint64_t high = compare32( a + 32, b + 32 );
    return low | ( high << 32 );
}

有效 load/compute/pack 64 双重比较结果在 uint64_t 位掩码中

Efficiently load/compute/pack 64 double comparison results in uint64_t bitmask

c++

optimization

simd

avx

avx2