计算数组中的整数,其中设置位是给定掩码的子集

Count integers in an array where the set bits are a subset of a given mask

给定一个掩码和一个值,如果值中的所有位都落入掩码中,则掩码覆盖该值。

例如:

mask:  0b011010
value: 0b010010

true

mask:  0b011010
value: 0b010110

false

对于int arr[arr_size],我需要计算给定掩码覆盖了多少数组元素。

我的代码:

int count = 0;

for (int index = 0; index < arr_size; index++)
{
    // note: always true because of operator precedence 
    if (arr[index] | mask == mask)
        count++;
}

或(较慢)

int count = 0;

for (int index = 0; index < arr_size; index++)
{
    // note: also broken because of operator precedence, but not always true
    if (arr[index] & mask == arr[index])
        count++;
}

我的程序经常需要计算这样的数组元素的个数。

你能告诉我是否有任何方法可以加快这种计算速度吗?例如使用 SSE、AVX 指令。


P.S.

我的代码被编译器转换为 5 条指令(启用了优化),但也许您应该使用组指令,这将提供额外的速度增益


P.P.S。最小代码:

constexpt block_size = 16;
int arr[] = {rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), rand(), }; // random values for example
int mask = rand(); 

int count = 0;

for (__int64 cycles = 0; cycles < 0xFFFFFFFF; cycles ++)
{
    for (int index = 0; index < block_size; index ++)
    {
        if (arr[index] | mask == mask)
            count++;
    }
}

AVX1 对此毫无用处,它只处理浮点数,而且那里有整数。这是您代码的 AVX2 端口。

#include <immintrin.h>
#include <assert.h>

// Compute horizontal sum of all int32 lanes of the vector
inline int horizontalSum( const __m256i v )
{
    // Add 8 scalars into 4, by extracting lower/higher pieces from the source vector
    const __m128i r4 = _mm_add_epi32( _mm256_castsi256_si128( v ),
        _mm256_extracti128_si256( v, 1 ) );
    // Add 4 scalars into 2, using shift by 8 bytes
    const __m128i r2 = _mm_add_epi32( r4, _mm_srli_si128( r4, 8 ) );
    // Add 2 lowest scalars, using shift by 4 bytes
    const __m128i r1 = _mm_add_epi32( r2, _mm_srli_si128( r2, 4 ) );
    // Move the lowest int32 into scalar register
    return _mm_cvtsi128_si32( r1 );
}

int countMaskedAvx2( const int* src, size_t length, int mask )
{
    assert( length <= INT_MAX );

    const int* const srcEndAligned = src + ( length / 8 ) * 8;
    const int* const srcEnd = src + length;
    // Broadcast mask into all 8 lanes of a vector register
    const __m256i maskVector = _mm256_set1_epi32( mask );
    // Initialize counters with zero vector
    __m256i counters = _mm256_setzero_si256();

    // Handle vast majority of data with AVX2 code
    while( src < srcEndAligned )
    {
        // Load 8 scalars from the array
        __m256i v = _mm256_loadu_si256( ( const __m256i * )src );
        src += 8;
        // val | mask
        v = _mm256_or_si256( v, maskVector );
        // ( val | mask ) == mask
        v = _mm256_cmpeq_epi32( v, maskVector );
        // Increment the counters.
        // The result of that comparison instruction is 0 for false, or 0xFFFFFFFF for true.
        // Conveniently, 0xFFFFFFFF bit pattern equals to -1 signed integer.
        counters = _mm256_sub_epi32( counters, v );
    }

    // Compute horizontal sum of the counters vector
    int counter = horizontalSum( counters );

    // Handle the final 1-7 scalars of the source array
    while( src < srcEnd )
    {
        const int v = mask | ( *src );
        src++;
        const int add = ( v == mask ) ? 1 : 0;
        counter += add;
    }
    return counter;
}

如果您没有 AVX2,同样的方法也适用于 SSE2。您可以依赖它的支持,所有 64 位 x86 处理器都需要至少支持 SSE1 和 SSE2。与 AVX2 相比,SSE2 版本将 运行 指令数量增加一倍,但它仍然会使用 128 位宽的内存加载和 4 位宽的算术指令。即使使用 SSE2,与您的原始代码相比,我也希望有可衡量的性能改进。

更新: 为获得上述代码的最佳性能,请确保您的输入向量按 32 字节 (AVX2) 或 16 字节 (SSE2) 对齐。如果您有 C++/11 或更新版本,alignas(32) 通常就足够了。