ARMv8-A NEON加速14元素uint32数组搜索优化及SVE2适配咨询

我们正在给开源项目cachegrand做AARCH64架构移植，目前大部分工作已完成，现在要基于NEON指令实现14个uint32元素数组的加速搜索功能。搜索逻辑是：输入目标值、uint32数组和忽略匹配的掩码，找出数组中与目标值匹配的元素，构建对应的位掩码（最低位对应数组第一个元素），再和掩码的取反值按位与，最后通过统计尾随零得到首个匹配的索引（99.9%的场景下掩码都是0）。

我自己写了一份NEON实现，但感觉代码冗余，想优化性能和实现方式，同时想问：这种简单操作有没有必要适配SVE2指令集（该指令集不强制要求256位寄存器支持）？

现有实现代码 #

NEON版本 #

uint8_t hashtable_mcmp_support_hash_search_armv8a_neon_14(
        uint32_t hash,
        volatile uint32_t* hashes,
        uint32_t skip_indexes_mask) {
    uint32x4_t tmp;
    uint32_t compacted_result_mask = 0;
    uint32_t skip_indexes_mask_inv = ~skip_indexes_mask;
    static const int32x4_t shift = {0, 1, 2, 3};

    uint32x4_t cmp_vector = vdupq_n_u32(hash);

    uint32x4_t ring_vector_0_3 = vld1q_u32((hashtable_hash_half_t*)hashes + 0);
    uint32x4_t cmp_vector_0_3 = vceqq_u32(ring_vector_0_3, cmp_vector);
    tmp = vshrq_n_u32(cmp_vector_0_3, 31);
    compacted_result_mask |=  vaddvq_u32(vshlq_u32(tmp, shift)) << 0;

    uint32x4_t ring_vector_4_7 = vld1q_u32((hashtable_hash_half_t*)hashes + 4);
    uint32x4_t cmp_vector_4_7 = vceqq_u32(ring_vector_4_7, cmp_vector);
    tmp = vshrq_n_u32(cmp_vector_4_7, 31);
    compacted_result_mask |=  vaddvq_u32(vshlq_u32(tmp, shift)) << 4;

    uint32x4_t ring_vector_8_11 = vld1q_u32((hashtable_hash_half_t*)hashes + 8);
    uint32x4_t cmp_vector_8_11 = vceqq_u32(ring_vector_8_11, cmp_vector);
    tmp = vshrq_n_u32(cmp_vector_8_11, 31);
    compacted_result_mask |=  vaddvq_u32(vshlq_u32(tmp, shift)) << 8;

    uint32x4_t ring_vector_10_13 = vld1q_u32((hashtable_hash_half_t*)hashes + 10);
    uint32x4_t cmp_vector_10_13 = vceqq_u32(ring_vector_10_13, cmp_vector);
    tmp = vshrq_n_u32(cmp_vector_10_13, 31);
    compacted_result_mask |=  vaddvq_u32(vshlq_u32(tmp, shift)) << 10;

    return __builtin_ctz(compacted_result_mask & skip_indexes_mask_inv);
}

参考AVX2版本 #

static inline uint8_t hashtable_mcmp_support_hash_search_avx2_14(
        uint32_t hash,
        volatile uint32_t* hashes,
        uint32_t skip_indexes_mask) {
    uint32_t compacted_result_mask = 0;
    uint32_t skip_indexes_mask_inv = ~skip_indexes_mask;
    __m256i cmp_vector = _mm256_set1_epi32(hash);

    // The second load, load from the 6th uint32 to the 14th uint32, _mm256_loadu_si256 always loads 8 x uint32
    for(uint8_t base_index = 0; base_index < 12; base_index += 6) {
        __m256i ring_vector = _mm256_loadu_si256((__m256i*) (hashes + base_index));
        __m256i result_mask_vector = _mm256_cmpeq_epi32(ring_vector, cmp_vector);

        // Uses _mm256_movemask_ps to reduce the bandwidth
        compacted_result_mask |= (uint32_t)_mm256_movemask_ps(_mm256_castsi256_ps(result_mask_vector)) << (base_index);
    }

    return _tzcnt_u32(compacted_result_mask & skip_indexes_mask_inv);
}

