使用 AVX512 或 AVX2 计算所有打包 32 位整数总和的最快方法

Fastest method to calculate sum of all packed 32-bit integers using AVX512 or AVX2

我正在寻找一种最佳方法来计算 __m256i__m512i 中所有打包的 32 位整数的总和。为了计算 n 元素的总和,我经常使用 log2(n) vpadddvpermd 函数,然后提取最后结果。但是,这不是我认为的最佳选择。

编辑:best/optimal 减少 speed/cycle。

相关:如果您正在寻找不存在的 _mm512_reduce_add_epu8,请参见 vpsadbw,因为 qwords 中的 hsum 比改组更有效。

没有 AVX512,请参阅下面的 hsum_8x32(__m256i) 了解没有英特尔 reduce_add 辅助函数的 AVX2。 reduce_add 无论如何都不一定使用 AVX512 进行最佳编译。


immintrin.h 中有一个 int _mm512_reduce_add_epi32(__m512i) 内联函数。你不妨使用它。 (它编译以随机播放和添加指令,但比 vpermd 更有效,就像我在下面描述的那样。)AVX512 没有引入任何新的 硬件 支持对于水平求和,只有这个新的辅助函数。 仍然需要尽可能避免或跳出循环。

GCC 9.2 -O3 -march=skylake-avx512 编译一个包装器,调用它如下:

        vextracti64x4   ymm1, zmm0, 0x1
        vpaddd  ymm1, ymm1, ymm0
        vextracti64x2   xmm0, ymm1, 0x1   # silly compiler, vextracti128 would be shorter
        vpaddd  xmm1, xmm0, xmm1
        vpshufd xmm0, xmm1, 78
        vpaddd  xmm0, xmm0, xmm1

        vmovd   edx, xmm0
        vpextrd eax, xmm0, 1              # 2x xmm->integer to feed scalar add.
        add     eax, edx
        ret

提取两次以提供标量加法是有问题的;它需要 p0 和 p5 的 uops,所以它等同于常规洗牌 + a movd.

Clang 不会那样做;它再进行一步洗牌/SIMD 添加以减少到 vmovd 的单个标量。两者的性能分析见下文


有一个 VPHADDD 但你不应该在两个输入相同的情况下使用它。 (除非您正在优化代码大小而不是速度)。转置和求和多个向量可能很有用,从而产生一些结果向量。您可以通过为 phadd 提供 2 个 不同的 输入来做到这一点。 (除了它在 256 位和 512 位上变得混乱,因为 vphadd 仍然只是在车道内。)

是的,你需要 log2(vector_width) 洗牌和 vpaddd 指令。(所以这不是很有效;避免内部循环内的水平求和。累加例如,垂直直到循环结束)。


所有 SSE / AVX / AVX512 的通用策略

您想从 512 -> 256、然后 256 -> 128 连续缩小,然后在 __m128i 内随机排列,直到缩小到一个标量元素 。据推测,未来的某些 AMD CPU 会将 512 位指令解码为两个 256 位微指令,因此减少宽度是一个巨大的胜利。更窄的指令可能会消耗更少的功率。

您的随机播放可以直接控制操作数,而不是 vpermd. 的向量,例如VEXTRACTI32x8vextracti128vpshufd。 (或 vpunpckhqdq 以保存立即数的代码大小。)

参见Fastest way to do horizontal SSE vector sum (or other reduction)(我的回答还包括一些整数版本)。

此通用策略适用于所有元素类型:浮点型、双精度型和任意大小的整数

特殊情况:

  • 8位整数:从vpsadbw开始,效率更高并避免溢出,但随后继续为64位整数。

  • 16 位整数:首先使用 pmaddwd 扩大到 32(_mm256_madd_epi16 使用 set1_epi16(1)): - 更少即使您不关心避免溢出的好处也是 uops,除了 Zen2 之前的 AMD,其中 256 位指令至少花费 2 uops。但是然后你继续 32 位整数。

