assembly 使用SIMD指令执行任意128/256/512位排列的最快方法是什么?

zaqlnxep  于 2023-06-23  发布在  其他
关注(0)|答案(2)|浏览(171)

我想对宽度为128、256或512位的CPU寄存器(xmm、ymm或zmm)执行单个位、位对和半字节(4位)的任意排列;这应该尽可能快。为此,我正在研究SIMD指令。有没有人知道一种方法来做到这一点/一个实现它的库?我在Windows上使用MSVC,在Linux上使用GCC,主机语言是C或C++。谢谢!
我给出了一个任意的排列,需要对大量的位向量/位向量对/半字节进行 Shuffle 。我知道如何对64位值中的位执行此操作,例如。使用Benes网络。
或者在更宽的SIMD寄存器上混洗8位或更大的块,例如使用Agner Fog的GPLed VectorClass库(https://www.agner.org/optimize/vectorclass.pdf)作为模板元编程函数,该函数在给定shuffle作为模板参数的情况下,从AVX 2通道内字节shuffle和/或较大元素通道交叉shuffle中构建shuffle。
然而,对排列进行更细粒度的细分(细分为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]];

如前所述,我有一个<= 64位的解决方案(即64位,32位对和16个半字节)。对于大小为8、16、32等的块也解决了该问题。在更宽的SIMD寄存器上。
编辑:澄清一下,置换是一个编译时常数(但不只是一个特定的,我将对每个给定的置换编译一次程序)。

vcudknz3

vcudknz31#

AVX 2 256位置换情况

我认为不可能编写一个有效的通用SSE 4/AVX 2/AVX-512算法,适用于所有向量大小(128,256,512位)和元素粒度(位,位对,半字节,字节)。一个问题是,对于例如字节大小元素存在的许多AVX 2指令对于双字元素不存在,反之亦然。
下面讨论AVX 2 256位置换情况。也许可以把这个案例的想法用于其他案例。
其思想是从输入向量x中每步提取32(置换)位。在每个步骤中,从置换向量pos读取32个字节。这些pos字节的位7..3确定需要x中的哪个字节。右字节由仿真的256位宽AVX 2通道交叉字节混洗coded here by Ermlg选择。pos字节的位2..0确定要查找的位。使用_mm256_movemask_epi8,32位被收集在一个_uint32_t中。此步骤重复8次,以获得所有256个置换位。
代码看起来不太优雅。尽管如此,如果有一个明显更快的AVX 2方法,比如说快两倍,我会感到惊讶。

