如何将位对从 uint64_t 解压缩为 __m256i？

Question

考虑

uint64_t

，其中每个连续的 2 位都是一个数字：

b00

代表

0

，

b01

代表

1

，

b11

代表

-1

并且

b10

未使用（从未发生，假设不进行处理）为了它）。 如何将这样的

uint64_t

解压缩为

__m256i

以便原始数字中的位对成为向量变量中的有符号字节？我需要将

b00

解压为

b00000000

，将

b01

解压为

b00000001

，将

b11

解压为

b11111111

，同时保留顺序。

到目前为止我已经尝试过： 受这个答案的启发，我只提出了逆过程，其中

__m256i

的字节被打包成

uint64_t

的位对：

const __m256i vBits0 = _mm256_slli_epi16(target, 7);
const int bits1 = _mm256_movemask_epi8(target);
const int bits0 = _mm256_movemask_epi8(vBits0);
const uint64_t withHigher = _pdep_u64(bits1, 0xAAAAAAAAAAAAAAAA);
const uint64_t withLower =  _pdep_u64(bits0, 0x5555555555555555);
return withHigher | withLower;

我可以访问 AVX2，但 AVX512 也会感兴趣。

Answer 1

当您的 BMI2 较快时，这是一种简单的方法。

// Expand signed 2-bit fields into signed bytes in the AVX2 vector, with BMI2
__m256i expand2_bmi( uint64_t bitmap )
{
    // Use BMI to extend 2 bit fields into bytes,
    // each instruction transforms 16 bits into 8 bytes
    constexpr uint64_t dep2 = 0x0303030303030303ull;
    uint64_t i0 = _pdep_u64( bitmap, dep2 );
    uint64_t i1 = _pdep_u64( bitmap >> 16, dep2 );
    uint64_t i2 = _pdep_u64( bitmap >> 32, dep2 );
    uint64_t i3 = _pdep_u64( bitmap >> 48, dep2 );

    // Assemble AVX2 vector with these numbers, probably 5 instructions
    __m256i v = _mm256_setr_epi64x( i0, i1, i2, i3 );

    // Lookup table to sign extend
    const __m256i perm = _mm256_set1_epi32( 0xFFFE0100 );
    v = _mm256_shuffle_epi8( perm, v );
    return v;
}

这是另一个不需要 BMI2 的想法。

// Expand signed 2-bit fields into signed bytes in the AVX2 vector
__m256i expand2_vec( uint64_t bitmap )
{
    // Move to SSE vector
    __m128i v16 = _mm_cvtsi64_si128( (int64_t)bitmap );
    // Zero-extend bytes into uint32_t lanes
    __m256i v = _mm256_cvtepu8_epi32( v16 );

    // Split pairs of 4-bit fields into uint16_t lanes
    __m256i tmp = _mm256_slli_epi32( v, 12 );
    v = _mm256_or_si256( v, tmp );

    // Split pairs of 2-bit fields into uint8_t lanes
    tmp = _mm256_slli_epi16( v, 6 );
    v = _mm256_or_si256( v, tmp );

    // The above 4 instructions left garbage in bits 2-7 of most bytes
    // Sadly, the pollution affects bits 7, 15 and 23 which cause vpshufb to output zeros
    // That's why need yet another constant vector
    v = _mm256_and_si256( v, _mm256_set1_epi8( 0b11 ) );

    // Lookup table to sign extend
    const __m256i perm = _mm256_set1_epi32( 0xFFFE0100 );
    v = _mm256_shuffle_epi8( perm, v );
    return v;
}

测试很少。

我不知道哪一个更快。我希望它取决于用例、周围的代码，甚至可能取决于编译器开关。

BMI2版本有更多指令，但许多部分会并行运行。仅 AVX2 版本的指令较少，但整个函数中存在单一数据依赖链，还有两个从内存加载的向量常量，而不是单个向量常量。