32位整数可以这样手动完成,用一个SSE2函数由AVX2函数在减少到__m128i后调用,在减少到__m256i后由AVX512函数调用。这些调用当然会在实践中内联。

#include <immintrin.h>
#include <stdint.h>

// from my earlier answer, with tuning for non-AVX CPUs removed
// static inline
uint32_t hsum_epi32_avx(__m128i x)
{
    __m128i hi64  = _mm_unpackhi_epi64(x, x);           // 3-operand non-destructive AVX lets us save a byte without needing a movdqa
    __m128i sum64 = _mm_add_epi32(hi64, x);
    __m128i hi32  = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1));    // Swap the low two elements
    __m128i sum32 = _mm_add_epi32(sum64, hi32);
    return _mm_cvtsi128_si32(sum32);       // movd
}

// only needs AVX2
uint32_t hsum_8x32(__m256i v)
{
    __m128i sum128 = _mm_add_epi32( 
                 _mm256_castsi256_si128(v),
                 _mm256_extracti128_si256(v, 1)); // silly GCC uses a longer AXV512VL instruction if AVX512 is enabled :/
    return hsum_epi32_avx(sum128);
}

// AVX512
uint32_t hsum_16x32(__m512i v)
{
    __m256i sum256 = _mm256_add_epi32( 
                 _mm512_castsi512_si256(v),  // low half
                 _mm512_extracti64x4_epi64(v, 1));  // high half.  AVX512F.  32x8 version is AVX512DQ
    return hsum_8x32(sum256);
}

请注意,这使用 __m256i hsum 作为 __m512i 的构建块;先进行车道内操作没有任何好处。

这可能是一个非常小的优势:车道内洗牌比车道交叉具有更低的延迟,因此它们可以更早执行 2 个周期并更早地离开 RS,并且类似地更早地退出 ROB。但是即使你这样做了,更高延迟的洗牌也会在几条指令之后出现。因此,如果此 hsum 在关键路径上(阻止退休),您可能会提前 2 个周期将一些独立指令放入后端。

但是越早减少到更窄的向量宽度通常是好的,也许越早从系统中获取 512 位 uops,这样 CPU 就可以重新激活端口 1 上的 SIMD 执行单元,如果你不是'马上做更多的 512 位工作。

使用 GCC9.2 -O3 -march=skylake-avx512

on Godbolt 编译为这些指令
hsum_16x32(long long __vector(8)):
        vextracti64x4   ymm1, zmm0, 0x1
        vpaddd  ymm0, ymm1, ymm0
        vextracti64x2   xmm1, ymm0, 0x1   # silly compiler uses a longer EVEX instruction when its available (AVX512VL)
        vpaddd  xmm0, xmm0, xmm1
        vpunpckhqdq     xmm1, xmm0, xmm0
        vpaddd  xmm0, xmm0, xmm1
        vpshufd xmm1, xmm0, 177
        vpaddd  xmm0, xmm1, xmm0
        vmovd   eax, xmm0
        ret

P.S.: GCC _mm512_reduce_add_epi32 与 clang 的 (相当于我的版本)的性能分析,使用来自 [=52= 的数据]:

内联到对结果执行某些操作的调用程序之后,它可以允许优化,例如使用 lea eax, [rax + rdx + 123] 或其他方式添加常量。

但除此之外,它似乎几乎总是比我在 Skylake-X 上实施结束时的 shuffle / vpadd / vmovd 差:

  • 总 uops:减少:4。我的:3
  • 端口:减少:2p0、p5(vpextrd 的一部分)、p0156(标量 add
  • 端口:我的:p5、p015(SKX 上vpadd)、p0(vmod

假设没有资源冲突,延迟在 4 个周期时相等:

  • shuffle 1 个周期 -> SIMD 添加 1 个周期 -> vmovd 2 个周期
  • vpextrd 3 个周期(与 2 个周期 vmovd 并行)-> 添加 1 个周期。