c - 使用 SIMD 搜索 key

Question

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

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

现在我想创建一个迭代 keys 的函数并返回对应的 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 是否然后我返回指针。 i可以溢出，它会返回 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?

有两种方法可以处理溢出(以及由此引起的潜在越界读取)。

最多只能使用 vector 实现 i可以被 vector 大小除以元素的数量。如果 vector 循环没有找到元素，则在标量代码中完成尾部处理。如果输入数据是从其他地方获得的，那么这个解决方案可能会很好，并且没有简单的方法来优化缓冲区末尾的内存分配和初始化。

允许读取越过缓冲区的末尾，并确保在那里读取的任何内容都不算作有效的(找到的)条目。过度分配缓冲区以确保您始终可以读取完整 vector 的数据。如果比较结果 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 对 vector ，则可以提高数据密度，然后将结果打包以便 vector 中的每个字节对应于不同的比较结果，然后应用一个 _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;
}

(请注意，如果 __builtin_ctz 为零，则 m 结果是未定义的，但是如果您检查 i 是否在范围内，您可以在退出循环时减轻这种情况。但是，如上所示，我宁愿测试 m在 __builtin_ctz 之前，并用它来快捷 __builtin_ctz 并作为打破循环的标志。)
这样做的问题是打包是按 128 位 channel 完成的，这意味着您必须先在 channel 之间打乱字节，然后才能使用结果。这和打包本身增加了开销，可能会抵消这种优化的好处。如果使用 128 位 vector ，则可以节省洗牌，并且可能会提高整体性能。我没有对代码进行基准测试，您将不得不进行测试。
要考虑的另一个可能的优化是缩短打包/改组和 _mm256_movemask_epi8如果没有一个比较是 true .您可以使用 _mm256_testz_si256检查所有比较结果 vector 是否为零并仅在它们不是时才跳出循环。

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 的循环显示更差的性能。 .

这是基于 建议的上述代码的更新版本彼得·科德斯 在评论中。

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考虑到每个周期可以发出两个负载，每个周期是可以实现的。也就是说，竞争端口5的两条pack指令不会成为瓶颈。

vpmovmskb指令只有一个微操作，延迟为 2-3 个周期和 vptest是两个微操作(3 个周期)。后续 test将与 jz 融合/jnz ，所以 _mm256_movemask_epi8 上的条件可以执行得稍微快一点。注意此时_mm256_movemask_epi8应用于虚拟 vector ymm5 ，稍后不会使用它来产生正确的结果。

我的代码版本中的两个 shuffle 可以用一个 vector 常量替换。在这里，我正在使用 _mm256_setr_epi32初始化常量，体面的编译器会将其转换为内存中的常量，而无需额外的指令。如果您的编译器不够智能，您可能需要手动执行此操作。另请注意，此常量是额外的内存访问，如果您的查找倾向于提前终止(即，如果条件背后的代码对算法的总执行时间有显着贡献)，则可能会起作用。您可以通过在进入循环之前尽早加载常量来缓解这种情况。该算法不使用很多 vector 寄存器，因此您必须有足够的空间来保持常量加载。