使用 SIMD 搜索密钥

Question

我有以下结构，它存储键和通用的用户指定值：

typedef struct {
        uint32_t  len;
        uint32_t  cap;
        int32_t  *keys;
        void     *vals;
} dict;

现在我想创建一个函数来迭代 keys 和 return 对应的 value.

非SIMD版本：

void*
dict_find(dict *d, int32_t k, size_t s) {
        size_t i;
        i = 0;

        while (i < d->len) {
                if (d->keys[i] == k) {
                        void *p;
                        p = (uint8_t*)d->vals + i * s;

                        return p;
                }

                ++i;
        }

        return NULL;
}

我试图对上面的代码片段进行矢量化并得出以下结论：

void*
dict_find_simd(dict *d, int32_t k, size_t s) {
        __m256i ymm0;
        ymm0 = _mm256_broadcastd_epi32(*(__m128i*)&k);

        __m256i  ymm1;
        uint32_t i;
        int      m;
        uint8_t  b;

        i = 0;
        while (i < d->len) { // [d->len] is aligned in 32 byte box.
                ymm1 = _mm256_load_si256((__m256i*)(d->keys + i));
                ymm1 = _mm256_cmpeq_epi32(ymm1, ymm0);

                m = _mm256_movemask_epi8(ymm1);
                b = __builtin_ctz(m) >> 2;

                i += (8 +  b * d->len); // Artificially break the loop. 
                                        // Remember [i] stores the modified value.
        }

        if (i <= d->len)
                return NULL;

        i -= (8 + b * d->len); // Restore the modified value.
        i += b;

        void *p;
        p = (uint8_t*)d->vals + i * s;

        return p;
}

该功能似乎工作正常（没有进行太多测试）？

但是，有两个问题：

注意：我正在检查 i > d->len 然后我 return 指针。 i 可以溢出，它会 return NULL 在那里。我该如何解决这个问题？
您可能会注意到我使用了 _mm256_movemask_epi8 和 __builtin_ctz 的组合来获取找到的键的索引。有没有更好的方法（也许是一条指令确实获得非零值的位置）来做到这一点（没有 AVX512）？

Answer 1

I'm checking if the i > d->len then I return the pointer. The i can be overflowed and it will return NULL there. How can I solve this issue?

有两种方法可以处理溢出（以及由此导致的潜在越界读取）。

仅使用向量实现最多 i，元素数量可被向量大小整除。如果向量循环没有找到元素，则在标量代码中完成尾部处理。如果输入数据是从其他地方获得的，这个解决方案可能会很好，并且没有简单的方法来优化缓冲区末尾的内存分配和初始化。
允许读取超过缓冲区的末尾，并确保在那里读取的任何内容都不算作有效（找到的）条目。过度分配缓冲区以确保您始终可以读取完整的向量数据。如果将结果 i 与容器中的元素数进行比较，这很容易做到——如果它更大，那么你的算法“找到”了一个超出末尾的元素，你应该指出没有找到任何东西。在某些情况下，这可以自然地来自数据的性质。例如，如果您使用永远不会有效的键值来填充结束元素，或者如果您的关联值可以用于相同的效果（例如，结束后的值是 NULL 指针，这也用于表示“未找到”结果。

You might noticed that I'm using a combination of _mm256_movemask_epi8 and __builtin_ctz in order to get the index of found key. Is there a better way (maybe a single instruction that does get the position of non zero value) to do this (without AVX512)?

我认为没有针对此的单一指令，但您可以提高此组合的性能。请注意，您正在比较 32 位值，这意味着 _mm256_movemask_epi8 为 8 个元素（每个 4 个相等的位）生成一个掩码。如果比较 4 对向量，则可以提高数据密度，然后打包结果，使向量中的每个字节对应一个不同的比较结果，然后应用一个 _mm256_movemask_epi8.

ymm1 = _mm256_load_si256((__m256i*)(d->keys + i));
ymm2 = _mm256_load_si256((__m256i*)(d->keys + i) + 1);
ymm3 = _mm256_load_si256((__m256i*)(d->keys + i) + 2);
ymm4 = _mm256_load_si256((__m256i*)(d->keys + i) + 3);

ymm1 = _mm256_cmpeq_epi32(ymm1, ymm0);
ymm2 = _mm256_cmpeq_epi32(ymm2, ymm0);
ymm3 = _mm256_cmpeq_epi32(ymm3, ymm0);
ymm4 = _mm256_cmpeq_epi32(ymm4, ymm0);

ymm1 = _mm256_packs_epi32(ymm1, ymm2);
ymm3 = _mm256_packs_epi32(ymm3, ymm4);
ymm1 = _mm256_packs_epi16(ymm1, ymm3);
ymm1 = _mm256_permute4x64_epi64(ymm1, _MM_SHUFFLE(3, 1, 2, 0));
ymm1 = _mm256_shuffle_epi32(ymm1, _MM_SHUFFLE(3, 1, 2, 0));

m = _mm256_movemask_epi8(ymm1);
if (m)
{
    b = __builtin_ctz(m); // no shift needed here
    break;
}

（请注意，如果 m 为零，则 __builtin_ctz 结果未定义，但如果检查 i 是否在范围内，则可以在退出循环时减轻这种情况。但是，如上所示，我宁愿在 __builtin_ctz 之前测试 m 并使用它来缩短 __builtin_ctz 并作为打破循环的标志。）

