使用 AVX-512 或 AVX-2 在大数据上计算 1 位（人口计数）

Question

我有一大块内存，比如 256 KiB 或更长。我想计算整个块中 1 的位数，或者换句话说：将所有字节的 "population count" 值相加。

我知道 AVX-512 有一个 VPOPCNTDQ instruction，它计算 512 位向量中每个连续 64 位中 1 位的数量，IIANM 应该可以在每个周期发出其中一个（如果有合适的 SIMD 向量寄存器可用）——但我没有任何编写 SIMD 代码的经验（我更像是一个 GPU 专家）。另外，我不是 100% 确定编译器对 AVX-512 目标的支持。

在大多数 CPU 上，仍然不（完全）支持 AVX-512；但 AVX-2 是广泛可用的。我一直无法找到类似于 VPOPCNTDQ 的小于 512 位的矢量化指令，所以即使在理论上我也不确定如何使用支持 AVX-2 的 CPU 快速计算位；也许存在这样的东西，我只是不知何故错过了它？

无论如何，对于两个指令集的每一个，我都希望有一个简短的 C/C++ 函数——使用一些内部包装库或内联汇编。签名是

uint64_t count_bits(void* ptr, size_t size);

备注：

与 How to quickly count bits into separate bins in a series of ints on Sandy Bridge? 相关，但不是骗子。
如果重要的话，我们可以假设输入对齐良好。
忘记多核或套接字，我想要单核（单线程）的代码。

Answer 1

AVX-2

@HadiBreis 的评论链接到 article on fast population-count with SSSE3, by Wojciech Muła; the article links to this GitHub repository; and the repository has 以下 AVX-2 实现。它基于向量化查找指令，并使用 16 值查找 table 来计算半字节的位数。

#   include <immintrin.h>
#   include <x86intrin.h>

std::uint64_t popcnt_AVX2_lookup(const uint8_t* data, const size_t n) {

    size_t i = 0;

    const __m256i lookup = _mm256_setr_epi8(
        /* 0 */ 0, /* 1 */ 1, /* 2 */ 1, /* 3 */ 2,
        /* 4 */ 1, /* 5 */ 2, /* 6 */ 2, /* 7 */ 3,
        /* 8 */ 1, /* 9 */ 2, /* a */ 2, /* b */ 3,
        /* c */ 2, /* d */ 3, /* e */ 3, /* f */ 4,

        /* 0 */ 0, /* 1 */ 1, /* 2 */ 1, /* 3 */ 2,
        /* 4 */ 1, /* 5 */ 2, /* 6 */ 2, /* 7 */ 3,
        /* 8 */ 1, /* 9 */ 2, /* a */ 2, /* b */ 3,
        /* c */ 2, /* d */ 3, /* e */ 3, /* f */ 4
    );

    const __m256i low_mask = _mm256_set1_epi8(0x0f);

    __m256i acc = _mm256_setzero_si256();

#define ITER { \
        const __m256i vec = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i)); \
        const __m256i lo  = _mm256_and_si256(vec, low_mask); \
        const __m256i hi  = _mm256_and_si256(_mm256_srli_epi16(vec, 4), low_mask); \
        const __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo); \
        const __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi); \
        local = _mm256_add_epi8(local, popcnt1); \
        local = _mm256_add_epi8(local, popcnt2); \
        i += 32; \
    }

    while (i + 8*32 <= n) {
        __m256i local = _mm256_setzero_si256();
        ITER ITER ITER ITER
        ITER ITER ITER ITER
        acc = _mm256_add_epi64(acc, _mm256_sad_epu8(local, _mm256_setzero_si256()));
    }

    __m256i local = _mm256_setzero_si256();

    while (i + 32 <= n) {
        ITER;
    }

    acc = _mm256_add_epi64(acc, _mm256_sad_epu8(local, _mm256_setzero_si256()));

