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

以下是上述文件中左包装问题的说明:

谢谢

Answer 1

AVX2 + BMI2.请参阅我对AVX512的其他答案.(更新:pdep在64位版本中保存了一个.)

我们可以使用AVX2 vpermps(_mm256_permutevar8x32_ps)(或等价的整数vpermd)来进行交叉变量shuffle.

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

对于具有32位或更宽元素的整数向量:1)pdep.
或者2)使用pext然后将第一个PDEP常量从0x0101010101010101更改为0x0F0F0F0F0F0F0F0F以分散4个连续位的块.将乘法值乘以0xFFU更改为pshufb或vpermilps (未测试).无论哪种方式,使用带有VPERMD的shuffle掩码而不是VPERMPS.

对于64位整数或_mm256_movemask_ps(_mm256_castsi256_ps(compare_mask))元素,一切仍然正常 ; 比较掩码恰好总是具有相同的32位元素对,因此生成的shuffle将每个64位元素的两半放在正确的位置.(因此,您仍然使用VPERMPS或VPERMD,因为VPERMPD和VPERMQ仅适用于即时控制操作数.)

算法:

从包含3位索引的常量开始,每个位置都有自己的索引.即_mm256_movemask_epi8每个元素的宽度为3位. expanded_mask |= expanded_mask<<4;.

用于expanded_mask *= 0x11;将我们想要的索引提取到整数寄存器底部的连续序列中.例如,如果我们想要索引0和2,我们的控制掩码double应该是[ 7 6 5 4 3 2 1 0 ]. 0b111'110'101'...'010'001'000将抓住pext和pext该线向上与在选择的1个比特索引的基团.选定的组被打包到输出的低位,因此输出将是0b000'...'111'000'111.(即pext)

请参阅注释代码,了解如何从输入矢量掩码生成010输入000.

现在我们和压缩LUT在同一条船上:打开最多8个打包索引.

当你把所有的碎片放在一起时,总共有三个0b000'...'010'000/ [ ... 2 0 ]秒.我从我想要的东西后退,所以在这个方向上理解它可能最容易.(即从shuffle行开始,然后从那里向后工作.)

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

在Godbolt Compiler Explorer上看到这个和32bit版本.我使用了0b111000111s,所以用pextor 编译最佳pext.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);
}

这将编译为没有内存加载的代码,只有立即常量.(请参阅godbolt链接以及32位版本).

    # 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的数字,这是6 uops(不计算常数,或内联时消失的零延伸mov).在Intel Haswell上,它的延迟为16c(vmovq为1,每个pdep/imul/pext/vpmovzx/vpermps为3).没有指令级并行性.然而,在这不是循环携带依赖的一部分的循环中(就像我在Godbolt链接中包含的那个),瓶颈有希望只是吞吐量,同时保持多次迭代.

