如何实现一个高效的_mm256_madd_epi8?

How to implement an efficient _mm256_madd_epi8?

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

__m256i _mm256_madd_epi16 (__m256i a, __m256i b)

Multiply packed signed 16-bit integers in a and b, producing intermediate signed 32-bit integers. Horizontally add adjacent pairs of intermediate 32-bit integers, and pack the results in dst.

现在我有两个 __m256i 变量,每个变量都有 32 个 8 位整数。

我想实现与 _mm256_madd_epi16 相同的功能,但结果 __m256i 中的每个 int32_t 元素都是 signed char 的四个乘积的 之和而不是两对带符号的 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);

注意乘法结果是符号-在加法之前扩展到int,并且_mm256_extract_epi8辅助函数1 returns signed __int8。不要介意总数是 uint32_t 而不是 int32_t;无论如何它都不会溢出,只有四个 8x8 => 16 位数字要添加。

它看起来非常丑陋,并且无法高效运行,除非编译器使用 SIMD 执行某些操作而不是像写入标量提取那样进行编译。


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

如果已知您的输入之一始终为非负数,则可以使用 pmaddubswpmaddwd 的 8->16 位等价物。如果总和溢出,它确实会将符号饱和到 16 位,这是可能的,所以如果这对你的情况来说是个问题,你可能需要避免它。

但除此之外,您可以pmaddubsw然后手动将 16 位元素符号扩展为 32 位并添加它们。或使用 pmaddwd_mm256_set1_epi16(1) 来 hsum 正确处理符号的元素对。


显而易见的解决方案是将您的输入字节解压缩为具有零或符号扩展名的 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); 是每个字内从低到高移动的更多 "obvious" 方式,并且如果随机播放周围的代码瓶颈可能会更好。但通常情况更糟,因为 this 代码在 shift + vec-int 乘法上存在严重瓶颈。

在 KNL 上,您可以使用 AVX512 vprold z,z,i(Agner Fog 不显示 AVX512 的时序 vpslld z,z,i),因为将什么移入或混入低字节并不重要每个字;这只是算术右移的设置。

执行端口瓶颈:

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

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

但是我们可以通过使用vpslldq 进行向量字节移位 来避免混洗控制向量。它仍然是通道内洗牌(在每个通道的末端移入零),因此它仍然具有单周期延迟。 (我的第一个想法是 vpshufb14,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 在小于 a 的元素上运行dword 都超级慢)。轮班和车道内洗牌是相同数量的 uops 和相同的延迟。

如果已知任何元素是非负的,则符号扩展 = 零扩展

零扩展比手动符号扩展成本更低,并且避免了端口瓶颈。 a_low and/or b_low 可以用 _mm256_and_si256(a, _mm256_set1_epi16(0x00ff)).

创建

a_high and/or b_high 可以用 shuffle 而不是 shift 创建。 (pshufb 当洗牌控制向量设置了高位时将元素归零)。

 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 瓶颈。但至少它与 multiply 不是同一个端口。如果只有高字节已知为非负数,则将 vpsra/lw 替换为 vpshufb 可能会有所帮助。未对齐的加载,因此那些高字节是低字节可能更有帮助,为 vpand 设置 a_low and/or 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,因此对吞吐量有好处.)