使用 SIMD 指令执行任意 128/256/512 位置换的最快方法是什么？答案

【问题标题】：What's the fastest way to perform an arbitrary 128/256/512 bit permutation using SIMD instructions?使用 SIMD 指令执行任意 128/256/512 位置换的最快方法是什么？
【发布时间】：2019-06-21 20:28:30
【问题描述】：

我想在宽度为 128、256 或 512 位的 CPU 寄存器（xmm、ymm 或 zmm）上执行单个位、位对和半字节（4 位）的任意排列；这应该尽可能快。为此，我正在研究 SIMD 指令。有谁知道这样做的方法/实现它的库？我在 Windows 上使用 MSVC，在 Linux 上使用 GCC，宿主语言是 C 或 C++。谢谢！

我得到了一个任意排列，需要对大量位向量/位向量对/半字节进行洗牌。我知道如何对 64 位值中的位执行此操作，例如using a Benes network.

或者在更宽的 SIMD 寄存器上混洗 8 位和更大的块，例如将 Agner Fog 的 GPLed VectorClass 库 (https://www.agner.org/optimize/vectorclass.pdf) 用于模板元编程函数，该函数在将 shuffle 作为模板参数的情况下，从 AVX2 通道内字节随机播放和/或更大元素的通道交叉随机播放构建随机播放。

不过，对排列进行更细粒度的细分（分为 1、2 或 4 位块）似乎很难在宽向量中实现。

我能够对排列进行预处理，例如提取位掩码，根据需要计算索引，例如对于 Benes 网络或其他任何东西——也很乐意用另一种高级语言来做这件事，所以假设排列是以最方便解决问题的任何格式给出的；包括小型查找表。

我希望代码比执行类似的操作要快得多

// actually 1 bit per element, not byte.  I want a 256-bit bit-shuffle
const uint8_t in[256] = get_some_vector(); // not a compile-time constant
const uint8_t perm[256] = ...;             // compile-time constant
uint8_t out[256];
for (size_t i = 0; i < 256; i ++)
    out[i] = in[perm[i]];

正如我所说，我有一个

编辑：澄清一下，排列是一个编译时常数（但不仅仅是一个特定的，一旦给出排列，我将编译程序）。

【问题讨论】：

为什么需要 Benes 网络来处理大型元素？ AVX2 有_mm256_permutevar8x32_epi32 (vpermd)，它为任意车道交叉置换采用控制向量。（对于 64 位索引，您需要对其进行预处理以创建 vpermd 的输入，但这只是将每个 i 映射到 2*i + 0..1 并添加相同，相同到双倍，然后是通道内随机播放，然后添加一个向量 0,1,0,1,...)
或者既然你提到了 512 位寄存器，AVX512F 为 vpermq 提供了矢量控制，而不仅仅是立即编译时常量。 你的置换是编译时常量吗？你展示的代码使用constin[]和perm[]，所以这没有帮助，编译器可以计算out的内容在如果两个输入都是编译时常量，则编译时间。对于 256 位寄存器，每个 uint8_t 应该是 1 位吗？或者你真的想要一个 256 字节的 shuffle（8 个 YMM 寄存器宽）？
谢谢！ in[] 不是编译时常量，但perm[] 是 - 我很乐意以任何合适的形式提供它（库的模板参数，汇编指令和立即数的掩码）。 vpermd 跨通道打乱 32 位整数块，但不打乱整数内的位。如前所述（可能不够清楚），大小 >= 8 的块的排列得到了处理。 所以要强调的是，我真的想要一个 256 位的 shuffle，而不是 BYTE---也就是说，也许你告诉我这样更快，并且每口存储一位/两位/一个半字节，而忽略其余部分是要走的路。
简单（但不切实际）的解决方案是手动编码每个可能的排列并通过模板专业化选择正确的排列。我怀疑是否有一个简单的解决方案对每种情况都是最佳的（例如，我确信对于某些排列，最有效的方法是乘法和位运算的组合——并且在编译时确定这听起来很难)。
确实有一个工具可以通过尝试全部找到programming.sirrida.de/bit_perm.html 上的最佳比特洗牌网络。但是，该代码仅适用于 64 位；我也更喜欢一个包罗万象的案例，我愿意为此牺牲几个周期。但有效的观点。

