使用 SIMD 搜索密钥

问题描述

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

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,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）？

解决方法

我正在检查 i > d->len 是否返回指针。 i 可以溢出，它会在那里返回 NULL。我该如何解决这个问题？

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

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

您可能注意到我使用了 _mm256_movemask_epi8 和 __builtin_ctz 的组合来获取找到的键的索引。有没有更好的方法（也许是一条可以获取非零值位置的指令）来做到这一点（没有 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,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 位通道完成的，这意味着您必须先在通道之间打乱字节，然后才能使用结果。这和打包本身增加了开销，可能会抵消这种优化的好处。如果使用 128 位向量，则可以节省洗牌，并且可能会提高整体性能。我没有对代码进行基准测试，您必须进行测试。

要考虑的另一个可能的优化是缩短打包/改组和 _mm256_movemask_epi8（如果所有比较都不是 true）。您可以使用 _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);

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,0));
    ymm1 = _mm256_shuffle_epi32(ymm1,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,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,5,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 被应用于虚拟向量 ymm5，稍后不会使用它来产生正确的结果。
我的代码版本中的两次 shuffle 可以替换为带有向量常量的一次。在这里，我使用 _mm256_setr_epi32 来初始化常量，体面的编译器会将它转换为内存中的常量，而无需额外的指令。如果您的编译器不够智能，您可能需要手动执行此操作。另请注意，此常量是额外的内存访问，如果您的查找倾向于提前终止（即，如果条件背后的代码对算法的总执行时间有显着贡献），则可能会起作用。您可以通过在进入循环之前尽早加载常量来缓解这种情况。该算法没有使用很多向量寄存器，所以你必须有足够的空间来保持常量加载。

avx c simd