/*     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的置换应该比应用具有变量xpos的置换更有效。
下面的代码说明了这一点:

/* 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个vpshufb s,而不是24个。此外,vpaddb s被提升出循环。请注意,循环中也只有一个vpermq
我不知道MSVC是否会在循环外提升这么多指令。如果没有,可以通过手动修改代码来提高循环的性能。应该这样做,使得只依赖于pos而不依赖于x的操作被提升到循环之外。
关于英特尔Skylake的性能:该循环的吞吐量可能受到每个循环迭代约32个端口5个微操作的限制。这意味着在循环上下文中,例如perm_array的吞吐量大约是每32个CPU周期256个置换位,或者每CPU周期8个置换位。

使用AVX 2指令的128位排列

该代码与256位置换的情况非常相似。虽然只有128位被置换,但AVX 2寄存器的全部256位宽用于实现最佳性能。这里不模拟字节混洗。这是因为在128位通道内存在有效的单个指令来进行字节混洗:vpshufb .
函数perm_array_128测试固定置换和可变输入x的比特置换的性能。如果我们假设一个Intel Skylake CPU,则汇编循环包含大约11个端口5(p5)微操作。这些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
wvyml7n5

wvyml7n52#

AVX2

上面的答案很好,但如果数据在内存中,我们可以做得更好一点:

void fill_avx2_perm_table(__m256i table[24], uint8_t idx[256]) {
    __m256i bit_idx_mask = _mm256_set1_epi8(0x7);
    __m256i byte_idx_mask = _mm256_set1_epi8(0xf);
    __m256i bit_mask_lookup = _mm256_set1_epi64x(0x8040201008040201);
    __m256i mask_out = _mm256_set1_epi8(-1);

    for (int i = 0; i < 8; ++i) {
        __m256i perm_32 = _mm256_loadu_si256((const __m256i*)(idx + 32 * i)); // 1 bit set -> comes from bits 127 .. 255
        __m256i shuf = _mm256_and_si256(_mm256_srli_epi32(perm_32, 3), byte_idx_mask);

        __m256i shuf_lo = _mm256_blendv_epi8(shuf, mask_out, perm_32);
        __m256i shuf_hi = _mm256_blendv_epi8(mask_out, shuf, perm_32);
        __m256i bit_mask = _mm256_shuffle_epi8(bit_mask_lookup, _mm256_and_si256(byte_idx_mask, perm_32));

        _mm256_store_si256(table + 2*i, shuf_lo);
        _mm256_store_si256(table + 2*i + 1, shuf_hi);
        _mm256_store_si256(table + 16 + i, bit_mask);
    }
}

void permute_256_array(char* arr, size_t len, uint8_t idx[256]) {
    __m256i perm_table[16 /* shuffles */ + 8 /* bit masks */];

    fill_avx2_perm_table(perm_table, idx);
    __m256i zero = _mm256_setzero_si256();

    char* end = arr + len * 32;
    for (; arr < end; arr += 32) {
        __m256i lo = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i*) arr));
        __m256i hi = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i*) arr + 1));
        __m256i lo_source, hi_source, bit_mask, bits;

        uint32_t result;

#define DO_ITER(i)  \
            lo_source = _mm256_shuffle_epi8(lo, _mm256_loadu_si256(perm_table + 2 * i)); \
            hi_source = _mm256_shuffle_epi8(hi, _mm256_loadu_si256(perm_table + 2 * i + 1)); \
\
            bit_mask = _mm256_loadu_si256(perm_table + 16 + i); \
            bits = _mm256_and_si256(_mm256_or_si256(lo_source, hi_source), bit_mask); \
            bits = _mm256_cmpeq_epi8(bits, bit_mask); \
\
            result = _mm256_movemask_epi8(bits); \
            memcpy(arr + 4 * i, &result, 4);

        DO_ITER(0) DO_ITER(1) DO_ITER(2) DO_ITER(3) DO_ITER(4) DO_ITER(5) DO_ITER(6) DO_ITER(7)
#undef DO_ITER
    }
}

vpermq不同,从内存向ymm寄存器的所有256位广播128位不使用混洗µop。不幸的是,编译器有时似乎会在存储vpmovmskb结果之前将其插入到向量寄存器中,这可以在接受的答案的编译器输出中看到。可以在存储之间插入asm内存占用器(asm volatile ("" ::: "memory");),不幸的是,这会破坏指令重新排序,或者(在我的应用程序中)链接汇编例程。
在使用-march=core-avx2的Cascade Lake上,我得到**~12.3位/周期**。

AVX512BW / VBMI

我们可以用AVX 512做得更好。大多数示例将使用以下函数,类似于原始AVX 2答案中get_32_128_bits的前几行。

const __m512i bit_idx_mask = _mm512_set1_epi8(0x7);
const __m512i bit_mask_lookup = _mm512_set1_epi64(0x8040201008040201);

void get_permute_constants(__m512i idx, __m512i* byte_idx, __m512i* bit_mask) {
    *byte_idx = _mm512_srli_epi32(_mm512_andnot_si512(bit_idx_mask, idx), 3);
    idx = _mm512_and_si512(idx, bit_idx_mask);
    *bit_mask = _mm512_shuffle_epi8(bit_mask_lookup, idx);
}

对于循环不变排列,这些值应该被提升而不是重复计算。

64位