问题是打包是按 128 位通道完成的，这意味着您必须在通道之间随机排列字节才能使用结果。这和打包本身会增加开销，可能会在某种程度上抵消此优化带来的好处。如果使用 128 位向量，则可以节省改组，并可能提高整体性能。我没有对代码进行基准测试，您必须进行测试。

另一个可能的优化方法是，如果比较的 none 是 true，则缩短 packing/shuffling 和 _mm256_movemask_epi8。您可以使用 _mm256_testz_si256 检查所有比较结果向量是否为零，只有当它们不是时才跳出循环。

ymm1 = _mm256_load_si256((__m256i*)(d->keys + i));
ymm2 = _mm256_load_si256((__m256i*)(d->keys + i) + 1);
ymm3 = _mm256_load_si256((__m256i*)(d->keys + i) + 2);
ymm4 = _mm256_load_si256((__m256i*)(d->keys + i) + 3);

ymm1 = _mm256_cmpeq_epi32(ymm1, ymm0);
ymm2 = _mm256_cmpeq_epi32(ymm2, ymm0);
ymm3 = _mm256_cmpeq_epi32(ymm3, ymm0);
ymm4 = _mm256_cmpeq_epi32(ymm4, ymm0);

ymm5 = _mm256_or_si256(ymm1, ymm2);
ymm6 = _mm256_or_si256(ymm3, ymm4);
ymm5 = _mm256_or_si256(ymm5, ymm6);

if (!_mm256_testz_si256(ymm5, ymm5))
{
    ymm1 = _mm256_packs_epi32(ymm1, ymm2);
    ymm3 = _mm256_packs_epi32(ymm3, ymm4);
    ymm1 = _mm256_packs_epi16(ymm1, ymm3);
    ymm1 = _mm256_permute4x64_epi64(ymm1, _MM_SHUFFLE(3, 1, 2, 0));
    ymm1 = _mm256_shuffle_epi32(ymm1, _MM_SHUFFLE(3, 1, 2, 0));

    m = _mm256_movemask_epi8(ymm1);
    b = __builtin_ctz(m);

    break;
}

在这里，3 次 OR 操作比 3 次打包 + 2 次洗牌更快，因此如果您的数据足够大（即，如果平均而言您不希望在初始元素中找到结果），您可能会节省一些周期.如果您发现元素主要位于第一个元素中，那么这将显示比没有 _mm256_testz_si256.

的循环更差的性能

这是根据 Peter Cordes 在评论中的建议对上述代码进行更新的版本。

ymm1 = _mm256_load_si256((__m256i*)(d->keys + i));
ymm2 = _mm256_load_si256((__m256i*)(d->keys + i) + 1);
ymm3 = _mm256_load_si256((__m256i*)(d->keys + i) + 2);
ymm4 = _mm256_load_si256((__m256i*)(d->keys + i) + 3);

ymm1 = _mm256_cmpeq_epi32(ymm1, ymm0);
ymm2 = _mm256_cmpeq_epi32(ymm2, ymm0);
ymm3 = _mm256_cmpeq_epi32(ymm3, ymm0);
ymm4 = _mm256_cmpeq_epi32(ymm4, ymm0);

ymm1 = _mm256_packs_epi32(ymm1, ymm2);
ymm3 = _mm256_packs_epi32(ymm3, ymm4);
ymm5 = _mm256_or_si256(ymm1, ymm3);  // cheap result to branch on 

if (_mm256_movemask_epi8(ymm5) != 0)
{
    ymm1 = _mm256_packs_epi16(ymm1, ymm3);     // now put the bits in order
    ymm1 = _mm256_permutevar8x32_epi32(ymm1,   // or vpermq + vpshufd like before
        _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7));

    m = _mm256_movemask_epi8(ymm1);
    b = __builtin_ctz(m);

    break;
}

改进时考虑到了 Skylake 或类似的微架构：

将两个包移动到条件之上。他们将能够高效地执行，因为每个周期只能执行两个 vpcmpeqd，这足以喂养一个 vpackssdw。鉴于每个周期可以发出两个负载，每个周期两个 vpcmpeqd 是可以实现的。也就是说，两条pack指令竞争端口5不会成为瓶颈。
vpmovmskb 指令只有一个 µop，有 2-3 个周期的延迟，vptest 是两个 µop（3 个周期）。后面的test会和jz/jnz融合，所以_mm256_movemask_epi8上的条件可以执行的稍微快一些。请注意，此时 _mm256_movemask_epi8 应用于虚拟向量 ymm5，稍后不会使用它来生成正确的结果。
我的代码版的两个shuffle可以换成一个vector常量。在这里，我使用 _mm256_setr_epi32 来初始化常量，并且体面的编译器会将其转换为内存中的常量，而无需额外的指令。如果您的编译器不够智能，您可能需要手动执行此操作。另请注意，此常量是额外的内存访问，如果您的查找倾向于提前终止（即，如果条件背后的代码对算法的总执行时间有显着影响），它可能会发挥作用。您可以通过在进入循环之前尽早加载常量来缓解这种情况。该算法不使用很多向量寄存器，因此您必须有足够的空间来保持常量加载。

使用 SIMD 搜索密钥

Searching for the key using SIMD

c

simd

avx