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

Question

AVX 具有将 16 位和 32 位整数插入和提取到 __m256i 向量的指令：_mm256_insert_epi16、_mm256_insert_epi32、_mm256_extract_epi16、_mm256_extract_epi32。

但是，AVX-512 似乎没有等效的说明。为__m512i 向量实现这些方法的适当方法是什么？即

• __m512i _mm512_insert_epi16(__m512i a, __int16 i, int index)
• __m512i _mm512_insert_epi32(__m512i a, __int32 i, int index)
• int _mm512_extract_epi16(__m512i a, int index)
• int _mm512_extract_epi32(__m512i a, int index)

Answer 1

相关：

• How can I write a QuadWord from AVX512 register zmm26 to the rax register? - 其中大部分也适用于提取 32 位和 16 位元素。
• Move an int64_t to the high quadwords of an AVX2 __m256i vector 显示插入，并且大部分内容应该适用于 32 位和可能的 16 位元素。 （尽管vpblendw 对两条车道重复混合控制，与vpblendd 不同）。这并没有利用 AVX512，例如合并掩码广播。
• Loading an xmm from GP regs 展示了 AVX512 如何使用合并屏蔽广播。但我没有费心为那里的 asm 编写内在函数。

AVX 具有指令，用于将 16 位和 32 位整数插入和提取到 __m256i 向量中：

不，它没有，_mm256_insert_epi16 和 epi32 内在函数是“假的”；它们必须用多条指令来模拟，同样_mm_set_epi32(a,b,c,d) 不是任何一条指令的内在函数。

IDK 为什么英特尔选择为 AVX1/2 而不是 AVX512 版本提供它们；也许他们后来意识到他们不应该为 AVX2 提供它们，以避免在假设这些内在函数只花费一次洗牌的情况下愚弄人们编写低效的代码。但他们无法在不破坏现有代码的情况下删除现有代码。

vpinsrd ymm_dst, ymm_src, r/m32, imm8（或 ZMM）很遗憾不存在，只有 xmm。 (https://www.felixcloutier.com/x86/pinsrb:pinsrd:pinsrq)。 XMM 版本在 __m256i 上不可用，因为它将高 128 位归零。请参阅Using ymm registers as a "memory-like" storage location（您可以使用pinsrd xmm, r/m32, imm 的传统 SSE 编码插入 YMM 的低 128 位，但在 Haswell 和 Ice Lake 上速度非常慢，因为 SSE/AVX transition penalties 在那里工作。但在 Skylake 或Ryzen。不过，编译器永远不会发出它。）

_mm256_insert_epi32 可以用 AVX2 编译到广播加载，vpblendd 可以从内存中插入一个双字。或者更糟的是，对于寄存器中的整数，编译器可能会将其vmovd 发送到 xmm reg，将其广播到 YMM，然后进行混合。 （就像我在 Move an int64_t to the high quadwords of an AVX2 __m256i vector 中展示的手工操作一样）

---

“适当的”实现取决于周围的代码。

如果要插入的元素超过 1 个，则可能需要在插入之前将它们混在一起。甚至考虑向量存储，多个标量存储，然后向量重新加载，尽管存储转发停止。或者，如果延迟关键路径通过向量而不是标量，则标量存储/向量重新加载以提供混合。如果你有很多小的标量元素，可能值得。

---

但是，对于单次插入，AVX512F 实际上有一些不错的功能：它具有 2 输入随机播放，例如 vpermt2d，您可以使用它从一个 x/y/ 的底部插入一个元素zmm 到另一个向量中的任何位置（将来自该另一个向量的所有其余目标元素作为源）。

但这里最有用的是屏蔽广播： uops.info confirms VPBROADCASTW zmm0{k1}, eax 是一条单微指令，从向量到向量（用于合并）和从屏蔽到向量。从 eax 到合并结果的

#include <immintrin.h>
#include <stdint.h>
__m512i _mm512_insert32(__m512i target, uint32_t x, const int pos)
{
    return _mm512_mask_set1_epi32(target, 1UL<<pos, x);
}

将on Godbolt 编译成这个asm：

# gcc8.3 -O3 -march=skylake-avx512
_mm512_insert32(long long __vector(8), unsigned int, int):
        mov     eax, 1
        shlx    eax, eax, esi
        kmovw   k1, eax                    # mask = 1<<pos
        vpbroadcastd    zmm0{k1}, edi
        ret

（gcc9无缘无故浪费了复制ESI的额外指令）。

使用编译时常量pos，您会得到类似mov eax,2 / kmovw k1, eax 的代码；掩蔽广播可能仍然是最好的选择。

这适用于 8、16、32 或 64 位元素。 8 和 16 当然需要 AVX512BW 用于vpbroadcastb/w 窄广播，而 32 和 64 只需要 AVX512F。

---

提取：

只需将您想要的元素移动到__m512i 的底部，您可以在其中使用_mm_cvtsi128_si32。 （在_mm512_castsi512_si128 之后）。一个有用的 shuffle 是 valignd 以按 dword 元素移动或旋转，让您有效地将任何元素移到向量的底部，而无需向量控制。 https://www.felixcloutier.com/x86/valignd:valignq

Answer 2

要完成Peter's answer，这里是 16 位和 32 位插入/提取方法的实现：

#if defined(__GNUC__)

int _mm512_cvtsi512_si32(__m512i a)
{
    __v16si b = (__v16si) a;
    return b[0];
}

#endif

__m512i _mm512_insert_epi16(__m512i target, const std::int16_t x, const int index)
{
    return _mm512_mask_set1_epi16(target, 1UL << index, x);
}
static inline __m512i _mm512_insert_epi32(__m512i target, const std::int32_t x, const int index)
{
    return _mm512_mask_set1_epi32(target, 1UL << index, x);
}

template <int index>
int _mm512_extract_epi32(__m512i target)
{
    return _mm512_cvtsi512_si32(_mm512_alignr_epi32(target, target, index));
}
template <int index>
int  _mm512_extract_epi16(__m512i target)
{
    return (_mm512_extract_epi32<index / 2>(target) >> (index % 2 ? 16 : 0)) & 0xFFFF;
}

见example。