在64位及更小的特殊情况下,我们可以使用vpshufbitqmb选择64位到掩码寄存器中,并使用kmovq m64, k直接存储到内存中。如果vpshufbitqmb不可用,我们可以用字节混洗来模拟它,然后用所有需要的位的掩码来模拟vptestmb。(vptestmb计算两个向量寄存器的逻辑与,并对每个非零字节向掩码寄存器写入1。)

#ifdef __AVX512BITALG__
void permute_64(uint64_t* in, uint64_t* out, __m512i idx) { 
    __m512i in_v = _mm512_set1_epi64(*in); 
    __mmask64 permuted_bits = _mm512_bitshuffle_epi64_mask(in_v, idx); 
    _store_mask64((__mmask64*) out, permuted_bits); 
} 
#else
void permute_64(uint64_t* in, uint64_t* out, __m512i idx) {
    __m512i byte_idx, bit_mask;
    get_permute_constants(idx, &byte_idx, &bit_mask);

    __m512i in_v = _mm512_set1_epi64(*in);
    __m512i in_shuffled = _mm512_shuffle_epi8(in_v, byte_idx);
    __mmask64 permuted_bits = _mm512_test_epi8_mask(in_shuffled, bit_mask);

    _store_mask64((__mmask64*) out, permuted_bits);
}
#endif

作为循环的一部分的主力:

1ed8:       62 f2 fd 48 59 00       vpbroadcastq zmm0,QWORD PTR [rax]
1ede:       48 83 c0 08             add    rax,0x8
1ee2:       62 f2 7d 48 8f c1       vpshufbitqmb k0,zmm0,zmm1
1ee8:       c4 e1 f8 91 40 f8       kmovq  QWORD PTR [rax-0x8],k0
1eee:       4c 39 e8                cmp    rax,r13
1ef1:       75 e5                   jne    1ed8

在Ice Lake上观察到的性能为~1.38周期/元素,或**~46位/周期**(由于无法访问性能计数器而出现一些错误)。回退实现以2周期/元素或32位/周期运行,因为vptestmbvpshufb竞争同一端口。

128位

void permute_128(const char* in, char* out,
    __m512i byte_idx1, __m512i bit_mask1,
    __m512i byte_idx2, __m512i bit_mask2) {
    __m512i in_v = _mm512_broadcast_i32x4(_mm_loadu_si128((const __m128i*) in));

    __mmask64 permuted_1 = _mm512_test_epi8_mask(
        _mm512_shuffle_epi8(in_v, byte_idx1), bit_mask1);
    __mmask64 permuted_2 = _mm512_test_epi8_mask(
        _mm512_shuffle_epi8(in_v, byte_idx2), bit_mask2);

    _store_mask64((__mmask64*) out, permuted_1);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 1, permuted_2);
}

void permute_128_array(char* arr, size_t count, uint8_t idx[128]) {
    __m512i idx1 = _mm512_loadu_si512(idx);
    __m512i idx2 = _mm512_loadu_si512(idx + 64);

    __m512i byte_idx1, bit_mask1, byte_idx2, bit_mask2;
    get_permute_constants(idx1, &byte_idx1, &bit_mask1);
    get_permute_constants(idx2, &byte_idx2, &bit_mask2);

    count *= 16;

    for (size_t i = 0; i < count; i += 16)
        permute_128(arr + i, arr + i, byte_idx1, bit_mask1, byte_idx2, bit_mask2);
}

循环始终以4个周期/元素或32位/周期运行,瓶颈与之前相同(vpshufb/vptestmb)。

256位

128位通道交叉字节粒度 Shuffle (vpermb和朋友)仅在AVX 512 VBMI(Ice Lake和更高版本; Zen 4). AVX 512 BW版本与AVX 2版本没有太大区别:256位输入在两个128位组块中广播,使用合并掩码vpshufb混洗。