这可以管理每3个周期一个的吞吐量,在port1上为pdep/pext/imul设置瓶颈.当然,对于加载/存储和循环开销(包括compare,movmsk和popcnt),总uop吞吐量很容易成为问题.(例如,我的godbolt链接中的过滤器循环是14个带有clang的uop,pdep以便更容易阅读.如果我们幸运的话,它可以维持每4c一次迭代,跟上前端,但我认为clang失败了考虑#ifdef到对其输出的错误依赖性,因此它会在-m64函数延迟的3/5处出现瓶颈.)

gcc使用多个指令乘以0xFF,左移8和a -m32.这需要额外的-fno-unroll-loops指令,但最终结果是乘以2的延迟.(Haswell处理popcnt寄存器重命名阶段,零延迟.)

由于支持AVX2的所有硬件也支持BMI2,因此没有必要为没有BMI2的AVX2提供版本.

如果你需要在很长的循环中执行此操作,如果初始缓存未命中在足够的迭代中分摊,并且只需解压缩LUT条目的开销较低,则LUT可能是值得的.你仍然需要compress256,所以你可以弹出掩码并将其用作LUT索引,但是你保存了一个pdep/imul/pexp.

您可以解压LUT条目与我用同样的整数序列,但@ Froglegs的sub/ mov/ mov可能是更好的,当LUT条目记忆开始,并不需要去到整数寄存器摆在首位.(32位广播负载在Intel CPU上不需要ALU uop).然而,Haswell的变速是3 uops(但Skylake只有1).

Answer 2

我想出了这种方法,它使用压缩的LUT,即768(+1填充)字节,而不是8k.它需要广播单个标量值,然后在每个通道中移位不同的量,然后屏蔽到低3位,这提供0-7 LUT.

这是内在函数版本,以及构建LUT的代码.

//Generate Move mask via: _mm256_movemask_ps(_mm256_castsi256_ps(mask)); etc
__m256i MoveMaskToIndices(u32 moveMask) {
    u8 *adr = g_pack_left_table_u8x3 + moveMask * 3;
    __m256i indices = _mm256_set1_epi32(*reinterpret_cast<u32*>(adr));//lower 24 bits has our LUT

   // __m256i m = _mm256_sllv_epi32(indices, _mm256_setr_epi32(29, 26, 23, 20, 17, 14, 11, 8));

    //now shift it right to get 3 bits at bottom
    //__m256i shufmask = _mm256_srli_epi32(m, 29);

    //Simplified version suggested by wim
    //shift each lane so desired 3 bits are a bottom
    //There is leftover data in the lane, but _mm256_permutevar8x32_ps  only examines the first 3 bits so this is ok
    __m256i shufmask = _mm256_srlv_epi32 (indices, _mm256_setr_epi32(0, 3, 6, 9, 12, 15, 18, 21));
    return shufmask;
}

u32 get_nth_bits(int a) {
    u32 out = 0;
    int c = 0;
    for (int i = 0; i < 8; ++i) {
        auto set = (a >> i) & 1;
        if (set) {
            out |= (i << (c * 3));
            c++;
        }
    }
    return out;
}
u8 g_pack_left_table_u8x3[256 * 3 + 1];

void BuildPackMask() {
    for (int i = 0; i < 256; ++i) {
        *reinterpret_cast<u32*>(&g_pack_left_table_u8x3[i * 3]) = get_nth_bits(i);
    }
}

这是VS2015生成的程序集:

  lea ecx, DWORD PTR [rcx+rcx*2]
  lea rax, OFFSET FLAT:unsigned char * g_pack_left_table_u8x3 ; g_pack_left_table_u8x3
  vpbroadcastd ymm0, DWORD PTR [rcx+rax]
  vpsrlvd ymm0, ymm0, YMMWORD PTR __ymm@00000015000000120000000f0000000c00000009000000060000000300000000

Answer 3

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

既然你提到了对AVX512的可扩展性的担忧:不用担心,有一个AVX512F指令就是这样:

VCOMPRESSPS - 将稀疏打包的单精度浮点值存储到密集存储器中.(还有double和32或64bit整数元素(vpcompressq)的版本,但不是字节或字(16bit)).它就像BMI2 pdep/ pext,但是对于向量而不是整数reg中的位.

目标可以是向量寄存器或内存操作数,而源是向量和掩码寄存器.使用寄存器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):

 # 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

Answer 4

将为 @PeterCordes 的精彩答案添加更多信息：/sf/answers/2586612801/。

我用它从 C++ 标准中实现了 std::remove 的整数类型。一旦可以进行压缩，该算法就相对简单：加载寄存器、压缩、存储。首先，我将展示变化，然后展示基准。

我最终对所提出的解决方案提出了两个有意义的变体：

1. __m128i寄存器，任何元素类型，使用_mm_shuffle_epi8指令
2. __m256i寄存器，元素类型至少4字节，使用_mm256_permutevar8x32_epi32

当类型小于 256 位寄存器的 4 个字节时，我将它们分成两个 128 位寄存器，并分别压缩/存储每个寄存器。

