基于SIMD 128/256/512实现std::countr_zero及全零检测优化咨询

针对海量512位数组的场景，我们可以利用SSE2、AVX2、AVX-512等SIMD指令集大幅优化最低连续零比特数的统计效率，同时针对半数以上数据为全零的特点，优先做全零快速检测以减少无效计算。

核心优化思路 #

• 全零优先检测：用SIMD指令一次性判断整个128/256/512位数据是否全零，命中则直接返回总比特数，跳过后续统计步骤。
• 并行统计：利用SIMD指令对所有64位元素的末尾零比特数进行并行计算，避免逐个循环处理的开销。
• 快速定位第一个非零元素：通过SIMD掩码快速找到第一个非零的64位元素，直接提取对应的统计结果。

AVX-512实现（最优512位支持） #

AVX-512原生支持512位寄存器，能一次性处理所有8个64位元素，效率最高：

#include <immintrin.h>

int CountRZero512_AVX512(const uint64_t* ptr) {
    // 加载512位数据到寄存器
    const __m512i data = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(ptr));
    
    // 快速检测全零：掩码为0表示所有位都是0
    if (_mm512_test_epi64_mask(data, data) == 0) {
        return 512;
    }
    
    // 并行统计每个64位元素的末尾零比特数
    const __m512i cnts = _mm512_cttz_epi64(data);
    // 生成每个元素对应的基础偏移量（0, 64, 128, ..., 448）
    const __m512i offsets = _mm512_setr_epi64(0, 64, 128, 192, 256, 320, 384, 448);
    // 计算每个元素对应的总连续零比特数
    const __m512i total = _mm512_add_epi64(cnts, offsets);
    
    // 获取非零元素的掩码，找到第一个非零元素的索引
    const __mmask8 non_zero_mask = _mm512_test_epi64_mask(data, data);
    const int first_non_zero_idx = _tzcnt_u32(non_zero_mask);
    
    // 提取对应位置的结果并返回
    return static_cast<int>(_mm512_extract_epi64(total, first_non_zero_idx));
}

AVX2实现（兼容多数现代CPU） #

AVX2支持256位寄存器，需将512位数据拆分为两个256位块处理：

#include <immintrin.h>

int CountRZero512_AVX2(const uint64_t* ptr) {
    // 加载前后两个256位块
    const __m256i low_chunk = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
    const __m256i high_chunk = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr + 4));
    
    // 检测是否全零
    if (_mm256_testz_si256(low_chunk, low_chunk) && _mm256_testz_si256(high_chunk, high_chunk)) {
        return 512;
    }
    
    // 先处理低256位块
    if (!_mm256_testz_si256(low_chunk, low_chunk)) {
        const __m256i cnts = _mm256_cttz_epi64(low_chunk);
        const __m256i offsets = _mm256_setr_epi64(0, 64, 128, 192);
        const __m256i total = _mm256_add_epi64(cnts, offsets);
        
        // 生成字节掩码，找到第一个非零元素的索引
        const int mask = _mm256_movemask_epi8(low_chunk);
        const int first_idx = _tzcnt_u32(mask) / 8;
        return static_cast<int>(_mm256_extract_epi64(total, first_idx));
    }
    
    // 低256位全零，处理高256位块
    const __m256i cnts = _mm256_cttz_epi64(high_chunk);
    const __m256i offsets = _mm256_setr_epi64(256, 320, 384, 448);
    const __m256i total = _mm256_add_epi64(cnts, offsets);
    
    const int mask = _mm256_movemask_epi8(high_chunk);
    const int first_idx = _tzcnt_u32(mask) / 8;
    return static_cast<int>(_mm256_extract_epi64(total, first_idx));
}

SSE2实现（兼容老旧CPU） #

SSE2仅支持128位寄存器，需拆分多个块处理，且无原生SIMD末尾零统计指令，需逐个提取64位元素处理：

128位版本 #

#include <immintrin.h>

int CountRZero128_SSE2(const uint64_t* ptr) {
    const __m128i data = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ptr));
    
    if (_mm_testz_si128(data, data)) {
        return 128;
    }
    
    // 提取两个64位元素，分别统计末尾零
    const uint64_t elem0 = _mm_extract_epi64(data, 0);
    const uint64_t elem1 = _mm_extract_epi64(data, 1);
    const __m128i cnts = _mm_set_epi64x(_tzcnt_u64(elem1), _tzcnt_u64(elem0));
    const __m128i offsets = _mm_set_epi64x(64, 0);
    const __m128i total = _mm_add_epi64(cnts, offsets);
    
    // 找到第一个非零元素的索引
    const __m128i non_zero_mask = _mm_cmpgt_epi64(data, _mm_setzero_si128());
    const int mask_val = _mm_movemask_epi8(non_zero_mask);
    const int first_idx = _tzcnt_u32(mask_val) / 8;
    
    return static_cast<int>(_mm_extract_epi64(total, first_idx));
}

512位版本 #

#include <immintrin.h>

int CountRZero512_SSE2(const uint64_t* ptr) {
    for (int i = 0; i < 4; ++i) {
        const __m128i chunk = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ptr + 2*i));
        if (!_mm_testz_si128(chunk, chunk)) {
            const uint64_t elem0 = _mm_extract_epi64(chunk, 0);
            const uint64_t elem1 = _mm_extract_epi64(chunk, 1);
            const __m128i cnts = _mm_set_epi64x(_tzcnt_u64(elem1), _tzcnt_u64(elem0));
            const __m128i offsets = _mm_set_epi64x(64 + i*128, 0 + i*128);
            const __m128i total = _mm_add_epi64(cnts, offsets);
            
            const __m128i non_zero_mask = _mm_cmpgt_epi64(chunk, _mm_setzero_si128());
            const int mask_val = _mm_movemask_epi8(non_zero_mask);
            const int first_idx = _tzcnt_u32(mask_val) / 8;
            
            return static_cast<int>(_mm_extract_epi64(total, first_idx));
        }
    }
    return 512;
}

性能优化注意事项 #

1. 编译选项：编译时需开启对应指令集的优化开关，例如：
   • AVX-512：-mavx512f
   • AVX2：-mavx2
   • SSE2：-msse2（多数编译器默认开启）
2. 内存对齐：如果数据是对齐的，可将_loadu替换为_load（如_mm512_load_si512），进一步提升加载速度。
3. 批量处理：针对海量数组，尽量批量加载数据到SIMD寄存器，减少单次调用的开销。

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