如何实现高效的_mm256_madd_epi8？

Question

Intel 提供了一个名为 _mm256_madd_epi16 的 C 风格函数，它基本上是

__m256i _mm256_madd_epi16 (__m256i a, __m256i b)

将 a 和 b 中的压缩有符号 16 位整数相乘，生成中间有符号 32 位整数。将相邻的中间 32 位整数对水平相加，并将结果打包到 dst 中。

现在我有两个 __m256i 变量，每个变量都有 32 个 8 位 int。

我想实现与 _mm256_madd_epi16 相同的功能，但结果 __m256i 中的每个 int32_t 元素是 四个带符号字符的乘积之和，而不是两对带符号的 int16_t。

我可以在标量循环中做到这一点：

  alignas(32) uint32_t res[8];
  for (int i = 0; i < 32; ++i)
      res[i / 4] += _mm256_extract_epi8(a, i) * _mm256_extract_epi8(b, i);
  return _mm256_load_si256((__m256i*)res);

请注意，相乘结果是 sign-在相加之前扩展为 int，并且 _mm256_extract_epi8 辅助函数¹returns signed __int8。没关系，总数是uint32_t 而不是int32_t；它无论如何都不会溢出，只需要添加四个 8x8 => 16 位数字。

它看起来很丑陋，并且运行效率不高，除非编译器使用 SIMD 做一些魔法，而不是按照写入标量提取的方式进行编译。

---

脚注 1：_mm256_extract_epi8 不是内在的。 vpextrb 仅适用于 256 位向量的低通道，并且此辅助函数可能允许索引不是编译时常量。

Answer 1

如果已知您的输入之一始终为非负数，您可以使用pmaddubsw； pmaddwd 的 8->16 位等价物。如果总和溢出，它确实有符号饱和到 16 位，这是可能的，所以如果这对您的情况有问题，您可能需要避免它。

但除此之外，您可以pmaddubsw 然后手动将 16 位元素符号扩展为 32 并添加它们。或者使用 pmaddwd 对 _mm256_set1_epi16(1) 来对正确处理符号的元素对进行求和。

---

显而易见的解决方案是将输入字节解压缩为 16 位元素，并带有零或符号扩展名。然后你可以使用pmaddwd 两次，并添加结果。

如果您的输入来自内存，使用vpmovsxbw 加载它们可能有意义。例如

__m256i a = _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)&arr1[i]);
__m256i b = _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)&arr2[i]);

但现在您有 4 个字节，您希望将其分布在 两个 个双字中，因此您必须对一个 _mm256_madd_epi16(a,b) 的结果进行洗牌。您也许可以使用vphaddd 洗牌并将两个 256 位产品向量添加到您想要的一个 256 位结果向量中，但这是很多洗牌。

因此，我认为我们希望从每个 256 位输入向量生成两个 256 位向量：一个将每个字中的高字节符号扩展为 16，另一个将低字节符号扩展。我们可以通过 3 个班次（每个输入）来做到这一点

 __m256i a = _mm256_loadu_si256(const  __m256i*)&arr1[i]);
 __m256i b = _mm256_loadu_si256(const  __m256i*)&arr2[i]);

 __m256i a_high = _mm256_srai_epi16(a, 8);     // arithmetic right shift sign extends
     // some compilers may only know the less-descriptive _mm256_slli_si256 name for vpslldq
 __m256i a_low =  _mm256_bslli_epi128(a, 1);   // left 1 byte = low to high in each 16-bit element
         a_low =  _mm256_srai_epi16(a_low, 8); // arithmetic right shift sign extends

    // then same for b_low / b_high

 __m256i prod_hi = _mm256_madd_epi16(a_high, b_high);
 __m256i prod_lo = _mm256_madd_epi16(a_low, b_low);

 __m256i quadsum = _m256_add_epi32(prod_lo, prod_hi);

