将__m256i寄存器转换为uint64_t位掩码，以使每个字节值处的值成为输出中的设置位

问题描述

基本上我有一个__m256i变量，其中每个字节代表一个需要在uint64_t中设置的位置。请注意，所有字节值均为

我什至不知道如何远程高效地执行此操作。

我正在考虑的一个选项是，在某些情况下，字节之间有很多重复项，因此类似于：

__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（w.o 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的速度更快。

解决方法

关键要素是_mm256_sllv_epi64，使用运行时可变移位距离在64位通道内移位位。

代码需要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,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 ) );
}