I want to merge elements of 2 AVX-512 vectors into two other vectors with the least possible number of clock cycles.
The problem specifics are as follows:
// 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};
The naive way is to copy the vectors (a and b) to memory and create vectors (A and B) by direct indexing as below:
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] );
I would also need to do similar merges but with different strides, for example:
// 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};
What are the most suitable AVX-512 intrinsics to solve this problem? Some explanation would be greatly appreciated as I am a newbie to AVX-512 intrinsics.
Thank you for your help!
Thanks to the comments mentioned above, one way to solve this problem is using vpermt2w or the intrinsic _mm512_mask_permutex2var_epi16.
On Skylake-avx512 and Ice Lake CPUs (https://uops.info/), vpermt2w decodes to 3 uops (2 of which can only run on port 5). Overall it has 7 cycle latency, with a throughput of one per 2 cycles.
The optimized code using vpermt2w is as follows:
#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 );
}
And the naive way is shown here for reference, but it compiles very inefficiently with GCC for input vectors that aren't compile-time constants.
#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]
);
}