标签： c++ assembly sse avx avx2

【解决方案1】：

AVX2 256 位置换案例

我认为不可能编写出高效的通用 SSE4/AVX2/AVX-512 算法适用于所有向量大小（128、256、512 位）和元素粒度（位、位对、半字节、字节）。一个问题是存在许多 AVX2 指令例如，字节大小的元素，对于双字元素不存在，反之亦然。

下面讨论了 AVX2 256 位排列的情况。或许可以将这个案例的想法再用于其他案例。

这个想法是每一步从输入向量x 中提取 32 个（置换）位。在每个步骤中，从排列向量pos 中读取 32 个字节。这些pos 字节的第7..3 位确定需要x 中的哪个字节。右字节由模拟的 256 位宽 AVX2 通道交叉字节选择随机播放coded here by Ermlg。 pos 字节的位 2..0 确定要查找的位。使用_mm256_movemask_epi8，32 位被收集在一个_uint32_t 中此步骤重复 8 次，得到所有 256 个置换位。

代码看起来不是很优雅。尽管如此，我还是会感到惊讶如果明显更快，比如快两倍，AVX2 方法就会存在。

/*     gcc -O3 -m64 -Wall -mavx2 -march=skylake bitperm_avx2.c     */
#include <immintrin.h>
#include <stdio.h>
#include <stdint.h>

inline __m256i shuf_epi8_lc(__m256i value, __m256i shuffle);
int print_epi64(__m256i  a);

uint32_t get_32_bits(__m256i x, __m256i pos){
    __m256i pshufb_mask  = _mm256_set_epi8(0,0,0,0, 0,0,0,0, 128,64,32,16, 8,4,2,1, 0,0,0,0, 0,0,0,0, 128,64,32,16, 8,4,2,1);
    __m256i byte_pos     = _mm256_srli_epi32(pos, 3);                       /* which byte within the 32 bytes    */
            byte_pos     = _mm256_and_si256(byte_pos, _mm256_set1_epi8(0x1F)); /* mask off the unwanted bits */
    __m256i bit_pos      = _mm256_and_si256(pos, _mm256_set1_epi8(0x07));   /* which bit within the byte         */
    __m256i bit_pos_mask = _mm256_shuffle_epi8(pshufb_mask, bit_pos);       /* get bit mask                      */
    __m256i bytes_wanted = shuf_epi8_lc(x, byte_pos);                       /* get the right bytes               */
    __m256i bits_wanted  = _mm256_and_si256(bit_pos_mask, bytes_wanted);    /* apply the bit mask to get rid of the unwanted bits within the byte */
    __m256i bits_x8      = _mm256_cmpeq_epi8(bits_wanted, bit_pos_mask);    /* check if the bit is set           */        
            return _mm256_movemask_epi8(bits_x8);
}

__m256i get_256_bits(__m256i x, uint8_t* pos){ /* glue the 32 bit results together */
    uint64_t t0 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[0]));
    uint64_t t1 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[32]));
    uint64_t t2 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[64]));
    uint64_t t3 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[96]));
    uint64_t t4 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[128]));
    uint64_t t5 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[160]));
    uint64_t t6 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[192]));
    uint64_t t7 = get_32_bits(x, _mm256_loadu_si256((__m256i*)&pos[224]));
    uint64_t t10 = (t1<<32)|t0;
    uint64_t t32 = (t3<<32)|t2;
    uint64_t t54 = (t5<<32)|t4;
    uint64_t t76 = (t7<<32)|t6;
    return(_mm256_set_epi64x(t76, t54, t32, t10));
}


