AVX2 基于面具的最有效打包方式是什么？

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 做到这一点的任何技巧？或者什么是最有效的方法？

以下是上述文档中左包装问题的说明：

谢谢

原文由 Froglegs 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

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

我们可以使用 AVX2 vpermps ( _mm256_permutevar8x32_ps ) （或整数等价物， vpermd ）进行车道交叉变量洗牌。

我们可以动态生成掩码，因为 BMI2 pext （并行位提取） 为我们提供了所需操作的按位版本。

请注意 pdep / pext 在 Zen 3 之前的 AMD CPU 上 非常 慢，例如 Ryzen Zen 1 和 Zen 2 上的 6 uops / 18 周期延迟和吞吐量。此实现将在那些 AMD CPU 上执行得非常糟糕。对于 AMD，您最好使用 pshufb 或 vpermilps LUT 或评论中讨论的一些 AVX2 可变移位建议来使用 128 位向量。特别是如果您的掩码输入是矢量掩码（不是内存中已经打包的位掩码）。

Zen2 之前的 AMD 反正只有 128 位向量执行单元，而且 256 位车道交叉洗牌很慢。所以 128 位向量在 Zen 1 上非常有吸引力。但是 Zen 2 有 256 位加载/存储和执行单元。 （而且微编码 pext/pdep 仍然很慢。）

---

对于具有 32 位或更宽元素的整数向量： 1) _mm256_movemask_ps(_mm256_castsi256_ps(compare_mask)) 。

或者 2) 使用 _mm256_movemask_epi8 然后将第一个 PDEP 常量从 0x0101010101010101 更改为 0x0F0F0F0F0F0F0F0F 以分散 4 个连续位的块。将乘以 0xFFU 更改为 expanded_mask |= expanded_mask<<4; 或 expanded_mask *= 0x11; （未测试）。无论哪种方式，使用带有 VPERMD 而不是 VPERMPS 的 shuffle 掩码。

对于 64 位整数或 double 元素，一切仍然正常；比较掩码恰好总是具有相同的 32 位元素对，因此生成的 shuffle 将每个 64 位元素的两半放在正确的位置。 （所以您仍然使用 VPERMPS 或 VPERMD，因为 VPERMPD 和 VPERMQ 仅适用于直接控制操作数。）

对于 16 位元素，您也许可以使用 128 位向量来调整它。

对于 8 位元素，请参阅 Efficient sse shuffle mask generation for left-packing byte elements 以获得不同的技巧，将结果存储在多个可能重叠的块中。

---

算法：

从一个压缩的 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 ] ）

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

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

当你把所有的部分放在一起时，总共有三个 pext / pdep s。我从我想要的东西向后工作，所以在那个方向上也可能最容易理解它。 （即从洗牌线开始，然后从那里向后工作。）

如果我们使用每个字节一个索引而不是打包的 3 位组，我们可以简化解包。由于我们有 8 个索引，这仅适用于 64 位代码。

在 Godbolt Compiler Explorer 上查看此版本和仅 32 位版本。我使用 #ifdef s，因此它与 -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

（后来 clang 像 GCC 一样编译，用 mov/shl/sub 代替 imul，见下文。）

因此，根据 Agner Fog 的数字 和 https://uops.info/ ，这是 6 uops（不计算常量，或内联时消失的零扩展 mov）。在 Intel Haswell 上，延迟为 16c（vmovq 为 1，每个 pdep/imul/pext / vpmovzx / vpermps 为 3）。没有指令级并行性。但是，在一个不属于循环携带依赖项的循环中（就像我在 Godbolt 链接中包含的那个），瓶颈希望只是吞吐量，同时保持多次迭代。

这可以管理每 4 个周期一个的吞吐量，在端口 1 上为 pdep/pext/imul 加上循环中的 popcnt 造成瓶颈。当然，由于加载/存储和其他循环开销（包括比较和 movmsk），总 uop 吞吐量也很容易成为问题。

例如，我的 Godbolt 链接中的过滤器循环是 14 uop，带有 clang，带有 -fno-unroll-loops 以使其更易于阅读。如果幸运的话，它可能会维持每 4c 一次迭代，跟上前端。

