c++ - AVX2是最有效的基于 mask 包装的方法吗？

Question

如果您有一个输入数组和一个输出数组，但是只想编写通过某个条件的那些元素，那么在AVX2中执行此操作的最有效方法是什么？

我在SSE中看到过这样做的地方：
（来自：https://deplinenoise.files.wordpress.com/2015/03/gdc2015_afredriksson_simd.pdf）

__m128i LeftPack_SSSE3(__m128 mask, __m128 val)
{
 // Move 4 sign bits of mask to 4-bit integer value.
 int mask = _mm_movemask_ps(mask);
 // Select shuffle control data
 __m128i shuf_ctrl = _mm_load_si128(&shufmasks[mask]);
 // Permute to move valid values to front of SIMD register
 __m128i packed = _mm_shuffle_epi8(_mm_castps_si128(val), shuf_ctrl);
 return packed;
}

这对于4宽的SSE似乎很好，因此仅需要16个条目的LUT，但是对于8宽的AVX，LUT变得非常大（256个条目，每个32字节，即8k）。

我很惊讶AVX似乎没有说明来简化此过程，例如带包装的遮罩存储。

我认为通过一些改组来计算设置在左边的符号位的数目，您可以生成必要的置换表，然后调用_mm256_permutevar8x32_ps。但这也是我认为的很多指示。

有谁知道使用AVX2进行此操作的任何技巧吗？或最有效的方法是什么？

这是上述文档中左包装问题的图示：



谢谢

Answer 1

AVX2 + BMI2。查看我对AVX512的其他答案。 （更新：在64位版本中保存了pdep。）

我们可以使用AVX2 vpermps (_mm256_permutevar8x32_ps)（或等效的整数vpermd）来进行车道交叉变量改组。

由于BMI2 pext (Parallel Bits Extract)为我们提供了所需操作的按位版本，因此我们可以动态生成掩码。

注意在AMD CPU上pdep / pext的运行速度非常慢，例如Ryzen上的6微秒/ 18个周期的延迟和吞吐量。此实现将在AMD上可怕地执行。对于AMD，如果掩码输入是矢量掩码（而不是已经计算出的位掩码），则最好使用使用pshufb或vpermilps LUT的128位向量，或者注释中讨论的一些AVX2可变移位建议。从记忆里）。在Zen2之前的AMD仍然只有128位向量执行单元，而256位交叉通道改组很慢。因此，在当前的AMD上，128位向量对此非常有吸引力。



对于具有32位或更大元素的整数向量：1）_mm256_movemask_ps(_mm256_castsi256_ps(compare_mask))。
或2）使用_mm256_movemask_epi8，然后将第一个PDEP常数从0x0101010101010101更改为0x0F0F0F0F0F0F0F0F0F，以分散4个连续位的块。将乘以0xFFU的乘数更改为expanded_mask |= expanded_mask<<4;或expanded_mask *= 0x11;（未测试）。无论哪种方式，都将改组掩码与VPERMD一起使用，而不要与VPERMPS一起使用。

对于64位整数或double元素，一切仍然有效。比较掩码恰好总是具有成对的32位元素对，因此所产生的混洗将每个64位元素的两半放在正确的位置。 （因此，您仍然使用VPERMPS或VPERMD，因为VPERMPD和VPERMQ仅可用于立即控制操作数。）

对于16位元素，您可能可以使用128位向量对其进行调整。



算法：

从一个压缩的3位索引常量开始，每个位置都有自己的索引。即[ 7 6 5 4 3 2 1 0 ]，其中每个元素为3位宽。 0b111'110'101'...'010'001'000。

使用pext将所需的索引提取到整数寄存器底部的连续序列中。例如如果我们想要索引0和2，则pext的控制掩码应为0b000'...'111'000'111。 pext将获取与选择器中的1位对齐的010和000索引组。选定的组被打包到输出的低位，因此输出将为0b000'...'010'000。 （即[ ... 2 0 ]）

有关如何从输入矢量掩码为0b111000111生成pext输入的信息，请参见注释的代码。

现在，我们与压缩LUT处于同一条船上：解压缩多达8个压缩索引。

当您将所有片段放在一起时，总共有三个pext / pdep。我从想要的工作中倒退了，因此也可能最容易朝那个方向理解它。 （即从洗牌线开始，然后从那里开始向后工作。）

