如何在 AVX 向量中找到元素的索引？

Posted 2023-02-16

【中文标题】如何在 AVX 向量中找到元素的索引？

【英文标题】：How to find the index of an element in the AVX vector?

【发布时间】：2019-08-06 14:32:18

【问题描述】：

我正在尝试使用 AVX 编写硬件加速哈希表，其中每个存储桶都有固定大小（AVX 向量大小）。问题是如何实现向量的快速搜索。

不完整的可能解决方案：

example target hash: 2

<1  7  8  9  2  6  3  5>  //  vector of hashes
<2  2  2  2  2  2  2  2>  //  mask vector of target hash
------------------------  //  equality comparison
<0  0  0  0 -1  0  0  0>  //  result of comparison
<0  1  2  3  4  5  6  7>  //  vector of indexes
------------------------  //  and operation
<0  0  0  0  4  0  0  0>  //  index of target hash

如何从最后一个向量中提取目标哈希的索引？

---

另一种（慢）使用标量积的可能解决方案：

<1  7  8  9  2  6  3  5>  //  vector of hashes
<2  2  2  2  2  2  2  2>  //  mask vector of target hash
------------------------  //  equality comparison
<0  0  0  0 -1  0  0  0>  //  result of comparison
<0  1  2  3  4  5  6  7>  //  vector of indexes
------------------------  //  dot
            -4

【参考方案1】：

对此合适的水平操作是 MOVMSKPS，它从 XMM/YMM 向量中提取掩码（基本上，它从每个通道中收集最高位）。完成后，您可以执行 TZCNT 或 LZCNT 来获取索引。

例子：

#include <intrin.h>
#include <immintrin.h>

int getIndexOf(int const values[8], int target)

    __m256i valuesSimd = _mm256_loadu_si256((__m256i const*)values);
    __m256i targetSplatted = _mm256_set1_epi32(target);
    __m256i equalBits = _mm256_cmpeq_epi32(valuesSimd, targetSplatted);
    unsigned equalMask = _mm256_movemask_ps(_mm256_castsi256_ps(equalBits));
    int index = _tzcnt_u32(equalMask);
    return index;