inline __m256i shuf_epi8_lc(__m256i value, __m256i shuffle){
/* Ermlg's lane crossing byte shuffle https://stackoverflow.com/a/30669632/2439725 */
const __m256i K0 = _mm256_setr_epi8(
    0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70,
    0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0);
const __m256i K1 = _mm256_setr_epi8(
    0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0, 0xF0,
    0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70, 0x70);
return _mm256_or_si256(_mm256_shuffle_epi8(value, _mm256_add_epi8(shuffle, K0)), 
    _mm256_shuffle_epi8(_mm256_permute4x64_epi64(value, 0x4E), _mm256_add_epi8(shuffle, K1)));
}


int main(){
    __m256i    input = _mm256_set_epi16(0x1234,0x9876,0x7890,0xABCD, 0x3456,0x7654,0x0123,0x4567,
                                        0x0123,0x4567,0x89AB,0xCDEF, 0xFEDC,0xBA98,0x7654,0x3210);
/* Example                                                                                         */
/*            240  224  208  192    176  160  144  128    112   96   80   64     48   32   16    0 */                        
/* input     1234 9876 7890 ABCD | 3456 7654 0123 4567 | 0123 4567 89AB CDEF | FEDC BA98 7654 3210 */
/* output    0000 0000 0012 00FF | 90AB 3210 7654 ABCD | 8712 1200 FF90 AB32 | 7654 ABCD 1087 7654 */
    uint8_t permutation[256] = {16,17,18,19,     20,21,22,23,      24,25,26,27,     28,29,30,31,
                                28,29,30,31,     32,33,34,35,      0,1,2,3,         4,5,6,7,
                                72,73,74,75,     76,77,78,79,      80,81,82,83,     84,85,86,87,      
                                160,161,162,163, 164,165,166,167,  168,169,170,171, 172,173,174,175,  
                                8,9,10,11,       12,13,14,15,      200,201,202,203, 204,205,206,207,
                                208,209,210,211, 212,213,214,215,  215,215,215,215, 215,215,215,215,
                                1,1,1,1,         1,1,1,1,          248,249,250,251, 252,253,254,255,
                                248,249,250,251, 252,253,254,255,  28,29,30,31,     32,33,34,35,
                                72,73,74,75,     76,77,78,79,      80,81,82,83,     84,85,86,87,
                                160,161,162,163, 164,165,166,167,  168,169,170,171, 172,173,174,175,
                                0,1,2,3,         4,5,6,7,          8,9,10,11,       12,13,14,15,
                                200,201,202,203, 204,205,206,207,  208,209,210,211, 212,213,214,215,
                                215,215,215,215, 215,215,215,215,  1,1,1,1,         1,1,1,1,
                                248,249,250,251, 252,253,254,255,  1,1,1,1,         1,1,1,1,
                                1,1,1,1,         1,1,1,1,          1,1,1,1,         1,1,1,1,
                                1,1,1,1,         1,1,1,1,          1,1,1,1,         1,1,1,1};
               printf("input = \n");
               print_epi64(input);
    __m256i    x = get_256_bits(input, permutation);
               printf("permuted input = \n");
               print_epi64(x);
               return 0;
}


int print_epi64(__m256i  a){
    uint64_t  v[4];
    int i;
    _mm256_storeu_si256((__m256i*)v,a);
    for (i = 3; i>=0; i--) printf("%016lX  ",v[i]);
    printf("\n");
    return 0;
}

具有示例排列的输出看起来正确：

$ ./a.out
input = 
123498767890ABCD  3456765401234567  0123456789ABCDEF  FEDCBA9876543210  
permuted input = 
00000000001200FF  90AB32107654ABCD  87121200FF90AB32  7654ABCD10877654

效率

如果你仔细看算法，你会发现只有一些操作取决于置换向量pos，而不是x。这意味着应用使用变量x 和固定pos 进行排列应该更有效而不是对变量x 和pos 应用排列。

以下代码说明了这一点：

/* apply the same permutation several times */
int perm_array(__m256i* restrict x_in, uint8_t* restrict pos, __m256i* restrict x_out){
    for (int i = 0; i<1024; i++){
            x_out[i]=get_256_bits(x_in[i], pos);
    }
    return 0;
}