#undef ITER

    uint64_t result = 0;

    result += static_cast<uint64_t>(_mm256_extract_epi64(acc, 0));
    result += static_cast<uint64_t>(_mm256_extract_epi64(acc, 1));
    result += static_cast<uint64_t>(_mm256_extract_epi64(acc, 2));
    result += static_cast<uint64_t>(_mm256_extract_epi64(acc, 3));

    for (/**/; i < n; i++) {
        result += lookup8bit[data[i]];
    }

    return result;
}

AVX-512

同一个存储库还有一个基于 VPOPCNT 的 AVX-512 实现。在列出它的代码之前，这里是简化且更易读的伪代码：

For every consecutive sequence of 64 bytes:

Load the sequence into a SIMD register with 64x8 = 512 bits

Perform 8 parallel population counts of 64 bits each on that register

Add the 8 population-count results in parallel, into an "accumulator" register holding 8 sums

Sum up the 8 values in the accumulator

If there's a tail of less than 64 bytes, count the bits there in some simpler way

Return the main sum plus the tail sum

现在开始真正的交易：

#   include <immintrin.h>
#   include <x86intrin.h>

uint64_t avx512_vpopcnt(const uint8_t* data, const size_t size) {
    
    const size_t chunks = size / 64;

    uint8_t* ptr = const_cast<uint8_t*>(data);
    const uint8_t* end = ptr + size;

    // count using AVX512 registers
    __m512i accumulator = _mm512_setzero_si512();
    for (size_t i=0; i < chunks; i++, ptr += 64) {
        
        // Note: a short chain of dependencies, likely unrolling will be needed.
        const __m512i v = _mm512_loadu_si512((const __m512i*)ptr);
        const __m512i p = _mm512_popcnt_epi64(v);

        accumulator = _mm512_add_epi64(accumulator, p);
    }

    // horizontal sum of a register
    uint64_t tmp[8] __attribute__((aligned(64)));
    _mm512_store_si512((__m512i*)tmp, accumulator);

    uint64_t total = 0;
    for (size_t i=0; i < 8; i++) {
        total += tmp[i];
    }

    // popcount the tail
    while (ptr + 8 < end) {
        total += _mm_popcnt_u64(*reinterpret_cast<const uint64_t*>(ptr));
        ptr += 8;
    }

    while (ptr < end) {
        total += lookup8bit[*ptr++];
    }

    return total;
}

lookup8bit 是针对字节而不是位的 popcnt 查找 table，定义为 here。 编辑： 正如评论者所说，在最后使用 8 位查找 table 不是一个好主意，可以改进。

Answer 2

Wojciech Muła's big-array popcnt functions 除了标量清理循环外看起来是最佳的。（有关主循环的详细信息，请参阅@einpoklum 的回答）。

最后只使用几次的 256 条目 LUT 可能会缓存未命中，并且即使缓存很热，也不是超过 1 个字节的最佳选择。我相信所有 AVX2 CPU 都有硬件 popcnt，我们可以很容易地隔离出最后最多 8 个尚未被计算的字节，以便为单个 popcnt.[=14= 设置我们]

与 SIMD 算法一样，在缓冲区的 last 字节结束的全角加载通常效果很好。但与矢量寄存器不同，完整整数寄存器的可变计数移位很便宜（尤其是 BMI2）。 Popcnt 不在乎位在哪里，所以我们可以只使用移位而不需要构造 AND 掩码或其他任何东西。

// untested
// ptr points at the first byte that hasn't been counted yet
uint64_t final_bytes = reinterpret_cast<const uint64_t*>(end)[-1] >> (8*(end-ptr));
total += _mm_popcnt_u64( final_bytes );
// Careful, this could read outside a small buffer.

或者更好的是，使用更复杂的逻辑来避免页面交叉。例如，这可以避免页面开始处的 6 字节缓冲区的页面交叉。

使用 AVX-512 或 AVX-2 在大数据上计算 1 位（人口计数）

Counting 1 bits (population count) on large data using AVX-512 or AVX-2

assembly

bitcount

avx2

avx512

population-count

AVX-2

AVX-512