从 __m512i/__m256i 向量中提取非零元素

Question

给定一个包含 64 个

__m512i

元素的

char

向量：

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

（注：元素的值在[-1, 1]之间）。

是否有任何优雅的方法来提取所有非零元素并将它们打包到另一个

__m512i

向量中，如下所示：

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

我幼稚的做法是：

non_zero_mask = _mm512_test_epi8_mask(X,X);

。然后用

while

循环遍历掩码，将每个元素一一添加到新向量中（是的，它相对较慢）

Answer 1

AVX-512VBMI2（Ice Lake 及更高版本）具有

vpcompressb

，可根据面罩（例如您的

_mm512_test_epi8_mask(X,X)

中的面罩）进行左侧包装。 它的成本为 2 uop（对于 Intel 上的端口 5），但仍然比没有它时可以做的任何事情要好得多。

在此之前，AVX-512F 仅支持 dword 和 qword 元素大小。 我的 AVX512 回答 AVX2 基于掩码打包左侧的最有效方法是什么？ 展示了如何使用

ps

版本；字节版本应该以相同的方式工作。 （我的 AVX2 答案在 8 字节整数上使用 BMI2

pdep

来创建随机播放掩码，但这不适用于较窄的元素。）

---

如果没有

compress

指令，左打包确实很难，这就是为什么它是一个有价值的原始操作作为构建块。

根据您的数据密度，在没有 AVX-512VBMI2 的情况下，您可能会考虑将字节解包为

vpcompressd

的双字，并在存储之前使用

vpmovdb

缩小到 8 位。

也许

test

一次 64 个字节并使用 3x

kshiftrq

为接下来的 4 次压缩进行输入？

vptestmd

和

kshift

在 Intel 上都是仅端口 5，与

vpcompress

的 2p5 uops 竞争，但 Zen 4 可以在不同端口上运行

kshiftrq

。 (https://uops.info/)但是如果你在加载数据时扩展数据，就像

vpmovzxbd

一样，你永远不会在单个向量寄存器中拥有64个字节，所以是的，你会想要分别对每个向量进行

_mm512_test_epi32_mask

，而不是花费更多的洗牌来扩大例如不带

__m512i

的

vpermb

的第二个 128 位通道需要 AVX-512VBMI，它是在与 VBMI2 相同的 CPU 上引入的。 (https://en.wikipedia.org/wiki/AVX-512#CPUs_with_AVX-512)

---

我不认为你的条件的简单性（非零）为以不同的方式做到这一点提供了任何好的可能性，例如洗牌和混合什么的。