Какой самый быстрый способ выполнить произвольную 128/256/512-битную перестановку с использованием SIMD-инструкций?

Question

Я хочу выполнить произвольную перестановку единичных битов, пар битов и кусков (4 бита) в регистре ЦП (xmm, ymm или zmm) шириной 128, 256 или 512 бит;это должно быть как можно быстрее.Для этого я искал инструкции SIMD.Кто-нибудь знает способ сделать это / библиотека, которая реализует это?Я использую MSVC в Windows и GCC в Linux, а язык хоста - C или C ++.Спасибо!

Мне дана произвольная перестановка, и мне нужно перетасовать большое количество битовых векторов / пар битовых векторов / полубайтов.Я знаю, как это сделать для битов в пределах 64-битного значения, например, , используя сеть Бенеша .

Или тасовать блоки 8-бит и более вокруг более широких регистров SIMD, напримериспользуя библиотеку Agner Fog GPLed VectorClass (https://www.agner.org/optimize/vectorclass.pdf) для функции метапрограммирования шаблона, которая создает тасования из тасов байтов 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]];

Как я уже сказал, у меня есть решение для <= 64 бит (что будет 64 бита, 32 пары битов и 16 полубайтов).Эта проблема также решается для блоков размером 8, 16, 32 и т. Д. В более широких регистрах SIMD. </p>

РЕДАКТИРОВАТЬ: для пояснения, перестановка является константой времени компиляции (но не одной конкретной, яскомпилирую программу один раз для данной перестановки).

Answer 1

Случай AVX2 256-битной перестановки

Я не думаю, что можно написать эффективный универсальный алгоритм SSE4 / AVX2 / AVX-512, который работает для всех векторных размеров (128,256, 512 бит) и гранулярность элементов (биты, пары битов, полубайты, байты).Одна проблема состоит в том, что многие инструкции AVX2, которые существуют, например, для элементов размера байта, не существуют для элементов двойного слова, и наоборот.

Ниже обсуждается случай 256-битной перестановки AVX2.Можно было бы использовать идеи этого случая для других случаев.

Идея состоит в том, чтобы извлечь 32 (переставленных) бита за шаг из входного вектора x.На каждом шаге считывается 32 байта из вектора перестановок pos.Биты 7..3 из этих pos байтов определяют, какой байт из x необходим.Правый байт выбирается с помощью эмулируемой перестановки байтов AVX2 шириной 256 битов , закодированной здесь Ermlg .Биты 2..0 из pos байтов определяют, какой бит ищется.При _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 это скомпилируется в действительно хороший код : цикл .L5 в строке 237 содержит только 16 vpshufb с вместо 24. Кроме того,vpaddb s подняты из петли.Обратите внимание, что внутри цикла есть только один vpermq.

Я не знаю, будет ли MSVC выводить такое количество инструкций за пределы цикла.Если нет, то возможно улучшить производительность цикла, изменив код вручную.Это должно быть сделано так, чтобы операции, которые зависят только от pos, а не от x, были подняты за пределы цикла.

Относительно производительности на Intel Skylake: пропускная способность этого цикласкорее всего, ограничено примерно 32 микрооперациями порта 5 на каждую итерацию цикла.Это означает, что пропускная способность в контексте цикла, такого как perm_array, составляет около 256 переставленных битов на 32 цикла ЦП или около 8 переставленных бит на цикл ЦП.

128-битные перестановки с использованием инструкций AVX2

Этот код очень похож на случай 256-битной перестановки.Хотя переставляются только 128 битов, полная 256-битная ширина регистров AVX2 используется для достижения наилучшей производительности.Здесь байтовые тасовки не эмулируются.Это связано с тем, что существует эффективная единая инструкция для перестановки байтов в пределах 128-битных полос: vpshufb.

Функция perm_array_128 проверяет производительность перестановки битов для фиксированной перестановки и ввода переменной x.Цикл сборки содержит около 11 портов 5 (p5) микроопераций, если предположить процессор Intel Skylake.Эти 11 микроопераций p5 занимают как минимум 11 циклов ЦП (пропускная способность).Таким образом, в лучшем случае мы получаем пропускную способность около 12 переставленных битов за цикл, что примерно в 1,5 раза быстрее, чем при 256-битной перестановке.

/*     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