如何优雅地从__m512i向量提取非零char元素并打包？

问题场景 #

给定一个包含64个char元素的__m512i向量，元素值范围为[-1, 1]，示例如下：

index: 0,  1, 2, 3, 4,  5, 6, 7, 8, 9, 10,...
value: 1, -1, 1, 0, 0, -1, 1, 1, 0, 0, 1,...

需要将所有非零元素按原顺序提取出来，打包到另一个__m512i向量中，预期输出如下：

expected output: 1, -1, 1, -1, 1, 1, 1,...

当前使用的朴素方法是通过non_zero_mask = _mm512_test_epi8_mask(X,X);生成掩码后，用while循环遍历掩码逐个添加元素，这种方法速度较慢，希望找到更优雅高效的实现方式。

高效解决方案 #

利用AVX-512的压缩指令可以直接完成这个需求，无需循环，完全通过SIMD硬件加速实现：

1. 生成非零元素掩码：使用_mm512_cmpneq_epi8_mask指令，将原向量与零向量比较，得到对应每个8位元素是否非零的__mmask64掩码。
2. 压缩非零元素：使用_mm512_compress_epi8指令，根据掩码将原向量中的非零元素按顺序打包到目标向量的前端，剩余位置自动补零。

代码示例 #

#include <immintrin.h>

__m512i extract_non_zero_epi8(__m512i input) {
    // 生成非零元素的掩码：元素不等于0则对应掩码位为1
    __mmask64 non_zero_mask = _mm512_cmpneq_epi8_mask(input, _mm512_setzero_si512());
    // 压缩非零元素到目标向量前端
    return _mm512_compress_epi8(_mm512_setzero_si512(), input, non_zero_mask);
}

说明 #

• _mm512_compress_epi8是AVX-512VL/BW指令集中的指令，支持将源向量中掩码为1的元素按原顺序复制到目标向量的连续位置，未被选中的位置用目标向量的原始值填充（这里用零向量，所以后面补零）。
• 整个过程通过硬件原生指令完成元素筛选和打包，效率远高于循环遍历的朴素方法，是最优雅的实现方式。

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