使用 clang 和 gcc 编译成真的 nice code：在第 237 行循环 .L5 仅包含 16 vpshufbs 而不是 24。此外，vpaddbs 被吊出循环。请注意，循环内也只有一个vpermq。

我不知道 MSVC 是否会在循环之外提升这么多指令。如果没有，也许有可能通过手动修改代码来提高循环的性能。应该这样做，以便仅依赖于pos 而不依赖于x 的操作被提升到循环之外。

关于英特尔 Skylake 的性能：此循环的吞吐量可能受每个循环迭代大约 32 个端口 5 微操作。这意味着吞吐量在诸如perm_array 之类的循环上下文中，每 32 个 CPU 周期大约有 256 个置换位，或每个 CPU 周期大约 8 个置换位。

使用 AVX2 指令的 128 位排列

此代码与 256 位排列的情况非常相似。虽然只置换了 128 位，但 AVX2 的完整 256 位宽度寄存器用于实现最佳性能。这里不模拟字节混洗。这是因为存在进行字节混洗的有效单指令在 128 位通道内：vpshufb。

函数perm_array_128 测试位置换的性能对于固定排列和可变输入x。组装循环包含大约 11 个端口 5 (p5) 微操作，如果我们假设一个英特尔 Skylake CPU。这 11 个 p5 微操作至少需要 11 个 CPU 周期（吞吐量）。因此，在最好的情况下，我们获得的吞吐量约为每个周期 12 个置换位，这大约是 256 位置换情况的 1.5 倍。

/*     gcc -O3 -m64 -Wall -mavx2 -march=skylake bitperm128_avx2.c     */
#include <immintrin.h>
#include <stdio.h>
#include <stdint.h>

int print128_epi64(__m128i  a);

uint32_t get_32_128_bits(__m256i x, __m256i pos){                           /* extract 32 permuted bits out from 2x128 bits   */
    __m256i pshufb_mask  = _mm256_set_epi8(0,0,0,0, 0,0,0,0, 128,64,32,16, 8,4,2,1, 0,0,0,0, 0,0,0,0, 128,64,32,16, 8,4,2,1);
    __m256i byte_pos     = _mm256_srli_epi32(pos, 3);                       /* which byte do we need within the 16 byte lanes. bits 6,5,4,3 select the right byte */
            byte_pos     = _mm256_and_si256(byte_pos, _mm256_set1_epi8(0xF)); /* mask off the unwanted bits (unnecessary if _mm256_srli_epi8 would have existed   */
    __m256i bit_pos      = _mm256_and_si256(pos, _mm256_set1_epi8(0x07));   /* which bit within the byte                 */
    __m256i bit_pos_mask = _mm256_shuffle_epi8(pshufb_mask, bit_pos);       /* get bit mask                              */
    __m256i bytes_wanted = _mm256_shuffle_epi8(x, byte_pos);                /* get the right bytes                       */
    __m256i bits_wanted  = _mm256_and_si256(bit_pos_mask, bytes_wanted);    /* apply the bit mask to get rid of the unwanted bits within the byte */
    __m256i bits_x8      = _mm256_cmpeq_epi8(bits_wanted, bit_pos_mask);    /* set all bits if the wanted bit is set     */        
            return _mm256_movemask_epi8(bits_x8);                           /* move most significant bit of each byte to 32 bit register */
}


__m128i permute_128_bits(__m128i x, uint8_t* pos){      /* get bit permutations in 32 bit pieces and glue them together */
    __m256i  x2 = _mm256_broadcastsi128_si256(x);   /* broadcast x to the hi and lo lane                            */
    uint64_t t0 = get_32_128_bits(x2, _mm256_loadu_si256((__m256i*)&pos[0]));
    uint64_t t1 = get_32_128_bits(x2, _mm256_loadu_si256((__m256i*)&pos[32]));
    uint64_t t2 = get_32_128_bits(x2, _mm256_loadu_si256((__m256i*)&pos[64]));
    uint64_t t3 = get_32_128_bits(x2, _mm256_loadu_si256((__m256i*)&pos[96]));
    uint64_t t10 = (t1<<32)|t0;
    uint64_t t32 = (t3<<32)|t2;
    return(_mm_set_epi64x(t32, t10));
}

