如何在 C (AVX2) 中向量化 int8 数组与 int16 常量的乘法，并扩展到 int32 结果数组

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

}

（int8 x int8非展宽的相关问答，生成int8的向量。）

Answer 1

GCC 已经通过检查重叠来自动矢量化它。 通过使用

int *restrict c

保证不存在重叠，让 GCC 删除该检查，并让 clang 决定自动矢量化。

但是，clang 扩展到 32 位并使用

vpmulld

，在 Haswell 及更高版本上为 2 uops。 （尽管它在 Zen 上完全有效。）GCC 使用

vpmullw

和

vpmulhw

来获取 16 位全乘法的低半部分和高半部分，并将它们混在一起。 (Godbolt) 这是一个相当笨拙的策略，尤其是对于

-march=znver2

，其中

vpmulld

是单个 uop。

GCC 只有 4 个单微指令乘法指令，但是需要大量的 shuffle 来实现它。 我们可以做得更好：

---

由于我们只需要 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);
    }
}

神箭：

# 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 代码时不会首先执行此操作。

如果有人想举报此事，请在此处发表评论。 否则我可能会抽出时间来做这件事。 我不知道这样的模式是否足够频繁，以至于 GCC / LLVM 的优化器想要寻找这种模式。 特别是 clang 已经做出了一个合理的选择，但由于 CPU 的怪癖，该选择只是次优（在最新的 Intel 微架构上，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)

代码被矢量化。参见神箭