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

问题描述

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

例如：

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 指令。

---

附言

我的代码被编译器变成了 5 条指令（启用了优化），但也许您应该使用组指令，这将提供额外的速度增益

---

P.P.S.最少的代码：

constexpt block_size = 16;
int arr[] = {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) 通常就足够了。