如何实现四个 i8 元素组的高效_mm256_madd_epi8点积

How to implement an efficient _mm256_madd_epi8 dot-products of groups of four i8 elements?

本文关键字：mm256 高效 madd 点积 epi8 i8 实现 何实现 四个 元素      更新时间：2023-10-16

英特尔提供了一个名为_mm256_madd_epi16的C风格函数，它基本上

__m256i _mm256_madd_epi16 (__m256i a， __m256i b)

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

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

我想实现与_mm256_madd_epi16相同的功能，但是结果__m256i中的每个int32_t元素都是有符号字符的四个乘积的总和，而不是两对有符号int16_t。每个 32 位块中四个int8_t元素的点积。

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

alignas(32) uint32_t res[8] = {0};
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返回有符号的__int8。 没关系，总数是uint32_t而不是int32_t;无论如何，它不能溢出，只有四个 8x8 => 16 位数字要添加。

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

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

pmaddubsw：

如果至少有一个输入是非负的(因此可以被视为无符号)，则可用

如果你的输入之一总是非负的，你可以使用它作为无符号输入来pmaddubsw;相当于pmaddwd的8->16位。 它添加了成对的u8*i8 -> i16产品，符号饱和度达到 16 位。 但是饱和是不可能的，一个输入最多是 127 而不是 255。 (127*-128 = -0x3f80，所以两倍仍然适合 i16。

pmaddubsw后，使用pmaddwd对_mm256_set1_epi16(1)对元素对进行正确处理。 (这通常比手动将 16 位元素签名扩展到 32 以添加它们更有效。

__m256i sum16 = _mm256_maddubs_epi16(a, b);   // pmaddubsw
__m256i sum32 = _mm256_madd_epi16(sum16, _mm256_set1_epi16(1)); // pmaddwd

( 对于水平 16=>32 位 4 字节元素中的对和，pmaddwd在某些 CPU 上的延迟高于 shift/和/add，但确实将两个输入都视为有符号以符号扩展到 32 位。 而且它只是一个 uop，因此它有利于吞吐量，特别是如果周围的代码在相同的执行端口上没有瓶颈。

<小时 />

一般情况(两个输入都可能是负数)

最近关于AVX-512BW仿真的AVX-512VNNI指令_mm512_dpbusd_epi32答案提出了一个很好的技巧，将一个输入拆分为MSB，低7位，以便可以使用vpmaddubsw(_mm256_maddubs_epi16)而不会溢出。 我们可以借用这个技巧并在求和时否定，因为 MSB 的位值是-2^7的，而不是vpmaddubsw的无符号输入将其视为的2^7。

// Untested.  __m128i version would need SSSE3
__m256i dotprod_i8_to_i32(__m256i v1, __m256i v2)
{
const __m256i highest_bit = _mm256_set1_epi8(0x80);
__m256i msb = _mm256_maddubs_epi16(_mm256_and_si256(v1, highest_bit), v2);     // 0 or 2^7
__m256i low7 = _mm256_maddubs_epi16(_mm256_andnot_si256(highest_bit, v1), v2);
low7 = _mm256_madd_epi16(low7, _mm256_set1_epi16(1));  // hsum i16 pairs to i32
msb  = _mm256_madd_epi16(msb,  _mm256_set1_epi16(1));
return _mm256_sub_epi32(low7, msb);  // place value of the MSB was negative
// equivalent to the below, but that needs an extra constant
//    msb = _mm256_madd_epi16(msb,  _mm256_set1_epi16(-1));   // the place-value was actually - 2^7
//    return _mm256_add_epi32(low7, msb);
// also equivalent to vpmaddwd with -1 for both parts
// return sub(msb, low7)
// which is cheaper because set1(-1) is just vpcmpeqd not a load.
}

这避免了有符号饱和：一侧的最大乘数为 128(MSB 被设置并被视为无符号)。128 * -128= -16384，两倍即 -32768 = -0x8000 = 位模式0x8000。 或128 * 127 * 2= 0x7f00 作为最高的阳性结果。

对于下面的版本，这是 7 uops(乘法单位为 4)，而以下版本为 9 uops(4 个移位 + 2 个乘法)。

AVX-512VNNI_mm256_dpbusd_epi32(或 512)或AVX_VNNI_mm256_dpbusd_avx_epi32(VPDPBUSD) 类似于vpmaddubsw(u8*i8个产品)，但添加到现有总和中，并在单个指令中将 4 个产品相加在一个字节内。 (i32 += four u8 * i8)。_mm256_sub_epi32(low7_prods, msb_prods)，同样的拆分技巧有效，但我们可以跳过madd_epi16(vpmaddwd)i16到i32的水平求和步骤。

(其他 VNNI 指令包括vpdpbusds(与vpdpbusd相同，但具有符号饱和而不是包装)。 无论哪种方式，饱和度都是 i32，而不是像vpmaddubsw那样的 i16 ，因此只有在累加器输入不为零时才饱和。 如果一个输入是非负的，因此可以被视为无符号，这将在一个指令中完成整个工作而不会拆分。vpdpwssd[s]，有或没有饱和的有符号的MAC，就像vpmaddwd但带有累加器操作数。

// Ice Lake (AVX-512 version only) or Alder Lake (AVX_VNNI), or Zen 4
__m256i dotprod_i8_to_i32_vnni(__m256i v1, __m256i v2)
{
const __m256i highest_bit = _mm256_set1_epi8(0x80);
__m256i msb = _mm256_and_si256(v1, highest_bit);
__m256i low7 = _mm256_andnot_si256(highest_bit, v1);
// or just _mm256_dpbusd_epi32 for the EVEX version
msb = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), msb, v2);     // 0 or 2^7
low7 = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), low7, v2);
return _mm256_sub_epi32(low7, msb);  // place value of the MSB was negative
}

