ARM NEON如何实现类似_mm512_mask_compress_epi的过滤压缩逻辑？

需求与标量实现 #

目标函数需保留数组中filter[i] != 0对应的nums[i]元素，剔除其余元素。标量实现代码如下：

static size_t filter_range(int* nums, const uint8_t* filter, size_t size) {
    size_t result_offset = 0;
    for (auto i = 0; i < size; i++) {
        if (filter[i]) {
            *(nums + result_offset) = *(nums + i);
            result_offset++;
        }
    }
    return result_offset;
}

x86 AVX512架构下，_mm512_mask_compress_epi指令可直接高效完成这类掩码压缩操作，但ARM NEON架构下能否用SIMD指令最大化该函数的性能？

现有NEON尝试实现 #

基于ARM位掩码优化思路，已实现filter_range2函数：批量处理16个filter元素，生成标记匹配位的掩码，用Kernighan算法遍历有效位完成元素拷贝，同时对全匹配/全不匹配场景做特殊优化。但该实现性能仍不及AVX512版本：

static size_t filter_range2(int* data, const uint8_t* filter, size_t size) {
    constexpr size_t data_type_size = sizeof(int);
    size_t start_offset = 0;
    size_t result_offset = 0;

#if defined(__ARM_NEON__) || defined(__aarch64__)
    constexpr size_t kBatchNums = 128 / (8 * sizeof(uint8_t));
    while (start_offset + kBatchNums < size) {
        uint8x16_t vfilter = vld1q_u8(filter);
        const uint16x8_t non_zero_mask = vreinterpretq_u16_u8(vmvnq_u8(vceqzq_u8(vfilter)));
        const uint8x8_t res = vshrn_n_u16(non_zero_mask, 4);
        // 每个字节对应4位全1，表示4个匹配的filter元素
        uint64_t matches = vget_lane_u64(vreinterpret_u64_u8(res), 0);

        if (matches == 0) {
            // 无匹配元素，跳过
        } else if (matches == 0xffff'ffff'ffff'ffffull) {
            // 全部匹配，批量拷贝
            memmove(data + result_offset, data + start_offset, kBatchNums * data_type_size);
            result_offset += kBatchNums;
        } else {
            matches &= 0x8888'8888'8888'8888ull;
            // *Kernighan算法*遍历所有匹配位
            for (; matches > 0; matches &= matches - 1) {
                uint32_t index = __builtin_ctzll(matches) >> 2;
                *(data + result_offset++) = *(data + start_offset + index);
            }
        }

        start_offset += kBatchNums;
        filter += kBatchNums;
    }
#endif

    // 处理剩余不足批量的元素
    for (auto i = start_offset; i < size; ++i) {
        if (filter[i]) {
            *(data + result_offset) = *(data + i);
            result_offset++;
        }
    }

    return result_offset;
}

更优的NEON优化方向 #

现有实现的瓶颈在于部分匹配场景下的逐个元素拷贝，可通过以下方式优化：

• 利用AArch64 NEON掩码加载/存储指令：使用vld1q_s32_x系列掩码加载指令，直接批量加载匹配元素到向量，再一次性存储到结果位置，避免逐个拷贝的开销。
• 预计算匹配元素索引：通过SIMD指令生成匹配元素的偏移量，用gather操作批量获取元素，减少标量循环次数。
• 优化掩码生成逻辑：直接将filter向量转换为位掩码，而非4位聚合的字节，更高效地利用位操作遍历匹配位置。

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