clang 6 和更早的版本使用 popcnt 对其输出的错误依赖 创建了一个循环携带依赖项，因此它将在 compress256 函数的 ³⁄₅ 延迟上成为瓶颈。 clang 7.0 及更高版本使用 xor-zeroing 来打破错误依赖（而不是仅使用 popcnt edx,edx 或类似 GCC 的东西：/）。

gcc（以及后来的clang）使用多条指令乘以0xFF，使用左移8和 sub ，而不是 imul 乘以255。这总共需要3个微指令vs。 1 用于前端，但延迟仅为 2 个周期，低于 3。（Haswell 处理 mov 在寄存器重命名阶段，延迟为零。）最重要的是， imul 只能在端口 1 上运行，与 pdep/pext/popcnt 竞争，因此避免该瓶颈可能是件好事。

---

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

如果您需要在一个非常长的循环中执行此操作，那么 LUT 可能是值得的，如果初始缓存未命中在足够多的迭代中分摊，并且仅解包 LUT 条目的开销较低。您仍然需要 movmskps ，因此您可以弹出掩码并将其用作 LUT 索引，但保存 pdep/imul/pext。

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

原文由 Peter Cordes 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

请参阅我对没有 LUT 的 AVX2+BMI2 的其他答案。

既然您提到了对 AVX512 的可扩展性的担忧：别担心， 有一个 AVX512F 指令就是为了这个：

VCOMPRESSPS — 将稀疏打包的单精度浮点值存储到密集内存 中。 （还有双精度、32 位或 64 位整数元素的版本（ vpcompressq ），但不是字节或字（16 位））。就像 BMI2 pdep / pext ，但是对于向量元素而不是整数寄存器中的位。

目标可以是向量寄存器或内存操作数，而源可以是向量和掩码寄存器。使用寄存器 dest，它可以合并或归零高位。使用内存 dest，“仅将连续向量写入目标内存位置”。

要确定将指针前进到下一个向量的距离，请弹出掩码。

假设您想从数组中过滤掉除值 >= 0 之外的所有内容：

 #include <stdint.h>
#include <immintrin.h>
size_t filter_non_negative(float *__restrict__ dst, const float *__restrict__ src, size_t len) {
    const float *endp = src+len;
    float *dst_start = dst;
    do {
        __m512      sv  = _mm512_loadu_ps(src);
        __mmask16 keep = _mm512_cmp_ps_mask(sv, _mm512_setzero_ps(), _CMP_GE_OQ);  // true for src >= 0.0, false for unordered and src < 0.0
        _mm512_mask_compressstoreu_ps(dst, keep, sv);   // clang is missing this intrinsic, which can't be emulated with a separate store

        src += 16;
        dst += _mm_popcnt_u64(keep);   // popcnt_u64 instead of u32 helps gcc avoid a wasted movsx, but is potentially slower on some CPUs
    } while (src < endp);
    return dst - dst_start;
}

这将（使用 gcc4.9 或更高版本）编译为（ Godbolt Compiler Explorer+%7B%0A++++const+float+*endp+%3D+src%2Blen%3B%0A++++float+*dst_start+%3D+dst%3B%0A++++do+%7B%0A++++++++m512++++++sv++%3D+_mm512_loadu_ps(src)%3B%0A++++++++mmask16+keep+%3D+_mm512_cmp_ps_mask(sv,+_mm512_setzero_ps(),+_CMP_GE_OQ)%3B++//+true+for+src+%3E%3D+0.0,+false+for+unordered+and+src+%3C+0.0%0A++++++++_mm512_mask_compressstoreu_ps(dst,+keep,+sv)%3B+++//+clang+is+missing+this+intrinsic,+which+can!’t+be+emulated+with+a+separate+store%0A%0A++++++++src+%2B%3D+16%3B%0A++++++++dst+%2B%3D+_mm_popcnt_u64(keep)%3B+++//+popcnt_u64+instead+of+u32+helps+gcc+avoid+a+wasted+movsx,+but+is+potentially+slower+on+some+CPUs%0A++++%7D+while+(src+%3C+endp)%3B%0A++++return+dst+-+dst_start%3B%0A%7D%0A’)),filterAsm:(commentOnly:!t,directives:!t,intel:!t,labels:!t),version:3) ）：

  # Output from gcc6.1, with -O3 -march=haswell -mavx512f.  Same with other gcc versions
    lea     rcx, [rsi+rdx*4]             # endp
    mov     rax, rdi
    vpxord  zmm1, zmm1, zmm1             # vpxor  xmm1, xmm1,xmm1 would save a byte, using VEX instead of EVEX