链接到编译器资源管理器，您可以在其中看到完整的程序集（底部有一个using typeand width（每个包的元素中），您可以插入它以获得不同的变体）： https: //gcc.godbolt.org/z/yQFR2t

注意：我的代码采用 C++17 编写，并且使用自定义 simd 包装器，因此我不知道它的可读性如何。如果你想阅读我的代码 -> 大部分代码都位于顶部包含在 godbolt 上的链接后面。或者，所有代码都在github上。

@PeterCordes 的实现回答了这两种情况

注意：与掩码一起，我还使用 popcount 计算剩余元素的数量。也许有不需要的情况，但我还没有看到。

面膜用于_mm_shuffle_epi8

1. 将每个字节的索引写入半字节：0xfedcba9876543210
2. 将成对的索引放入 8 个短裤中__m128i
3. 使用将它们展开x << 4 | x & 0x0f0f
   

分散索引的示例。假设选择了第七个和第六个元素。这意味着相应的短路将是：0x00fe。之后我们<< 4就会|得到0x0ffe. 然后我们清除第二个f。

完整掩码代码：

// helper namespace
namespace _compress_mask {

// mmask - result of `_mm_movemask_epi8`, 
// `uint16_t` - there are at most 16 bits with values for __m128i. 
inline std::pair<__m128i, std::uint8_t> mask128(std::uint16_t mmask) {
    const std::uint64_t mmask_expanded = _pdep_u64(mmask, 0x1111111111111111) * 0xf;

    const std::uint8_t offset = 
        static_cast<std::uint8_t>(_mm_popcnt_u32(mmask));  // To compute how many elements were selected

    const std::uint64_t compressed_idxes = 
        _pext_u64(0xfedcba9876543210, mmask_expanded); // Do the @PeterCordes answer

    const __m128i as_lower_8byte = _mm_cvtsi64_si128(compressed_idxes); // 0...0|compressed_indexes
    const __m128i as_16bit = _mm_cvtepu8_epi16(as_lower_8byte);         // From bytes to shorts over the whole register
    const __m128i shift_by_4 = _mm_slli_epi16(as_16bit, 4);             // x << 4
    const __m128i combined = _mm_or_si128(shift_by_4, as_16bit);        // | x
    const __m128i filter = _mm_set1_epi16(0x0f0f);                      // 0x0f0f
    const __m128i res = _mm_and_si128(combined, filter);                // & 0x0f0f

    return {res, offset};
}

}  // namespace _compress_mask

template <typename T>
std::pair<__m128i, std::uint8_t> compress_mask_for_shuffle_epi8(std::uint32_t mmask) {
     auto res = _compress_mask::mask128(mmask);
     res.second /= sizeof(T);  // bit count to element count
     return res;
}

面膜用于_mm256_permutevar8x32_epi32

这几乎是一对一的@PeterCordes 解决方案 - 唯一的区别是_pdep_u64一点（他建议将此作为注释）。

我选择的面膜是0x5555'5555'5555'5555。这个想法是 - 我有 32 位的 mmask，8 个整数每个 4 位。我想要获得 64 位 => 我需要将 32 位中的每一位转换为 2 => 因此 0101b = 5。乘数也从 0xff 更改为 3，因为我将为每个整数获得 0x55，而不是 1。

完整掩码代码：

