首页 > 解决方案 > 将 __m256i 寄存器转换为 uint64_t 位掩码,以便每个字节值处的值是输出中的设置位

问题描述

基本上我有一个__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(wo 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 的速度更快。

标签: c++simdavxmicro-optimizationavx2

解决方案


关键因素是_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, 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 ) );
}

推荐阅读