优化后的NEON实现 #

原代码存在重复加载、分散掩码压缩的问题，以下是优化版本：

uint8_t hashtable_mcmp_support_hash_search_armv8a_neon_14(
        uint32_t hash,
        volatile uint32_t* hashes,
        uint32_t skip_indexes_mask) {
    uint32_t skip_indexes_mask_inv = ~skip_indexes_mask;
    uint32x4_t cmp_vec = vdupq_n_u32(hash);
    uint32_t result_mask = 0;

    // 加载0-11索引的12个元素，无重复
    uint32x4_t vec0 = vld1q_u32((uint32_t*)hashes);
    uint32x4_t vec1 = vld1q_u32((uint32_t*)hashes + 4);
    uint32x4_t vec2 = vld1q_u32((uint32_t*)hashes + 8);

    // 批量比较生成匹配标记（匹配为0xFFFFFFFF，否则为0）
    uint32x4_t cmp0 = vceqq_u32(vec0, cmp_vec);
    uint32x4_t cmp1 = vceqq_u32(vec1, cmp_vec);
    uint32x4_t cmp2 = vceqq_u32(vec2, cmp_vec);

    // 合并比较结果并压缩为位掩码：提取每个uint32的最高位作为对应索引的bit
    uint8x16_t packed_cmp = vcombine_u8(vreinterpret_u8_u32(cmp0), vreinterpret_u8_u32(cmp1));
    packed_cmp = vext_u8(packed_cmp, vreinterpret_u8_u32(cmp2), 8);
    uint16_t tmp_mask = vgetq_lane_u16(vmovl_u8(vmovn_u16(vreinterpret_u16_u8(packed_cmp))), 0);
    result_mask |= tmp_mask & 0xFFF; // 保留前12位对应0-11索引

    // 处理剩余12-13索引的2个元素
    uint32x2_t vec3 = vld1_u32((uint32_t*)hashes + 12);
    uint32x2_t cmp3 = vceqq_u32(vec3, vdup_n_u32(hash));
    uint8_t tmp_mask2 = (vget_lane_u32(cmp3, 0) >> 31) | ((vget_lane_u32(cmp3, 1) >> 31) << 1);
    result_mask |= tmp_mask2 << 12;

    return __builtin_ctz(result_mask & skip_indexes_mask_inv);
}

优化点说明 #

• 消除冗余加载：原代码重复加载10-11索引的元素，优化后改为仅加载未覆盖的12-13索引，减少内存操作开销
• 批量掩码压缩：模仿AVX2的_mm256_movemask_ps逻辑，通过NEON的向量合并、窄化指令一次性提取12个元素的匹配位，替代原有的分散移位累加操作
• 精简剩余元素处理：用双字寄存器处理最后2个元素，避免使用四字寄存器的冗余操作

SVE2适配建议 #

是否值得适配SVE2需结合以下几点判断：

• 硬件部署场景：如果目标设备以支持SVE2的ARM硬件为主（如AWS Graviton3、新一代移动端旗舰芯片），适配能获得性能收益；若面向仅支持NEON的旧设备，则无需适配
• 性能收益潜力：SVE2的可变长度寄存器可一次性加载全部14个元素，减少加载次数；且针对"找首个匹配项"的高频场景，SVE2的first_true类指令可直接返回索引，跳过掩码构建和ctz操作，能节省少量周期
• 维护成本：需同时维护NEON和SVE2两套实现，增加代码复杂度；若项目已有自动向量化构建逻辑，可优先依赖编译器优化而非手动实现
• 业务优先级：如果该搜索操作是项目的性能热点，且目标用户群体大量使用SVE2硬件，适配是值得的；否则优化后的NEON版本已足够高效

内容的提问来源于stack exchange，提问作者Daniele Salvatore Albano