Como executar OR bit a bit em vetores grandes de u64 da maneira mais eficiente?

Como um experimento mental, estou tentando descobrir qual é a maneira mais eficiente de executar o bit a bit ou em máscaras de bits de tamanho relativamente grande (~1M elementos u64) em um único thread.

Estou executando-o no meu laptop com CPU i7-8750H a 2,2 GHz, que suporta AVX2, mas não AVX-512. Como mencionei antes, estou tentando descobrir o que posso extrair de um único thread. Além disso, tenho quase certeza de que compactar minhas máscaras de bits em algo como bitmaps rugindo provavelmente produzirá acelerações adicionais, pois eu teria que ingerir menos bytes na CPU. Mas gostaria de entender primeiro se há algo obviamente errado com o que estou fazendo na minha implementação atual bastante simples. Estou mirando x86-64.

Tentei estimar qual é o teto de velocidade dessa operação. Vamos começar com a suposição obviamente muito errada de que todos os dados já estão nos registradores. _mm256_or_si256pode executar uma operação bit a bit ou em 4x u64s. Vou assumir que posso executar apenas uma operação desse tipo por ciclo, o que provavelmente também está errado e posso fazer 2-3. Vamos usar u64vetores de 1M s para nossos testes. Isso nos dá 1M / 4 per cycle = 250K cycles. Sob carga, a taxa de clock da minha CPU fica em ~3GHz. Isso nos dá 250K cycles / 3Ghz ~ 0.00008(3)sec ~ 83µs. Esse é o nosso teto ingenuamente rápido que usaremos como valor de referência.

Vamos transformá-lo em código. Usei o ChatGPT para gerar um grande pedaço dele, desculpe antecipadamente, não sou um desenvolvedor C++ e sou relativamente ignorante em desenvolvimento de baixo nível.

#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;
}

Eu o construo com g++ -O3 -mavx2 -o experiment main.cpp. Quando estou executando, leva cerca de 9 segundos para executar 10.000 iterações, o que é 900 µs por iteração. Isso é ~10+ vezes mais lento do que nossos cálculos ingênuos de trás do envelope. Tentei desenrolar os loops, mas não me deu nenhum resultado de desempenho tangível. Talvez eu tenha feito errado de alguma forma.

Vamos contrastar com o código que evita a leitura da memória:

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;

Foram necessários cerca de 9 segundos para uma iteração de 100K para o mesmo 1M x u64vetor, o que equivale a 90 µs por iteração, muito próximo do nosso valor de referência ingênuo de 83 µs.

Então, eu me pergunto o que posso fazer para aumentar ainda mais o desempenho do meu código que lê da memória? Há algo que eu possa fazer?