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

Posted 2023-02-16

【中文标题】将 __m256i 寄存器转换为 uint64_t 位掩码，以便每个字节值处的值是输出中的设置位

【英文标题】：Convert __m256i register to uint64_t bitmask such that that value at each byte's value is a set bit in the output

【发布时间】：2020-09-05 18:14:23

【问题描述】：

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

【问题讨论】：

元素是否已排序，以防万一？以某种方式使用pshufb 将i 替换为8 位块中的1<<i，8 次8 宽度？此外，_mm256_extract_epi8 不是真正的硬件指令，vpextrb 只需要立即数。你不妨把它写成循环外的一个向量存储+在里面索引一个 char 数组，因为这是它可以编译的唯一理智的方式，你想排除疯狂的方式。

也许将元素加宽到 8x 32 位元素以处理移位计数 0..31 和 4x vpsllvd，或者一起，并且还分别计数 32..63？然后将它们与vpunpckl/hdq 一起解压缩为 64 位位，或一起，并对 1 个 qword 元素进行 OR 缩减。或者只使用 8x vpsllvq，每个处理 4 个索引元素。 （vpmovzxbq 设置计数）。不确定这是否能击败部分或完全展开的bts 循环，其中有几个累加器来隐藏 ALU 延迟，最后有一个 OR。无论是延迟还是吞吐量。

我喜欢 vpsllvd 的想法，实际上我之前正在打包到 epi8 以获取其他逻辑，但这可能是一个中间步骤。

如果我们必须为每个元素移动 2 种不同的方式以防它是 =32，我不确定它是否有助于执行 32 位。可能只是扩大到 64 位更有意义，除非有某种方法可以对 idx 元素进行分组，我们知道其中一些组都是 =32。就像他们被分类一样。 （请注意，x86 SIMD 移位使计数饱和，而不是隐式屏蔽它，这意味着高半移位需要屏蔽或减去输入，但低半移位将自行归零。上半部分仍然需要工作正确。）

使用可能的方法进行编辑。

【参考方案1】：

关键要素是_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 ) );

【讨论】：

quick-bench.com/q/X2GHxuaaPcjCCvk8Yi9UfseL9zw original、this 和 andrey's。与 GCC 的结果几乎相同

我的回答是不正确的，所以不要费心去做。很抱歉造成混乱。

@Sopel 您的链接中的 clang 程序集输出很有趣。它用带有恒定洗牌掩码的 vpshufb 替换了我的 6 个固定班次：godbolt.org/z/jxsz93 虽然从技术上讲指令较少，但我不确定这是一件好事，它交换了一些非常快速的指令以供 RAM 参考，还有代码由于这些常数，大小要大得多。我想我更喜欢 gcc 输出，它完成了我要求它做的事情。

@Sopel 还要感谢那个基准测试，我没想到我的版本会提高 3.7 倍，我以为可能会是 2-3。