将 32 位解压缩为 32 字节 SIMD 向量的最快方法

Question

将 32 位存储在内存中的 uint32_t 中，将每个位解压缩到 AVX 寄存器的单独字节元素的最快方法是什么？这些位可以位于其各自字节内的任何位置。

编辑：澄清一下，我的意思是位 0 到字节 0，位 1 到字节 1。显然，字节内的所有其他位都为零。目前我能做的最好的事情是 2 PSHUFB 并且每个位置都有一个掩码寄存器。

如果uint32_t是位图，那么对应的向量元素应该是0或者非0。 （也就是说，我们可以得到一个带有vpcmpeqb 的向量掩码，以对抗全零向量）。

https://software.intel.com/en-us/forums/topic/283382

Answer 1

将 32 位整数 x 的 32 位“广播”到 256 位 YMM 寄存器 z 的 32 字节或两个 128 位 XMM 寄存器 z_low 和 z_high 的 16 字节您可以执行以下操作。

使用 AVX2：

__m256i y = _mm256_set1_epi32(x);
__m256i z = _mm256_shuffle_epi8(y,mask1);
z = _mm256_and_si256(z,mask2);

如果没有 AVX2，最好使用 SSE：

__m128i y = _mm_set1_epi32(x);      
__m128i z_low  = _mm_shuffle_epi8(y,mask_low);
__m128i z_high = _mm_shuffle_epi8(y,mask_high); 
z_low  = _mm_and_si128(z_low ,mask2);
z_high = _mm_and_si128(z_high,mask2);

掩码和工作示例如下所示。如果您打算多次执行此操作，您可能应该 在主循环之外定义掩码。

#include <immintrin.h>
#include <stdio.h>

int main() {
    int x = 0x87654321;

    static const char mask1a[32] = {
        0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00,
        0x01, 0x01, 0x01, 0x01,
        0x01, 0x01, 0x01, 0x01,
        0x02, 0x02, 0x02, 0x02,
        0x02, 0x02, 0x02, 0x02,
        0x03, 0x03, 0x03, 0x03,
        0x03, 0x03, 0x03, 0x03
    };

    static const char mask2a[32] = {
        0x01, 0x02, 0x04, 0x08,
        0x10, 0x20, 0x40, 0x80,
        0x01, 0x02, 0x04, 0x08,
        0x10, 0x20, 0x40, 0x80,
        0x01, 0x02, 0x04, 0x08,
        0x10, 0x20, 0x40, 0x80,
        0x01, 0x02, 0x04, 0x08,
        0x10, 0x20, 0x40, 0x80,
    };

char out[32];

#if defined ( __AVX2__ )
    __m256i mask2 = _mm256_loadu_si256((__m256i*)mask2a);
    __m256i mask1  = _mm256_loadu_si256((__m256i*)mask1a);

    __m256i y =    _mm256_set1_epi32(x);
    __m256i z =    _mm256_shuffle_epi8(y,mask1);
    z = _mm256_and_si256(z,mask2);

    _mm256_storeu_si256((__m256i*)out,z);

#else
    __m128i mask2 = _mm_loadu_si128((__m128i*)mask2a);
    __m128i mask_low  = _mm_loadu_si128((__m128i*)&mask1a[ 0]);
    __m128i mask_high = _mm_loadu_si128((__m128i*)&mask1a[16]);    

    __m128i y = _mm_set1_epi32(x); 
    __m128i z_low  = _mm_shuffle_epi8(y,mask_low);
    __m128i z_high = _mm_shuffle_epi8(y,mask_high);
    z_low  = _mm_and_si128(z_low,mask2);
    z_high = _mm_and_si128(z_high,mask2);

    _mm_storeu_si128((__m128i*)&out[ 0],z_low);
    _mm_storeu_si128((__m128i*)&out[16],z_high);
#endif
    for(int i=0; i<8; i++) {
        for(int j=0; j<4; j++) {        
            printf("%x ", out[4*i+j]);
        }printf("\n");
    } printf("\n");
}

---

在每个向量元素中获取 0 或 -1：

它需要一个额外的步骤_mm256_cmpeq_epi8 来对抗全零。任何非零都变成0，零变成-1。如果我们不想要这种反转，请使用andnot 而不是and。它反转它的第一个操作数。

__m256i expand_bits_to_bytes(uint32_t x)
{
    __m256i xbcast = _mm256_set1_epi32(x);    // we only use the low 32bits of each lane, but this is fine with AVX2

    // Each byte gets the source byte containing the corresponding bit
    __m256i shufmask = _mm256_set_epi64x(
        0x0303030303030303, 0x0202020202020202,
        0x0101010101010101, 0x0000000000000000);
    __m256i shuf  = _mm256_shuffle_epi8(xbcast, shufmask);

    __m256i andmask  = _mm256_set1_epi64x(0x8040201008040201);  // every 8 bits -> 8 bytes, pattern repeats.
    __m256i isolated_inverted = _mm256_andnot_si256(shuf, andmask);

    // this is the extra step: compare each byte == 0 to produce 0 or -1
    return _mm256_cmpeq_epi8(isolated_inverted, _mm256_setzero_si256());
     // alternative: compare against the AND mask to get 0 or -1,
     // avoiding the need for a vector zero constant.
}

在Godbolt Compiler Explorer 上查看。

有关其他元素大小，另请参阅 is there an inverse instruction to the movemask instruction in intel avx2?。