c - 2 个 AVX-512 vector 元素的交错合并 - C 内在函数

Question

我想以尽可能少的时钟周期将 2 个 AVX-512 vector 的元素合并到另外两个 vector 中。

问题具体如下:

// 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};

最简单的方法是将 vector (a 和 b)复制到内存并通过直接索引创建 vector (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 内在函数的新手，所以一些解释将不胜感激。

感谢您的帮助!

Answer 1

感谢上面提到的评论，解决此问题的一种方法是使用 vpermt2w 或内在的 _mm512_mask_permutex2var_epi16。

在 Skylake-avx512 和 Ice Lake CPU ( https://uops.info/ ) 上，vpermt2w 解码为 3 uops(其中 2 只能在端口 5 上运行)。总体而言，它有 7 个周期的延迟，每 2 个周期的吞吐量为 1 个。

使用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 );
}

此处显示的简单方法仅供引用，但对于不是编译时常量的输入 vector ，它使用 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]
            );

}