如何以最高效的方式对 u64 的大向量运行按位或运算？

作为一个思想实验，我试图找出在单个线程上运行按位或相对较大尺寸（~1M u64 元素）的位掩码的最有效方法是什么。

我在配备 2.2GHz i7-8750H CPU 的笔记本电脑上运行它，它支持 AVX2，但不支持 AVX-512。正如我之前提到的，我正在尝试弄清楚我能从单个线程中挤出什么。此外，我有点确定将我的位掩码压缩成咆哮位图之类的东西可能会产生额外的速度提升，因为我必须将更少的字节摄入 CPU。但我想首先了解我目前相当简单的实现中所做的是否明显有问题。我的目标是 x86-64。

我试图估计这个操作的速度上限。让我们从显然非常错误的假设开始，即所有数据都已在寄存器中。_mm256_or_si256可以对 4x 执行按位或操作u64。我假设我每个周期只能执行一次这样的操作，这也可能错了，我可以做 2-3 次。我们将使用 1Mu64向量进行测试。这给了我们1M / 4 per cycle = 250K cycles。在负载下，我的 CPU 时钟频率保持在 ~3GHz。这给了我们250K cycles / 3Ghz ~ 0.00008(3)sec ~ 83µs。这是我们将用作参考值的天真的快速上限。

让我们把它变成代码。我使用 ChatGPT 生成了大部分代码，提前说声抱歉，我不是 C++ 开发人员，对底层开发相对无知。

#include <iostream>
#include <vector>
#include <random>
#include <chrono>
#include <cstdint>
#include <immintrin.h>

std::vector<uint64_t> generate_random_u64s(size_t amount) {
    std::vector<uint64_t> random_u64s;
    random_u64s.reserve(amount);
    std::random_device rd; // Obtain a random number from hardware
    std::mt19937_64 eng(rd()); // Seed the generator
    std::uniform_int_distribution<uint64_t> distr; // Define the range

    for (size_t i = 0; i < amount; ++i) {
        random_u64s.push_back(distr(eng));
    }
    return random_u64s;
}

void avx_bitwise_or(
    const std::vector<uint64_t>& vector1, 
    const std::vector<uint64_t>& vector2, 
    std::vector<uint64_t>& result
) {
    size_t size = vector1.size();
    
    // Ensure result has enough space
    if (result.size() != size) {
        result.resize(size);
    }

    size_t i = 0;
    for (; i + 4 <= size; i += 4) {
        // Prefetch data into CPU cache
        _mm_prefetch(reinterpret_cast<const char*>(&vector1[i + 256]), _MM_HINT_T0);
        _mm_prefetch(reinterpret_cast<const char*>(&vector2[i + 256]), _MM_HINT_T0);

        // Load vectors using AVX
        __m256i vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vector1[i]));
        __m256i vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vector2[i]));

        // Perform bitwise OR operation
        __m256i vec_result = _mm256_or_si256(vec1, vec2);

        // Use non-temporal store to write result back to memory bypassing CPU cache
        _mm256_stream_si256(reinterpret_cast<__m256i*>(&result[i]), vec_result);
    }

    // Handle remaining elements that don't fit in AVX registers
    for (; i < size; ++i) {
        result[i] = vector1[i] | vector2[i];
    }
}

int main() {
    std::cout << "Starting" << std::endl;
    const size_t size = 1'000'000;
    auto vector1 = generate_random_u64s(size);
    auto vector2 = generate_random_u64s(size);

    auto result = std::vector<uint64_t>(size);

    auto start = std::chrono::high_resolution_clock::now();
    const int repetitions = 10000;

    for (int i = 0; i < repetitions; ++i) {
        avx_bitwise_or(vector1, vector2, result);
    }

    auto duration = std::chrono::high_resolution_clock::now() - start;

    uint32_t popcnt = 0;
    for (const auto& x : result) {
        popcnt += __builtin_popcountll(x); // Count the number of set bits (1s)
    }

    std::cout << "Popcnt is: " << popcnt << std::endl;
    std::cout << "Time elapsed is: " << std::chrono::duration_cast<std::chrono::milliseconds>(duration).count() << " ms" << std::endl;

    return 0;
}

我使用 构建它g++ -O3 -mavx2 -o experiment main.cpp。当我运行它时，运行 10,000 次迭代大约需要 9 秒，即每次迭代 900µs。这比我们简单的粗略计算慢了约 10 倍以上。我尝试展开循环，但它没有给我任何切实的性能结果。也许我只是做错了什么。

我们将其与避免从内存中读取的代码进行对比：

uint64_t single_register_test (
    const std::vector<uint64_t>& vector1, 
    const std::vector<uint64_t>& vector2
) {
    size_t size = vector1.size();

    // Load vectors using AVX, use only first 4 elements and ignore everything else
    __m256i vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vector1[0]));
    __m256i vec_result = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vector2[0]));

    // Perform bitwise OR operation on the same data over and over again
    for (size_t i = 0; i + 4 <= size; i += 4) {
        vec_result = _mm256_or_si256(vec1, vec_result);
    }

    // Get first u64 from the result
    auto result = _mm256_extract_epi64(vec_result, 0);

    return result;
}

// ...........

    auto start = std::chrono::high_resolution_clock::now();
    const int repetitions = 100'000;

    uint32_t popcnt_single = 0;
    for (int i = 0; i < repetitions; ++i) {
        auto x = single_register_test(vector1, vector2);
        popcnt_single += __builtin_popcountll(x);
    }
    std::cout << "Popcnt single is: " << popcnt_single << std::endl;
    auto duration = std::chrono::high_resolution_clock::now() - start;

对于同一个向量，100K 次迭代花费约 9 秒1M x u64，即每次迭代 90µs，非常接近我们的 83µs 简单参考值。

所以，我想知道我能做些什么来进一步提高从内存读取的代码的性能？我能做些什么吗？