c++ - 将 __m256i 寄存器转换为 uint64_t 位掩码，使得每个字节的值是输出中的一个设置位

Question

基本上我有一个 __m256i变量，其中每个字节代表一个需要在 uint64_t 中设置的位置.请注意，所有字节值都将 < 64。
我对如何远程有效地做到这一点感到有些茫然。
我正在考虑的一种选择是在某些情况下字节之间有很多重复项，因此类似于:

__m256i indexes = foo();

uint64_t result         = 0;
uint32_t aggregate_mask = ~0;
do {
    uint32_t idx = _mm256_extract_epi8(indexes, __tzcnt_u32(aggregate_mask));

    uint32_t idx_mask =
        _mm256_movemask_epi8(_mm256_cmpeq_epi(indexes, _mm256_set1_epi8(idx)));
    aggregate_mask ^= idx_mask;
    result |= ((1UL) << idx);
} while (aggregate_mask);

有了足够多的重复项，我相信这可能会有些效率，但我不能保证总是有足够的重复项来使这比仅遍历字节并按顺序设置更快。
我的目标是找到一些东西，这总是比感觉最坏的情况要快:

__m256i indexes = foo();
uint8_t index_arr[32];
_mm256_store_si256((__m256i *)index_arr, indexes);

uint64_t result = 0;
for (uint32_t i = 0; i < 32; ++i) {
    result |= ((1UL) << index_arr[i];
}

如果可能的话，我正在寻找一种可以在 Skylake 上运行的解决方案(没有 AVX512)。如果 AVX512 是必要的(我认为分组可能有一些半有效的东西，然后使用 _mm256_shldv_epi16 )总比没有好:)
这就是我的想法。从 Epi32 开始:

    // 32 bit
    __m256i lo_shifts = _mm256_sllv_epi32(_mm256_set1_epi32(1), indexes);
    __m256i t0 = _mm256_sub_epi32(indexes, _mm256_set1_epi32(1));
    __m256i hi_shifts = _mm256_sllv_epi32(_mm256_set1_epi32(1), t0);
    __m256i lo_shifts_lo = _mm256_shuffle_epi32(lo_shifts, 0x5555);
    __m256i hi_shifts_lo = _mm256_shuffle_epi32(hi_shifts, 0x5555);
    
    __m256i hi_shifts_hi0 = _mm256_slli_epi64(hi_shifts, 32);
    __m256i hi_shifts_hi1 = _mm256_slli_epi64(hi_shifts_lo, 32);
    __m256i all_hi_shifts = _mm256_or_epi64(hi_shifts_hi0, hi_shifts_hi1);
    
    __m256i all_lo_shifts_garbage = _mm256_or_epi64(lo_shifts_lo, lo_shifts);
    __m256i all_lo_shifts = _mm256_and_epi64(all_lo_shifts_garbage, _mm256_set1_epi64(0xffffffff));

    __m256i all_shifts = _mm256_or_epi64(all_lo_shifts, all_hi_shifts);

或者从epi64位开始:

    // 64 bit
    __m256i indexes0 = _m256_and_epi64(indexes, _mm256_set1_epi64(0xffffffff));
    __m256i indexes1 = _m256_shuffle_epi32(indexes, 0x5555);

    __m256i shifts0 = _m256_sllv_epi64(_mm256_set1_epi64x(1), indexes0);
    __m256i shifts1 = _m256_sllv_epi64(_mm256_set1_epi64x(1), indexes1);

    __m256i all_shifts = _m256_or_epi64(shifts0, shifts1);

我的猜测是来自epi64 更快。

Answer 1

关键要素是 _mm256_sllv_epi64 使用运行时可变的移位距离在 64 位 channel 内移位位。
代码需要 C++/17，仅在 VC++ 2019 中测试。
虽然不确定它是否会比标量代码快得多，但大多数指令都是 1 周期延迟，但对我来说太多了，VC++ 在关键路径上产生了大约 35 个。

// Move a single bit within 64-bit lanes
template<int index>
inline __m256i moveBit( __m256i position )
{
    static_assert( index >= 0 && index < 8 );

    // Extract index-th byte from the operand
    if constexpr( 7 == index )
    {
        // Most significant byte only needs 1 instruction to shift into position
        position = _mm256_srli_epi64( position, 64 - 8 );
    }
    else
    {
        if constexpr( index > 0 )
        {
            // Shift the operand by `index` bytes to the right.
            // On many CPUs, _mm256_srli_si256 is slightly faster than _mm256_srli_epi64
            position = _mm256_srli_si256( position, index );
        }
        const __m256i lowByte = _mm256_set1_epi64x( 0xFF );
        position = _mm256_and_si256( position, lowByte );
    }
    const __m256i one = _mm256_set1_epi64x( 1 );
    return _mm256_sllv_epi64( one, position );
}

inline uint64_t setBitsAvx2( __m256i positions )
{
    // Process each of the 8 bytes within 64-bit lanes
    const __m256i r0 = moveBit<0>( positions );
    const __m256i r1 = moveBit<1>( positions );
    const __m256i r2 = moveBit<2>( positions );
    const __m256i r3 = moveBit<3>( positions );
    const __m256i r4 = moveBit<4>( positions );
    const __m256i r5 = moveBit<5>( positions );
    const __m256i r6 = moveBit<6>( positions );
    const __m256i r7 = moveBit<7>( positions );
    // vpor instruction is very fast with 1 cycle latency,
    // however modern CPUs can issue and dispatch multiple instructions per cycle,
    // it still makes sense to try reducing dependencies.
    const __m256i r01 = _mm256_or_si256( r0, r1 );
    const __m256i r23 = _mm256_or_si256( r2, r3 );
    const __m256i r45 = _mm256_or_si256( r4, r5 );
    const __m256i r67 = _mm256_or_si256( r6, r7 );
    const __m256i r0123 = _mm256_or_si256( r01, r23 );
    const __m256i r4567 = _mm256_or_si256( r45, r67 );
    const __m256i result = _mm256_or_si256( r0123, r4567 );

    // Reduce 4 8-byte values to scalar
    const __m128i res16 = _mm_or_si128( _mm256_castsi256_si128( result ), _mm256_extracti128_si256( result, 1 ) );
    const __m128i res8 = _mm_or_si128( res16, _mm_unpackhi_epi64( res16, res16 ) );
    return (uint64_t)_mm_cvtsi128_si64( res8 );
};

inline uint64_t setBitsScalar( __m256i positions )
{
    alignas( 32 ) std::array<uint8_t, 32> index_arr;
    _mm256_store_si256( ( __m256i * )index_arr.data(), positions );

    uint64_t result = 0;
    for( uint32_t i = 0; i < 32; i++ )
        result |= ( ( 1ull ) << index_arr[ i ] );
    return result;
}

static void testShuffleBits()
{
    const __m128i src16 = _mm_setr_epi8( 0, 0, 0, 0, 1, 4, 5, 10, 11, 12, 13, 14, 15, 16, 17, 31 );
    const __m256i src32 = _mm256_setr_m128i( src16, _mm_setzero_si128() );
    printf( "AVX2: %" PRIx64 "\n", setBitsAvx2( src32 ) );
    printf( "Scalar: %" PRIx64 "\n", setBitsScalar( src32 ) );
}