作为vplldq 1 字节的替代方法，vpsllw 8 位 __m256i a_low = _mm256_slli_epi16(a, 8); 是在每个单词中从低到高移动的更“明显”的方式，如果周围的代码瓶颈在洗牌时可能会更好.但通常情况会更糟，因为 此 代码在 shift + vec-int 乘法上严重瓶颈。

在 KNL 上，您可以使用 AVX512 vprold z,z,i（Agner Fog 没有显示 AVX512 vpslld z,z,i 的时间），因为无论您将什么移动或洗牌到每个字的低字节都无关紧要；这只是为算术右移设置。

执行端口瓶颈：

Haswell 仅在端口 0 上运行向量移位和向量整数乘法，因此这严重成为瓶颈。 （Skylake 更好：p0/p1）。 http://agner.org/optimize/.

我们可以使用 shuffle（端口 5）代替左移作为算术右移的设置。这提高了吞吐量，甚至通过减少资源冲突来减少延迟。

但我们可以通过使用vpslldq 进行矢量字节移位来避免随机播放控制矢量。它仍然是一个in-lane shuffle（在每个lane的末端移入零），所以它仍然具有单周期延迟。 （我的第一个想法是vpshufb 带有一个控制向量，如14,14, 12,12, 10,10, ...，然后是vpalignr，然后我记得简单的旧pslldq 有一个AVX2 版本。同一指令有两个名称。 我喜欢_mm256_bslli_epi128，因为用于字节​​移位的b 将其区分为随机播放，这与元素内移位不同。我没有检查哪个编译器支持 128 位或 256 位版本的内在函数的名称。）

这也有助于 AMD Ryzen。向量移位仅在一个执行单元 (P2) 上运行，但随机播放可以在 P1 或 P2 上运行。

我没有看过 AMD Ryzen 执行端口冲突，但我很确定这在任何 CPU 上都不会更糟（KNL Xeon Phi 除外，其中 AVX2 对小于 dword 的元素的操作都非常慢） .移位和通道内随机播放具有相同的微指令数和相同的延迟。

如果任何元素已知为非负数，则符号扩展 = 零扩展

零扩展比手动符号扩展更便宜，并且避免了端口瓶颈。 a_low 和/或b_low 可以使用_mm256_and_si256(a, _mm256_set1_epi16(0x00ff)) 创建。

a_high 和/或b_high 可以使用随机播放而不是移位创建。 （pshufb 在 shuffle-control 向量的高位设置时将元素归零）。

 const _mm256i pshufb_emulate_srl8 = _mm256_set_epi8(
               0x80,15, 0x80,13, 0x80,11, ...,
               0x80,15, 0x80,13, 0x80,11, ...);

 __m256i a_high = _mm256_shuffle_epi8(a, pshufb_emulate_srl8);  // zero-extend

在主流 Intel 上，Shuffle 吞吐量也被限制为每个时钟 1 个，因此如果过度使用，可能会成为 shuffle 的瓶颈。但至少它与乘法的端口不同。如果只有高字节已知为非负数，则将 vpsra/lw 替换为 vpshufb 可能会有所帮助。未对齐的加载，因此那些高字节是低字节可能会更有帮助，为vpand 设置a_low 和/或b_low。

---

pmaddubsw：我认为如果至少有一个输入是非负的（因此可以视为无符号），这是可用的

它将一个输入视为有符号，另一个视为无符号，并执行 i8 x u8 => i16，然后添加水平对以生成 16 位整数（有符号饱和，因为总和可能会溢出。这也可能排除它适合您的用例）。

但可能只是使用它，然后使用 pmaddwd 与常量 1 添加水平对：

__m256i sum16 = _mm256_maddubs_epi16(a, b);
__m256i sum32 = _mm256_madd_epi16(sum16, _mm256_set1(1));

（pmaddwd 对于水平 16=>32 位求和可能比移位 / 和 / 加法的延迟更高，但确实将所有内容都视为有符号。而且它只是一个 uop，因此对吞吐量有好处。）