AVX2下同一SIMD寄存器串行累加前缀和 实现并行数组指针分配

实现方案 #

你已经完成了Nums为0对应位置Capacities清零的逻辑，下面直接补充else分支的AVX2实现，完全符合你要求的4个lane计算规则，仅用少量AVX2指令即可完成：

完整循环代码 #

#include <immintrin.h>
#include <stdint.h>

// 你已定义的全零常量 __m256i ZEROES256 = _mm256_setzero_si256();

void process(int n, int16_t* Nums, int16_t* Capacities, int** Data, int* freePointer) {
    int i = 0;
    // 每次处理4个元素（对应4个64位指针）
    for (; i <= n - 4; i += 4) {
        const __m256i nums = _mm256_loadu_si256((__m256i*)&Nums[i]);
        const __m256i bZeroes = _mm256_cmpeq_epi16(nums, ZEROES256);
        const __m256i capacities = _mm256_loadu_si256((__m256i*)&Capacities[i]);
        const __m256i zeroedCapacities = _mm256_andnot_si256(bZeroes, capacities);
        _mm256_storeu_si256((__m256i*)&Capacities[i], zeroedCapacities);

        // 1. 提取前4个int16_t容量，零扩展为4个64位整数适配指针运算
        const __m128i cap_16 = _mm256_castsi256_si128(zeroedCapacities);
        const __m256i cap_64 = _mm256_cvtepu16_epi64(cap_16);

        // 2. 2条指令完成4个64位元素的前缀和：得到 [c0, c0+c1, c0+c1+c2, 总容量]
        __m256i prefix = cap_64;
        prefix = _mm256_add_epi64(prefix, _mm256_slli_si256(prefix, 8));
        prefix = _mm256_add_epi64(prefix, _mm256_slli_si256(prefix, 16));

        // 3. 右移8字节得到你需要的偏移向量：[0, c0, c0+c1, c0+c1+c2]
        const __m256i offsets = _mm256_srli_si256(prefix, 8);

        // 4. 广播当前freePointer到4个lane，计算所有Data[i]的指针值
        const __m256i base_ptr = _mm256_set1_epi64x((intptr_t)freePointer);
        const __m256i data_ptrs = _mm256_add_epi64(base_ptr, offsets);

        // 5. 存储结果并更新freePointer
        _mm256_storeu_si256((__m256i*)&Data[i], data_ptrs);
        freePointer += (size_t)_mm256_extract_epi64(prefix, 3);
    }
    // 处理n不是4倍数的剩余元素，标量兜底
    for (; i < n; i++) {
        if (Nums[i] == 0)
            Capacities[i] = 0;
        else {
            Data[i] = freePointer;
            freePointer += Capacities[i];
        }
    }
}

注意说明 #

• 如果Capacities可能存在负值，把_mm256_cvtepu16_epi64替换为_mm256_cvtepi16_epi64做符号扩展即可
• 如果数组是16/32字节对齐的，把loadu/storeu替换为对齐版本load/store可获得更高性能

内容的提问来源于stack exchange，提问作者GlassBeaver