void fill_avx512_perm_table(__m512i table[8], uint8_t idx[256], __mmask64 keeps[4]) {
    for (int i = 0; i < 4; ++i) {
        __m512i perm_64 = _mm512_loadu_si512(idx + 64 * i);
        get_permute_constants(perm_64, table + i, table + i + 4);

        if (keeps) // we'll use this function later w/o needing masks
            keeps[i] = _mm512_movepi8_mask(perm_64);  // hi bit -> 128 .. 255
    }
}

void permute_256(const char* in, char* out, __m512i perm_table[8], __mmask64 keeps[4]) {
    __m512i in_v1 = _mm512_broadcast_i32x4(_mm_loadu_si128((const __m128i*) in));
    __m512i in_v2 = _mm512_broadcast_i32x4(_mm_loadu_si128((const __m128i*) in + 1));
 
    __m512i ones = _mm512_set1_epi8(1);

#define DO_INTEL_OPT 0

#define DO_ITER(i, mask_reg) __mmask64 mask_reg; { \
    __m512i shuffle = perm_table[i]; \
    __m512i perm = _mm512_shuffle_epi8(in_v1, shuffle); \
    perm = _mm512_mask_shuffle_epi8(perm, keeps[i], in_v2, shuffle); \
    if (DO_INTEL_OPT) { \
        perm = _mm512_andnot_si512(perm, perm_table[i + 4]); \
        perm = _mm512_sub_epi8(perm, ones); \
        mask_reg = _mm512_movepi8_mask(perm); \
    } else mask_reg = _mm512_test_epi8_mask(perm, perm_table[i + 4]); \
}

    DO_ITER(0, permuted_1) DO_ITER(1, permuted_2) DO_ITER(2, permuted_3) DO_ITER(3, permuted_4)

    _store_mask64((__mmask64*) out, permuted_1);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 1, permuted_2);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 2, permuted_3);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 3, permuted_4);
#undef DO_ITER
#undef DO_INTEL_OPT
}

void permute_256_array(char* arr, size_t len, uint8_t idx[256]) {
    __m512i table[8];
    __mmask64 keeps[4];
    fill_avx512_perm_table(table, idx, keeps);

    char* end = arr + 32 * len;
    for (; arr < end; arr += 32)
        permute_256(arr, arr, table, keeps);
}

我得到了12个周期/ 256位元素(21位/周期),这甚至没有AVX 2纯版本的两倍,因为比较到掩码寄存器与混洗竞争。* 不 * 使用端口5的向量到掩码操作似乎是vpmovb2m等。似乎没有一个有利可图的方法来获得HSB的相关位使用单个指令(没有可变字节移位),但我们可以在两个以上启用DO_INTEL_OPT-这将效率提高到理论上10个周期/元素(25.6位/周期)。这可能是或可能不值得的复杂性,我敢肯定是严格更糟的AMD。
VBMI解决方案不起眼,以32位/周期运行:

void permute_256(const char* in, char* out, __m512i perm_table[8], __mmask64 keeps[4] /* unused */) {
    __m512i in_v = _mm512_broadcast_i64x4(_mm256_loadu_si256((const __m256i*) in));

#define DO_ITER(i, mask_reg) __mmask64 mask_reg; { \
    __m512i shuffle = perm_table[i]; \
    __m512i perm = _mm512_permutexvar_epi8(shuffle, in_v); \
    mask_reg = _mm512_test_epi8_mask(perm, perm_table[i + 4]); \
}

    DO_ITER(0, permuted_1) DO_ITER(1, permuted_2) DO_ITER(2, permuted_3) DO_ITER(3, permuted_4)

    _store_mask64((__mmask64*) out, permuted_1);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 1, permuted_2);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 2, permuted_3);
    asm volatile ("" ::: "memory");
    _store_mask64((__mmask64*) out + 3, permuted_4);
#undef DO_ITER
}

512位

我不需要这么广泛的排列,但用VBMI也不难。唯一的问题是你必须将排列索引存储为每个16位,但是一旦你将其转换为字节+位索引的形式,它应该和以前的解决方案一样--只需用普通的旧的512位vmovdqu替换vbroadcasti64x4加载即可。

相关问题