如何实现AVX-512寄存器仅非零整数元素的条件自增操作

云上老周

2026-10-12

优化方案

核心问题

你当前实现的冗余点完全来自手动构造掩码的流程：逐个标量读取32个元素、拼接成32位掩码的操作，既引入了大量不必要的标量运算，还额外多了一轮内存访问，完全没有用到AVX512原生的掩码生成能力。

优化后实现

直接用AVX512的向量比较指令一步生成非零元素的掩码，整个处理逻辑全走向量指令，无额外标量开销，优化后的完整代码如下：

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

int main()
{
    std::vector<uint16_t> testValues{0};
    testValues.resize(65'536);
    for (size_t i{ 0 }; i < testValues.size(); i += 4)
    {
        testValues[i] = static_cast<uint16_t>(i);
    }
    auto start{ std::chrono::high_resolution_clock::now() };
    const __m512i zero_reg = _mm512_setzero_si512();
    const __m512i one_reg = _mm512_set1_epi16(1);
    for (size_t i{ 0 }; i < testValues.size(); i += 32)
    {
        // 一次性加载32个uint16_t到向量寄存器
        const __m512i vals = _mm512_loadu_epi16(&testValues[i]);
        // 直接生成非零元素的掩码：元素不等于0的位置对应掩码位为1
        const __mmask32 non_zero_mask = _mm512_cmpneq_epi16_mask(vals, zero_reg);
        // 仅对掩码为1的位置执行加1操作
        const __m512i res = _mm512_mask_add_epi16(vals, non_zero_mask, vals, one_reg);
        // 存回内存
        _mm512_storeu_epi16(&testValues[i], res);
    }
    /*
    for (auto& iter : testValues)
    {
        if (iter)
            iter += 1;
    }
    */
    auto end{ std::chrono::high_resolution_clock::now() };
    std::cout << "Summation took: " << std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count() << "ns" << std::endl;
    return 0;
}