使用 AVX-512 或 AVX-2 對大數據計算 1 位（人口計數）

Question

我有很長的內存塊，比如 256 KiB 或更長。 我想計算整個塊中 1 位的數量，或者換句話說：將所有字節的“人口計數”值相加。

我知道 AVX-512 有一個VPOPCNTDQ 指令，它計算 512 位向量內每個連續 64 位中 1 位的數量，而 IIANM 應該可以在每個周期發出這些指令之一（如果合適的 SIMD 向量寄存器是可用） - 但我沒有任何編寫 SIMD 代碼的經驗（我更像是一個 GPU 人）。 另外，我不是 100% 確定編譯器支持 AVX-512 目標。

在大多數 CPU 上，仍然不（完全）支持 AVX-512； 但 AVX-2 是廣泛可用的。 我找不到類似於 VPOPCNTDQ 的小於 512 位的向量化指令，所以即使從理論上講，我也不確定如何使用支持 AVX-2 的 CPU 快速計算位； 也許像這樣的東西存在而我只是以某種方式錯過了它？

無論如何，對於兩個指令集中的每一個，我都希望有一個簡短的 C/C++ 函數 - 要么使用一些 intristics-wrapper 庫，要么使用內聯匯編。 簽名是

uint64_t count_bits(void* ptr, size_t size);

筆記：

• 與如何在 Sandy Bridge 上的一系列整數中快速將位計數到單獨的 bin 中有關？ ，但不是騙子。
• 如果這很重要，我們可以假設輸入是對齊的。
• 忘記多個內核或套接字，我想要一個（單個線程上的）內核的代碼。

Answer 1

AVX-2

@HadiBreis 的評論鏈接到 Wojciech Muła 的一篇關於使用 SSSE3 進行快速人口計數的文章； 文章鏈接到此 GitHub 存儲庫； 並且存儲庫具有以下 AVX-2 實現。 它基於矢量化查找指令，並使用 16 值查找表來獲取半字節的位數。

#   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 實現。 在列出它的代碼之前，這里是簡化且更易讀的偽代碼：

• 對於每個連續的 64 字節序列：

  • 將序列加載到 64x8 = 512 位的 SIMD 寄存器中
  • 在該寄存器上執行 8 次 64 位的並行填充計數
  • 將 8 個人口計數結果並行添加到保存 8 個和的“累加器”寄存器中

• 將累加器中的 8 個值相加

• 如果尾部少於 64 個字節，請以更簡單的方式計算那里的位數

• 返回主和加上尾和

現在是真正的交易：

#   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 查找表，並在此處定義。 編輯：正如評論者所指出的，在最后使用 8 位查找表不是一個好主意，可以改進。

Answer 2

除了標量清理循環之外， Wojciech Muła 的大數組 popcnt 函數看起來是最佳的。 （有關主循環的詳細信息，請參閱@einpoklum 的回答）。

最后只使用幾次的 256 個條目的 LUT 可能會緩存未命中，即使緩存很熱，也不是超過 1 個字節的最佳選擇。 我相信所有 AVX2 CPU 都有硬件popcnt ，我們可以輕松地隔離最后最多 8 個尚未計數的字節，以便為我們設置單個popcnt 。

與 SIMD 算法一樣，在緩沖區的最后一個字節處進行全寬加載通常效果很好。 但與向量寄存器不同，完整整數寄存器的可變計數移位很便宜（尤其是使用 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 字節緩沖區的頁面交叉。