// helper namespace
namespace _compress_mask {

// mmask - result of _mm256_movemask_epi8
inline std::pair<__m256i, std::uint8_t> mask256_epi32(std::uint32_t mmask) {
    const std::uint64_t mmask_expanded = _pdep_u64(mmask, 0x5555'5555'5555'5555) * 3;

    const std::uint8_t offset = static_cast<std::uint8_t(_mm_popcnt_u32(mmask));  // To compute how many elements were selected

    const std::uint64_t compressed_idxes = _pext_u64(0x0706050403020100, mmask_expanded);  // Do the @PeterCordes answer

    // Every index was one byte => we need to make them into 4 bytes
    const __m128i as_lower_8byte = _mm_cvtsi64_si128(compressed_idxes);  // 0000|compressed indexes
    const __m256i expanded = _mm256_cvtepu8_epi32(as_lower_8byte);  // spread them out
    return {expanded, offset};
}

}  // namespace _compress_mask

template <typename T>
std::pair<__m256i, std::uint8_t> compress_mask_for_permutevar8x32(std::uint32_t mmask) {
    static_assert(sizeof(T) >= 4);  // You cannot permute shorts/chars with this.
    auto res = _compress_mask::mask256_epi32(mmask);
    res.second /= sizeof(T);  // bit count to element count
    return res;
}

基准测试

处理器：Intel Core i7 9700K（现代消费级 CPU，不支持 AVX-512）
编译器：clang，从版本 10 发行版附近的 trunk 构建
编译器选项：--std=c++17 --stdlib=libc++ -g -Werror -Wall -Wextra -Wpedantic -O3 -march=native -mllvm -align-all-functions=7
微基准测试库：google benchmark

控制代码对齐：
如果您不熟悉这个概念，请阅读此内容或观看此
基准测试二进制文件中的所有函数均与 128 字节边界对齐。每个基准测试函数都会重复 64 次，并在函数开头（进入循环之前）使用不同的 noop 幻灯片。我显示的主要数字是每次测量的最小值。我认为这是有效的，因为算法是内联的。我得到了截然不同的结果，这一事实也验证了我的观点。在答案的最底部，我展示了代码对齐的影响。
注意：基准测试代码。BENCH_DECL_ATTRIBUTES 只是非内联

Benchmark 从数组中删除一定比例的 0。我用 {0, 5, 20, 50, 80, 95, 100}% 的零来测试数组。
我测试了 3 个大小：40 字节（看看这是否适用于非常小的数组）、1000 字节和 10'000 字节。我按大小分组，因为 SIMD 取决于数据的大小而不是元素的数量。元素计数可以从元素大小得出（1000 字节是 1000 个字符，但 500 个短整型和 250 个整型）。由于非 simd 代码所需的时间主要取决于元素计数，因此字符的胜利应该更大。

绘图：x - 零的百分比，y - 以纳秒为单位的时间。padding : min 表示这是所有对齐方式中最小的。

40 个字节的数据，40 个字符

对于 40 字节，即使对于字符来说，这也没有意义 - 当使用 128 位寄存器而不是非 simd 代码时，我的实现速度会慢 8-10 倍。因此，例如，编译器应该小心执行此操作。

1000 字节数据，1000 个字符

显然，非 simd 版本主要由分支预测主导：当我们获得少量零时，我们获得较小的加速：对于没有 0 - 大约 3 倍，对于 5% 零 - 大约 5-6 倍加速。当分支预测器无法帮助非 simd 版本时 - 速度大约提高 27 倍。simd 代码的一个有趣的特性是它的性能对数据的依赖性要小得多。使用 128 与 256 寄存器几乎没有什么区别，因为大部分工作仍然分为 2 128 个寄存器。

1000 字节数据，500 个短数据

空头的结果类似，只是增益小得多 - 最多 2 倍。我不知道为什么对于非 simd 代码，shorts 比 chars 好得多：我预计 Shorts 会快两倍，因为只有 500 个 Shorts，但差异实际上高达 10 倍。

1000 字节数据，250 个整数

对于 1000 个仅 256 位版本来说是有意义的 - 20-30% 获胜，不包括任何 0 来删除任何内容（完美的分支预测，不删除非 simd 代码）。

10'000 字节数据，10'000 个字符

与 1000 个字符相同的数量级获胜：当分支预测器有帮助时快 2-6 倍，而当分支预测器无帮助时快 27 倍。

相同的图，只有 simd 版本：

在这里，我们可以看到使用 256 位寄存器并将它们分成 2 128 位寄存器大约有 10% 的优势：大约快 10%。它的大小从 88 条指令增加到 129 条指令，这并不是很多，因此根据您的用例可能有意义。对于基线 - 非 SIMD 版本有 79 条指令（据我所知 - 这些指令比 SIMD 指令要小）。

10'000 字节数据，5'000 条短信息

从 20% 到 9 次获胜，具体取决于数据分布。没有显示 256 和 128 位寄存器之间的比较 - 它几乎与字符相同的汇编，并且 256 位寄存器的相同胜利约为 10%。

10'000 字节数据，2'500 个整数

使用 256 位寄存器似乎很有意义，这个版本比 128 位寄存器快大约 2 倍。与非 simd 代码进行比较时，如果分支预测完美，则获胜 20%，而如果分支预测不完美，则获胜 3.5 - 4 倍。

结论：当您有足够的数据量（至少 1000 字节）时，对于没有 AVX-512 的现代处理器来说，这可能是非常值得的优化

附：

要删除的元素的百分比

一方面，过滤一半元素的情况并不常见。另一方面，类似的算法可以在排序 => 期间的分区中使用，实际上预计有大约 50% 的分支选择。

代码对齐影响

问题是：如果代码恰好对齐得不好（一般来说 - 对此几乎无能为力），那么它有多少价值。
我只显示 10'000 字节。
该图对于每个百分点有两条线表示最小值和最大值（这意味着 - 这不是最好/最差的代码对齐 - 它是给定百分比的最佳代码对齐）。

代码对齐影响 - 非 simd

字符：

从分支预测较差时的 15-20% 到分支预测有很大帮助时的 2-3 倍。（已知分支预测器会受到代码对齐的影响）。

短裤：

由于某种原因 - 0% 根本不受影响。可以通过std::remove首先进行线性搜索来找到要删除的第一个元素来解释。显然，短裤的线性搜索不受影响。除此之外 - 从 10% 到 1.6-1.8 倍价值

整数：

与短裤相同 - 没有 0 不受影响。一旦我们进入删除部分，它的价值就会从最佳情况对齐的 1.3 倍变为 5 倍。

代码对齐影响 - simd 版本

不显示 Shorts 和 ints 128，因为它与字符的汇编几乎相同

字符 - 128 位寄存器 大约慢1.2倍

字符 - 256 位寄存器 大约慢 1.1 - 1.24 倍

整数 - 256 位寄存器 慢 1.25 - 1.35 倍

我们可以看到，对于 simd 版本的算法，与非 simd 版本相比，代码对齐的影响明显较小。我怀疑这是因为实际上没有分支机构。

Answer 5

如果有人对这里感兴趣，那么这是SSE2的解决方案，它使用指令LUT代替数据LUT（又称跳转表）。对于AVX，这将需要256个案例。

每次您LeftPack_SSE2在下面进行调用时，它基本上使用三个指令：jmp，shufps，jmp。16种情况中有5种不需要修改向量。

static inline __m128 LeftPack_SSE2(__m128 val, int mask)  {
  switch(mask) {
  case  0:
  case  1: return val;
  case  2: return _mm_shuffle_ps(val,val,0x01);
  case  3: return val;
  case  4: return _mm_shuffle_ps(val,val,0x02);
  case  5: return _mm_shuffle_ps(val,val,0x08);
  case  6: return _mm_shuffle_ps(val,val,0x09);
  case  7: return val;
  case  8: return _mm_shuffle_ps(val,val,0x03);
  case  9: return _mm_shuffle_ps(val,val,0x0c);
  case 10: return _mm_shuffle_ps(val,val,0x0d);
  case 11: return _mm_shuffle_ps(val,val,0x34);
  case 12: return _mm_shuffle_ps(val,val,0x0e);
  case 13: return _mm_shuffle_ps(val,val,0x38);
  case 14: return _mm_shuffle_ps(val,val,0x39);
  case 15: return val;
  }
}

__m128 foo(__m128 val, __m128 maskv) {
  int mask = _mm_movemask_ps(maskv);
  return LeftPack_SSE2(val, mask);
}