/* Test loop performance with the following loop (see assembly) -> 11 port5 uops inside the critical loop */
/* Use gcc -O3 -m64 -Wall -mavx2 -march=skylake -S bitperm128_avx2.c to generate the assembly             */
int perm_array_128(__m128i* restrict x_in, uint8_t* restrict pos, __m128i* restrict x_out){
    for (int i = 0; i<1024; i++){
            x_out[i]=permute_128_bits(x_in[i], pos);
    }
    return 0;
}


int main(){
    __m128i    input = _mm_set_epi16(0x0123,0x4567,0xFEDC,0xBA98,  0x7654,0x3210,0x89AB,0xCDEF);
/* Example                                                                                         */
/*             112   96   80   64     48   32   16    0 */                        
/* input      0123 4567 FEDC BA98   7654 3210 89AB CDEF */
/* output     8FFF CDEF DCBA 08EF   CDFF DCBA EFF0 89AB */
    uint8_t permutation[128] = {16,17,18,19,     20,21,22,23,      24,25,26,27,     28,29,30,31,
                                32,32,32,32,     36,36,36,36,      0,1,2,3,         4,5,6,7,
                                72,73,74,75,     76,77,78,79,      80,81,82,83,     84,85,86,87,      
                                0,0,0,0,         0,0,0,0,          8,9,10,11,       12,13,14,15,      
                                0,1,2,3,         4,5,6,7,          28,29,30,31,     32,33,34,35,
                                72,73,74,75,     76,77,78,79,      80,81,82,83,     84,85,86,87,
                                0,1,2,3,         4,5,6,7,          8,9,10,11,       12,13,14,15,
                                1,1,1,1,         1,1,1,1,          1,1,1,1,         32,32,32,1};
               printf("input = \n");
               print128_epi64(input);
    __m128i    x = permute_128_bits(input, permutation);
               printf("permuted input = \n");
               print128_epi64(x);
               return 0;
}


int print128_epi64(__m128i  a){
  uint64_t  v[2];
  int i;
  _mm_storeu_si128((__m128i*)v,a);
  for (i = 1; i>=0; i--) printf("%016lX  ",v[i]);
  printf("\n");
  return 0;
}

一些任意排列的示例输出：

$ ./a.out
input = 
01234567FEDCBA98  7654321089ABCDEF  
permuted input = 
8FFFCDEFDCBA08EF  CDFFDCBAEFF089AB

【讨论】：

同样的想法可用于置换 256 位宽 AVX2 寄存器中的半字节。有 64 个半字节，所以我们有 2 个步骤。在每个步骤中，基于排列索引的位 5..1 排列 32 个字节（通道交叉）。这样的字节包含一个高半字节和一个低半字节。置换索引的位 0 用作此高半字节和低半字节之间的 vpblendvb 的选择器。这给出了 32 个半字节。在下一步中，从 256 位寄存器中选择另外 32 个半字节。一些简单的位/字节操作最终会导致结果。
这太棒了！非常感谢，这已经非常有帮助，而且解释得很好！我还可以看到如何有效地进行两位随机播放：只需将偶数位和奇数位以交替方式存储在两个 256 位宽的 AVX2 寄存器中并置换两次。整洁的。你有任何关于如何达到 128 位的指示吗？
很好用！关于 128 位：您对使用 SSEx 或使用 AVX2 的 128 位排列感兴趣吗？尽管输出只有 128 位，但使用 256 位宽的 AVX2 寄存器和指令可以使计算受益（更好的吞吐量）。 128 位情况的一个优点是不需要模拟 256 位宽的字节洗牌，因为 128 位宽的字节洗牌只有一条指令：pshufb。
AVX2 非常好！一旦我有使用这个实现的位，我将发布一些真实世界的基准。
@JBausch 添加了 128 位大小写。