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

Question

如何用 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];
    }

}

Answer 1

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

但是，clang 扩展到 32 位并使用 vpmulld，这是 Haswell 及更高版本上的 2 uop。 （尽管它在 Zen 上完全有效。） GCC 使用 vpmullw 和 vpmulhw 获得 16 位全乘法的低半部分和高半部分，并将它们混洗在一起。 (Godbolt) 这是一个相当笨拙的策略，尤其是-march=znver2，其中vpmulld 是单个微指令。

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

---

由于我们只需要 8x16 => 32 位乘法，我们可以改用 vpmaddwd，它是 Haswell/Skylake 和 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 会产生可衡量的性能差异，那么在源中手动进行向量化可能是值得的。

在自动矢量化纯 C 代码时，GCC / clang 一开始就没有这样做，这是一个错过的优化。

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

Answer 2

您需要添加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