Блог

Я хочу объединить элементы из 2 векторов AVX-512 в два других вектора с минимально возможным количеством тактов.

Специфика проблемы заключается в следующем:

 // inputs
__m512i a = {a0, a1, ..., a31}; // 32x 16-bit int16_t integers
__m512i b = {b0, b1, ..., b31}; // 32x 16-bit int16_t integers

// desired output
__m512i A = {a0 , b0 , a1 , b1 , ..., a15, b15};
__m512i B = {a16, b16, a17, b17, ..., a31, b31};
  

Наивный способ состоит в том, чтобы скопировать векторы (a и b) в память и создать векторы (A и B) путем прямого индексирования, как показано ниже:

 union U512i {
    __m512i vec;
    alignas(64) int16_t vals[32];
};

U512i ta = { a };
U512i tb = { b }

U512i A = _mm512_set_epi16( tb.vals[15], ta.vals[15], ... tb.vals[0], ta.vals[0] );
U512i B = _mm512_set_epi16( tb.vals[31], ta.vals[31], ... tb.vals[16], ta.vals[16] );
  

Мне также нужно было бы выполнить аналогичные слияния, но с разными шагами, например:

 // inputs
__m512i a = {a0, a1, ..., a31}; // 32x 16-bit int16_t integers
__m512i b = {b0, b1, ..., b31}; // 32x 16-bit int16_t integers

// desired output
__m512i A = {a0 , a1 , b0 , b1 , ..., a14, a15, b14, b15};
__m512i B = {a16, a17, b16, b17, ..., a30, a31, b30, b31};
  

Каковы наиболее подходящие встроенные компоненты AVX-512 для решения этой проблемы? Я был бы весьма признателен за некоторые пояснения, поскольку я новичок во встроенных компонентах AVX-512.

Спасибо вам за вашу помощь!

Благодаря комментариям, упомянутым выше, одним из способов решения этой проблемы является использование vpermt2w или встроенного _mm512_mask_permutex2var_epi16 .

На процессорах Skylake-avx512 и Ice Lake (https://uops.info /), vpermt2w декодирует до 3 операций ввода-вывода (2 из которых могут выполняться только на порту 5). В целом он имеет задержку в 7 циклов с пропускной способностью один на 2 цикла.

Оптимизированный код с использованием vpermt2w выглядит следующим образом:

 #include <immintrin.h>
#include <inttypes.h>

void foo(__m512i a, __m512i b) {

    __m512i A, B;
    __m512i idx1 = _mm512_set_epi16( 47, 15, 46, 14, 45, 13, 44, 12, 43, 11, 42, 10, 41, 9, 40, 8, 39, 7, 38, 6, 37, 5, 36, 4, 35, 3, 34, 2, 33, 1, 32, 0 );
    __m512i idx2 = _mm512_set_epi16(
        47   16, 15   16, 46   16, 14   16, 45   16, 13   16, 44   16, 12   16, 43   16, 11   16, 42   16, 10   16, 41   16, 9   16, 40   16, 8   16,
        39   16, 7   16, 38   16, 6   16, 37   16, 5   16, 36   16, 4   16, 35   16, 3   16, 34   16, 2   16, 33   16, 1   16, 32   16, 0   16 );

    A = _mm512_mask_permutex2var_epi16( a, 0xFFFFFFFF, idx1, b );
    B = _mm512_mask_permutex2var_epi16( a, 0xFFFFFFFF, idx2, b );
}
  

И наивный способ показан здесь для справки, но он очень неэффективно компилируется с GCC для входных векторов, которые не являются константами времени компиляции.

 #include <immintrin.h>
#include <inttypes.h>

union U512i {
    __m512i vec;
    alignas(64) int16_t vals[32];
};

void foo(__m512i a, __m512i b) {

    __m512i A, B;

    U512i u_a = { a };
    U512i u_b = { b };
    A = _mm512_set_epi16 (
            u_b.vals[15], u_a.vals[15], u_b.vals[14], u_a.vals[14],
            u_b.vals[13], u_a.vals[13], u_b.vals[12], u_a.vals[12],
            u_b.vals[11], u_a.vals[11], u_b.vals[10], u_a.vals[10],
            u_b.vals[9], u_a.vals[9], u_b.vals[8], u_a.vals[8],
            u_b.vals[7], u_a.vals[7], u_b.vals[6], u_a.vals[6],
            u_b.vals[5], u_a.vals[5], u_b.vals[4], u_a.vals[4],
            u_b.vals[3], u_a.vals[3], u_b.vals[2], u_a.vals[2],
            u_b.vals[1], u_a.vals[1], u_b.vals[0], u_a.vals[0]
            );

    B = _mm512_set_epi16 (
            u_b.vals[31], u_a.vals[31], u_b.vals[30], u_a.vals[30],
            u_b.vals[29], u_a.vals[29], u_b.vals[28], u_a.vals[28],
            u_b.vals[27], u_a.vals[27], u_b.vals[26], u_a.vals[26],
            u_b.vals[25], u_a.vals[25], u_b.vals[24], u_a.vals[24],
            u_b.vals[23], u_a.vals[23], u_b.vals[22], u_a.vals[22],
            u_b.vals[21], u_a.vals[21], u_b.vals[20], u_a.vals[20],
            u_b.vals[19], u_a.vals[19], u_b.vals[18], u_a.vals[18],
            u_b.vals[17], u_a.vals[17], u_b.vals[16], u_a.vals[16]
            );

}