如果我们使用每个字节一个索引而不是打包的3位组，则可以简化解压缩。由于我们有8个索引，因此只有64位代码才有可能。

请参见this and a 32bit-only version on the Godbolt Compiler Explorer。我使用了#ifdef，因此可以使用-m64或-m32进行最佳编译。 gcc浪费了一些指令，但是clang编写了非常不错的代码。

#include <stdint.h>
#include <immintrin.h>

// Uses 64bit pdep / pext to save a step in unpacking.
__m256 compress256(__m256 src, unsigned int mask /* from movmskps */)
{
  uint64_t expanded_mask = _pdep_u64(mask, 0x0101010101010101);  // unpack each bit to a byte
  expanded_mask *= 0xFF;    // mask |= mask<<1 | mask<<2 | ... | mask<<7;
  // ABC... -> AAAAAAAABBBBBBBBCCCCCCCC...: replicate each bit to fill its byte

  const uint64_t identity_indices = 0x0706050403020100;    // the identity shuffle for vpermps, packed to one index per byte
  uint64_t wanted_indices = _pext_u64(identity_indices, expanded_mask);

  __m128i bytevec = _mm_cvtsi64_si128(wanted_indices);
  __m256i shufmask = _mm256_cvtepu8_epi32(bytevec);

  return _mm256_permutevar8x32_ps(src, shufmask);
}

这可以编译为代码，而不会产生来自内存的任何负载，而只有立即常量。 （有关此版本和32位版本，请参见godbolt链接）。

    # clang 3.7.1 -std=gnu++14 -O3 -march=haswell
    mov     eax, edi                   # just to zero extend: goes away when inlining
    movabs  rcx, 72340172838076673     # The constants are hoisted after inlining into a loop
    pdep    rax, rax, rcx              # ABC       -> 0000000A0000000B....
    imul    rax, rax, 255              # 0000000A0000000B.. -> AAAAAAAABBBBBBBB..
    movabs  rcx, 506097522914230528
    pext    rax, rcx, rax
    vmovq   xmm1, rax
    vpmovzxbd       ymm1, xmm1         # 3c latency since this is lane-crossing
    vpermps ymm0, ymm1, ymm0
    ret

因此，根据 Agner Fog's numbers，这是6 oups（不计算常数，或内联时消失的零扩展mov）。在Intel Haswell上，延迟为16c（vmovq为1，每pdep / imul / pext / vpmovzx / vpermps为3）。没有指令级的并行性。但是，在一个循环中，它不是循环依赖项的一部分（就像我在Godbolt链接中包含的那样），瓶颈希望仅仅是吞吐量，可以一次进行多次迭代。

这可能可以管理每3个周期之一的吞吐量，这在pdep / pext / imul的端口1上是瓶颈。当然，由于加载/存储和循环开销（包括比较，movmsk和popcnt），总uop吞吐量很容易成为问题。 （例如，我的godbolt链接中的过滤器循环使用clang时为14微秒，使用 -fno-unroll-loops使其更易于阅读。如果幸运的话，它可能每4c维持一次迭代，与前端保持同步，但是如果我很幸运，认为clang无法解释 popcnt对输出的错误依赖，因此它将成为 compress256函数延迟的3/5的瓶颈。）

gcc使用左移8和 sub用多条指令乘以0xFF。这需要一条额外的 mov指令，但是最终结果是一个乘积，延迟为2。（Haswell在寄存器重命名阶段以零延迟处理 mov。）



由于所有支持AVX2的硬件也都支持BMI2，因此为没有BMI2的AVX2提供版本可能没有意义。

如果您需要在一个很长的循环中执行此操作，那么如果可以通过足够的迭代来摊销初始缓存缺失，并且只需拆开LUT条目的开销就较小，则LUT可能值得。您仍然需要 movmskps，因此可以弹出掩码并将其用作LUT索引，但是可以保存pdep / imul / pexp。

您可以使用与我使用的整数序列相同的整数序列来解压缩LUT条目，但是当LUT条目在内存中开始并且不需要先进入整数寄存器时，@ Froglegs的 set1() / vpsrlvd / vpand可能更好。地点。 （32位广播负载在Intel CPU上不需要ALU uop）。但是，在Haswell上，可变移位是3 oups（在Skylake上只有1 uops）。