.L2:
    vmovups zmm0, ZMMWORD PTR [rsi]
    add     rsi, 64
    vcmpps  k1, zmm0, zmm1, 29           # AVX512 compares have mask regs as a destination
    kmovw   edx, k1                      # There are some insns to add/or/and mask regs, but not popcnt
    movzx   edx, dx                      # gcc is dumb and doesn't know that kmovw already zero-extends to fill the destination.
    vcompressps     ZMMWORD PTR [rax]{k1}, zmm0
    popcnt  rdx, rdx
    ## movsx   rdx, edx         # with _popcnt_u32, gcc is dumb.  No casting can get gcc to do anything but sign-extend.  You'd expect (unsigned) would mov to zero-extend, but no.
    lea     rax, [rax+rdx*4]             # dst += ...
    cmp     rcx, rsi
    ja      .L2

    sub     rax, rdi
    sar     rax, 2                       # address math -> element count
    ret

---

性能：256 位向量在 Skylake-X / Cascade Lake 上可能更快

理论上，加载位图并将一个数组过滤到另一个数组的循环应该在 SKX / CSLX 上以每 3 个时钟 1 个向量运行，无论向量宽度如何，在端口 5 上成为瓶颈。（ kmovb/w/d/q k1, eax 在 p5 上运行，和 vcompressps 根据 IACA 和 http://uops.info/ 的测试，内存是 2p5 + 一个存储）。

@ZachB 在评论中报告说，在 实际 CSLX 硬件上，使用 ZMM _mm512_mask_compressstoreu_ps 的循环比 _mm256_mask_compressstoreu_ps 稍慢。 （我不确定这是否是允许 256 位版本退出“512 位矢量模式”并提高时钟频率的微基准，或者是否有周围的 512 位代码。）

我怀疑未对齐的商店正在损害 512 位版本。 vcompressps 可能有效地进行了屏蔽的 256 位或 512 位向量存储，如果它跨越了缓存线边界，那么它必须做额外的工作。由于输出指针通常不是 16 个元素的倍数，因此全行 512 位存储几乎总是未对齐。

由于某种原因，未对齐的 512 位存储可能比缓存行拆分的 256 位存储更糟糕，而且发生得更频繁；我们已经知道，其他事物的 512 位矢量化似乎对对齐更加敏感。这可能只是因为每次都发生拆分加载缓冲区时用完，或者处理缓存行拆分的回退机制对于 512 位向量的效率较低。

将 vcompressps 基准测试到寄存器中会很有趣，具有单独的全向量重叠存储。这可能是相同的微指令，但是当它是一个单独的指令时，存储可以微融合。如果蒙面商店与重叠商店之间存在一些差异，这将揭示它。

---

下面评论中讨论的另一个想法是使用 vpermt2ps 为对齐的商店构建完整向量。这 将很难做到无 分支，并且当我们填充向量时分支可能会错误预测，除非位掩码具有非常规则的模式，或者全 0 和全 1 的大量运行。

一个无分支的实现是可能的，在构建的向量中循环携带 4 或 6 个循环的依赖链，使用 vpermt2ps 和混合或在它“满”时替换它的东西。每次迭代都使用对齐的向量存储，但仅在向量已满时移动输出指针。

这可能比当前 Intel CPU 上未对齐存储的 vcompressps 慢。

原文由 Peter Cordes 发布，翻译遵循 CC BY-SA 4.0 许可协议