如何实现高效的_mm256_madd_epi8？

英特尔提供了一个名为_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元素是四个signed char产品的总和,而不是两对signed 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辅助函数¹ 返回signed__int8.没关系总数uint32_t而不是int32_t; 无论如何它只能添加四个8x8 => 16位数字才能溢出.

它看起来非常难看,并且不能高效运行,除非编译器使用SIMD做一些魔术,而不是编写为标量提取.

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

如果您知道其中一个输入始终为非负数,则可以使用pmaddubsw; 8-> 16位相当于pmaddwd.如果总和溢出,它会将饱和度写入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个字节要分散在两个 dword中,所以你必须改变一个结果_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);

作为vplldq1字节的替代,vpsllw8位__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矢量字节移位来避免shuffle控制向量.它仍然是一个通道内的shuffle(在每个通道的末尾以零的形式移动),因此它仍然具有单周期延迟.(我的第一个想法是vpshufb使用控制向量14,14, 12,12, 10,10, ...,然后vpalignr,我记得那个简单的旧pslldq版本有一个AVX2版本.同一条指令有两个名字.我喜欢_mm256_bslli_epi128因为bfor byte-shift将它区分为shuffle,不像内部 - 元素位移.我没有检查哪个编译器支持内核的128位或256位版本的名称.)

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

我没有看过AMD Ryzen执行端口冲突,但我很确定在任何CPU上都不会更糟(除了KNL Xeon Phi,其中AVX2操作小于dword的元素都非常慢).换档和车道内混音是相同的uop数和相同的延迟.

如果任何元素已知非负,则sign-extend = zero-extend

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

a_high和/或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

在主流英特尔上,随机播放吞吐量也限制为每时钟1个,因此如果您过度使用,可能会出现混乱的瓶颈.但至少它与乘法的端口不同.如果只有高字节被称为非负,替换vpsra/lw用vpshufb能有所帮助.未对齐的加载使得那些高字节是低字节可能更有帮助,设置vpandfor 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位总和可能比shift /和/ add更高的延迟,但确实将所有内容视为已签名.而且它只是一个uop,因此它对吞吐量有利.)