问题描述
基本上我有一个__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 ) );
}