使用 AVX 内在函数计算具有允许的标量值 0、1 和 2 的向量的内积

Question

我正在做两列维度的内积以万计。这些值只能是 0、1 或 2。因此它们可以存储为字符。如果要在带有 avx 标志的 CPU 上对计算进行矢量化，我希望它快 32 倍左右。但问题是乘法会自动将字符转换为整数，即 4 个字节。因此，仅在速度上可以获得最大 8 倍。能达到32倍的速度吗？

顺便说一句，我正在使用带有 g++ 5.1 的 Linux（迄今为止的 Fedora 22）。

Answer 1

假设你有AVX2（不仅仅是AVX，它只用于浮点），那么你可以使用vpmaddubsw指令，其内在是：

__m256i _mm256_maddubs_epi16 (__m256i a, __m256i b)

这将执行 8 位 x 8 位乘法（有符号 x 无符号，但这对您的情况无关紧要），然后添加成对的相邻项以给出 16 位结果。 [1] 这有效地在一条指令中为您提供 32 x 8 x 8 位乘法。

如果您没有 AVX2，那么您可以使用 128 位 SSE 版本 (_mm_maddubs_epi16) 在一条指令中获得 16 x x 8 x 8 位乘法。

请注意，对 16 位项进行水平求和可能需要几条指令，但由于您的输入范围非常小，因此您只需要相对不频繁地执行此水平求和。一种可能的方法（用于 SSE）：

v = _mm_madd_epi16(v, _mm_set1_epi16(1));       // unpack/sum 16 -> 32
v = _mm_add_epi32(v, _mm_srli_si128(v, 8));     // shift and add 32 bit terms
v = _mm_add_epi32(v, _mm_srli_si128(v, 4));
sum = _mm_cvtsi128_si32(v);                     // extract sum as scalar

上面的 AVX2 实现留给读者作为练习。

Answer 2

看起来 AVX 指令集没有 8 位乘法，只有加法。 The Intel intrinsics guide 不包含任何以 _mm_mul* 开头的 8 位操作。 （编辑：实际上有一个 8 位乘法，但它有一个误导性的名称 - 请参阅 @PaulR 的回答）

但是，还有另一种方法。由于唯一允许的值是 0、1 和 2，并且您正在计算内积，因此您可以使用位运算而不是乘法。

在第一个向量A中，使用如下编码：

0 = 0b00000000 = 0x00
1 = 0b00010011 = 0x13
2 = 0b00001111 = 0x0F

在第二个向量B中，使用如下编码：

0 = 0b00000000 = 0x00
1 = 0b00011100 = 0x1C
2 = 0b00001111 = 0x0F

现在计算popcount(A & B)。 AND-ing 将导致相应的 8 位单元设置 0、1、2 或 4 位，popcount 会将它们加在一起。您可以每 5 位整数打包一个值，因此如果您可以使用 256 位整数，您可以获得 51 倍的吞吐量。

Answer 3

我想通过位操作来尝试是值得的。

假设所有数字都是 0 或 1。 然后您可以将两个向量打包成位数组。然后通过以下方式计算内积：

for (int i = 0; i < N; i += 256)
  res += popcount(A[i..i+255] & B[i..i+255]);

and 操作自然存在于 AVX/AVX2 中。最难的问题是如何快速计算 YMM 寄存器的 popcount。

现在假设给定 0、1 和 2。对于每个整数向量 A 组成两个位向量 A1 和 A2：

A1[i] = (A[i] >= 1);    
A2[i] = (A[i] >= 2);

现在我们可以注意到：

A[i] * B[i] = A1[i] * B1[i] + A1[i] * B2[i] + A2[i] * B1[i] + A2[i] * B2[i];

所以我们可以用下面的伪代码计算内积：

for (int i = 0; i < N; i += 256) {
  res += popcount(A1[i..i+255] & B1[i..i+255]);
  res += popcount(A2[i..i+255] & B1[i..i+255]);
  res += popcount(A1[i..i+255] & B2[i..i+255]);
  res += popcount(A2[i..i+255] & B2[i..i+255]);
}

这允许每次迭代处理 256 个元素，但每次迭代会慢 4 倍。有效地，每个操作有 64 个元素。由于 popcount 可能是计算中最慢的部分，我们可以说它需要 N/64 popcount_256 次操作来计算内积。

编辑：我决定为这个想法添加一个小例子：

A = {01212012210};  //input array A
B = {21221100120};  //input array B
A1 = {01111011110};  //A should be stored in two halves like this
A2 = {00101001100};
B1 = {11111100110};  //B is stored in similar two halves
B2 = {10110000010};
A1 & B1 = {01111000110}, popcount = 6;  //computing pairwise and-s + popcounts
A1 & B2 = {00110000010}, popcount = 3;
A2 & B1 = {00101000100}, popcount = 3;
A2 & B2 = {00100000000}, popcount = 1;
res = 6 + 3 + 3 + 1 = 13   //summing all the popcounts