如何用AVX-512实现16位和32位整数的插入与提取操作？

嘿，这个问题我之前也碰到过！AVX-512确实没有直接对应AVX那些_mm256_insert/extract_epi16/32的内置函数，但咱们可以用现有AVX-512指令组合出完全等效的功能，下面我给你逐个实现并解释：

1. 实现_mm512_insert_epi16 #

这个函数负责把16位整数插入到__m512i向量的指定位置，核心思路是用掩码混合原向量和广播后的插入值：

#include <immintrin.h>

__m512i _mm512_insert_epi16(__m512i a, int16_t i, int index) {
    // 生成仅目标位置为1的掩码（__m512i包含32个16位元素，index范围0~31）
    __mmask32 mask = 1 << index;
    // 将插入值广播为全是i的512位向量
    __m512i val_vec = _mm512_set1_epi16(i);
    // 掩码混合：掩码位为1的位置用val_vec的元素，其余保留原向量a的元素
    return _mm512_mask_blend_epi16(mask, a, val_vec);
}

2. 实现_mm512_insert_epi32 #

逻辑和16位版本完全一致，只是针对32位元素调整掩码和广播函数：

__m512i _mm512_insert_epi32(__m512i a, int32_t i, int index) {
    // __m512i包含16个32位元素，index范围0~15
    __mmask16 mask = 1 << index;
    __m512i val_vec = _mm512_set1_epi32(i);
    return _mm512_mask_blend_epi32(mask, a, val_vec);
}

3. 实现_mm512_extract_epi16 #

要提取指定位置的16位元素，我们可以用AVX-512的掩码提取指令，避免内存拷贝：

int _mm512_extract_epi16(__m512i a, int index) {
    __mmask32 mask = 1 << index;
    // 将选中的元素提取到128位向量的低位，其余位置置0
    __m128i extracted = _mm512_mask_extractu_epi16(_mm_setzero_si128(), mask, a);
    // 取出128位向量的第一个16位元素
    return _mm_extract_epi16(extracted, 0);
}

4. 实现_mm512_extract_epi32 #

同样用掩码提取指令，针对32位元素调整参数：

int _mm512_extract_epi32(__m512i a, int index) {
    __mmask16 mask = 1 << index;
    __m128i extracted = _mm512_mask_extractu_epi32(_mm_setzero_si128(), mask, a);
    return _mm_extract_epi32(extracted, 0);
}

一些注意事项 #

• 索引范围：必须保证index在合法区间内：epi16是0_{31，epi32是0}15，否则会触发未定义行为。
• 编译选项：编译时需要添加AVX-512相关选项，比如GCC/Clang用-mavx512bw -mavx512f，MSVC用/arch:AVX512。
• 立即数优化：如果index是编译时已知的常量，其实可以直接用编译器内置的_mm512_extract_epi16/_mm512_extract_epi32（这些函数支持立即数索引），上面的实现主要针对变量索引的场景。

内容的提问来源于stack exchange，提问作者Daniel