使用AVX/SIMD查找16位值首次出现的索引

要在__m256i类型的256位数据中定位指定16位值的首次出现位置，可通过以下步骤基于AVX2指令集实现：

1. 广播目标值：用_mm256_set1_epi16将待查找的16位值广播到__m256i寄存器，生成每个元素均为目标值的向量。
2. 逐元素比较：用_mm256_cmpeq_epi16将原数据向量与广播后的目标向量做相等比较，结果向量中与目标值相等的16位元素会被设为0xFFFF，不等的设为0x0000。
3. 生成有效掩码并定位索引：由于没有直接的16位版本movemask，可通过以下两种方式处理：

方法一：256位掩码整体处理 #

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

int find_first_epi16_index(__m256i data, uint16_t target) {
    __m256i target_broadcast = _mm256_set1_epi16(target);
    __m256i cmp_result = _mm256_cmpeq_epi16(data, target_broadcast);
    
    // 将每个16位相等标记(0xFFFF)右移8位，得到低8位为0xFF、高8位为0x00的元素
    __m256i shifted_cmp = _mm256_srli_epi16(cmp_result, 8);
    // 生成掩码：偶数位对应每个16位元素是否相等
    int mask = _mm256_movemask_epi8(shifted_cmp);
    // 仅保留偶数位的有效标志
    mask &= 0x55555555;
    
    if (mask == 0) {
        return -1; // 未找到目标值
    }
    
    // 找到第一个置位的位，转换为16位元素索引
    int first_set_bit = __builtin_ctz(mask);
    return first_set_bit / 2;
}

方法二：拆分128位部分处理（更高效） #

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

int find_first_epi16_index(__m256i data, uint16_t target) {
    __m256i target_broadcast = _mm256_set1_epi16(target);
    __m256i cmp_result = _mm256_cmpeq_epi16(data, target_broadcast);
    
    // 处理低128位（前8个16位元素）
    __m128i cmp_low = _mm256_castsi256_si128(cmp_result);
    int mask_low = _mm_movemask_epi8(cmp_low);
    // 将每两位掩码合并为一位：相等则为1，不等为0
    mask_low = (mask_low & 0x5555) | ((mask_low >> 1) & 0x5555);
    if (mask_low != 0) {
        return __builtin_ctz(mask_low);
    }
    
    // 处理高128位（后8个16位元素）
    __m128i cmp_high = _mm256_extracti128_si256(cmp_result, 1);
    int mask_high = _mm_movemask_epi8(cmp_high);
    mask_high = (mask_high & 0x5555) | ((mask_high >> 1) & 0x5555);
    if (mask_high != 0) {
        return 8 + __builtin_ctz(mask_high);
    }
    
    return -1; // 未找到目标值
}

补充说明 #

• __builtin_ctz是GCC/Clang的内置函数，用于查找整数中第一个置位的最低位位置；MSVC环境可替换为__tzcnt_u32。
• 两种方法均通过向量并行比较生成掩码，再定位首个匹配的16位元素索引，时间复杂度为O(1)，充分利用AVX2的并行计算能力。

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