如何在 C (AVX2) 中向量化 int8 乘法

How to vectorise int8 multiplcation in C (AVX2)

如何使用 AVX2 向量化此 C 函数?

static void propogate_neuron(const short a, const int8_t *b, int *c) {

    for (int i = 0; i < 32; ++i){
        c[i] += a * b[i];
    }

}

您需要添加 restrict 限定符以标记 c 它不能与 b 别名。

问题是 int8_t 很可能 signed char 可以根据严格的别名规则与任何其他类型别名。因此编译器不能确定设置 c[i] 不会修改 b[i]。 强制编译器在每次迭代时获取数据。

存在const没有任何意义,因为它只是限制程序员通过指针修改数据b

替换原型后为:

void propogate_neuron(const short a, const int8_t *b, int * restrict c)

代码被向量化。参见 godbolt

GCC 已经通过检查重叠自动矢量化它。通过使用 int *restrict c 承诺没有重叠让 GCC 删除该检查,并让 clang 决定自动矢量化。

但是,clang 扩展到 32 位并使用 vpmulld,在 Haswell 及更高版本上为 2 微指令。 (虽然它在 Zen 上完全有效。)GCC 使用 vpmullwvpmulhw 来获得 16 位全乘法的低半部分和高半部分,并将它们混合在一起。 (Godbolt) 这是一个非常笨拙的策略,尤其是对于 -march=znver2,其中 vpmulld 是单个 uop。

GCC 确实只有四个单 uop 乘法指令,但实现它需要大量的 shuffle。我们可以做得更好:


因为我们只需要 8x16 => 32 位乘法,我们可以使用 vpmaddwd which is single-uop on Haswell/Skylake as well as Zen. https://uops.info/table.html

遗憾的是,我们无法利用加法部分,因为我们需要加法到一个完整的 32 位值。我们需要在每对 16 位元素的高半部分中使用零,以便将其用作每个 32 位元素中的 16x16 => 32 位乘法。

#include <immintrin.h>

void propogate_neuron_avx2(const short a, const int8_t *restrict b, int *restrict c) {
   __m256i va = _mm256_set1_epi32( (uint16_t)a );    // [..., 0, a, 0, a] 16-bit elements

   for (int i = 0 ; i < 32 ; i+=8) {
       __m256i vb = _mm256_cvtepi8_epi32( _mm_loadl_epi64((__m128i*)&b[i]) );
       __m256i prod = _mm256_madd_epi16(va, vb);
       __m256i sum = _mm256_add_epi32(prod, _mm256_loadu_si256((const __m256i*)&c[i]));
       _mm256_storeu_si256((__m256i*)&c[i], sum);
    }
}

Godbolt:

# clang13.0 -O3 -march=haswell
        movzx   eax, di
        vmovd   xmm0, eax                     # 0:a  16-bit halves
        vpbroadcastd    ymm0, xmm0            # repeated to every element

        vpmovsxbd       ymm1, qword ptr [rsi]  # xx:b 16-bit halves
        vpmaddwd        ymm1, ymm0, ymm1       # 0 + a*b in each 32-bit element
        vpaddd  ymm1, ymm1, ymmword ptr [rdx]
        vmovdqu ymmword ptr [rdx], ymm1

... repeated 3 more times, 8 elements per vector

        vpmovsxbd       ymm1, qword ptr [rsi + 8]
        vpmaddwd        ymm1, ymm0, ymm1
        vpaddd  ymm1, ymm1, ymmword ptr [rdx + 32]
        vmovdqu ymmword ptr [rdx + 32], ymm1

如果每个向量乘法保存一个 uop 会产生可衡量的性能差异,那么在源代码中手动向量化的麻烦可能是值得的。

GCC / clang 在自动向量化您的纯 C 代码时首先不这样做,这是一个错过的优化。

如果有人想举报此事,请在此处发表评论。否则我可能会解决它。 IDK 如果像这样的模式足够频繁以至于 GCC / LLVM 的优化器想要寻找这种模式。特别是 clang 已经做出了一个合理的选择,它只是次优的,因为 CPU 怪癖(32x32 => 32 位 SIMD 乘法在最近的英特尔微体系结构上比 2x 16x16 => 32 位具有水平添加的成本更高)。