没有AVX-512VNNI的AVX-512可以使用AVX2版本不变，或者扩大到512。 或者可以通过将其转换为掩码(vptestmb)并将输入的某些字节(零屏蔽vpmovdqu8)归零来应用符号位，将4字节块的水平和转换为32位元素(使用恒等式洗牌控制vdbpsadbw对零)。 但是不，这不会在添加 8 位输入之前对其进行签名扩展，因为它是无符号差异。 也许首先将范围偏移到无符号(例如，带有0x80的零掩蔽异或)，然后添加4*128？ 无论如何，msb = _mm256_slli_epi32(dword_hsums_of_input_b, 7)上面的代码使用其msb变量的方式相同。 如果这甚至有效，IDK 如果它节省了 uops。 欢迎反馈，或发布AVX-512BW答案。

另一种方式：解压缩并签名扩展到 16 位

显而易见的解决方案是将输入字节解压缩为具有零或符号扩展名的 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);

作为vplldqx 1 字节的替代方案，__m256i a_low = _mm256_slli_epi16(a, 8);vpsllw8 位是在每个单词中从低到高的更"明显"的方法，如果周围的代码瓶颈在随机播放时可能会更好。 但通常情况下情况更糟，因为这段代码严重阻碍了 shift + vec-int 乘法。

在KNL上，你可以使用AVX512vprold z,z,i(Agner Fog没有显示AVX512vpslld z,z,i的时间)，因为你在每个单词的低字节中移位或随机播放什么并不重要;这只是算术右移的设置。

执行端口瓶颈：

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

我们可以使用随机播放(端口 5)而不是左移作为算术右移的设置。 这提高了吞吐量，甚至通过减少资源冲突来减少延迟。

但是我们可以通过使用vpslldq进行向量字节偏移来避免随机控制向量。 它仍然是车道内洗牌(在每个车道末端以零换档)，因此它仍然具有单周期延迟。 (我的第一个想法是用像14,14, 12,12, 10,10, ...这样的控制向量vpshufb，然后vpalignr，然后我想起了简单的旧pslldq有一个AVX2版本。同一指令有两个名称。 我喜欢_mm256_bslli_epi128因为字节移位b将其区分为洗牌，与元素内位移不同。 我没有检查哪个编译器支持 128 位或 256 位版本的内部函数的名称。

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

我没有看过AMD Ryzen执行端口冲突，但我很确定这在任何CPU上都不会更糟(除了KNL Xeon Phi，其中AVX2在小于dword的元素上的操作都非常慢)。 班次和车道内随机播放是相同数量的 uop 和相同的延迟。

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

(或者更好的是，使用第一部分中所示的pmaddubsw。

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

a_high和/或b_high可以通过随机播放而不是移位来创建。 (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 个，因此如果您过火，您可能会在随机播放时遇到瓶颈。 但至少它与乘法不是同一个端口。 如果只有高字节是已知的非负字节，那么用vpshufb替换vpsra/lw可能会有所帮助。 未对齐的加载，因此这些高字节是低字节可能会更有帮助，为a_low和/或b_low设置vpand。

相关文章：