使用 SIMD 搜索密钥

Posted 2023-02-16

【中文标题】使用 SIMD 搜索密钥

【英文标题】：Searching for the key using SIMD

【发布时间】：2021-04-23 09:12:58

【问题描述】：

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

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;

我尝试对上面的 sn-p 进行矢量化处理，结果如下：

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

【问题讨论】：

为什么不简单地for(i=0; i<d-len; i+=8) /* SIMD stuff */ if(b) break;？

我认为使用 SIMD 进行分支是“不可行的”。也许我错了。

其实你也可以用if(m) break;break，在循环之后做b的计算。如果您要处理许多非常小的数组（通常适合几个寄存器），则分支将是低效的。在这种情况下，会有更有效的解决方案。 （这个问题肯定有重复，但我暂时不想找）

人为地修改循环计数器通常对性能来说是一个更大的问题，因为它破坏了循环展开、预加载数据等的许多可能性。如果你只是引入一个新的会稍微好一点变量break_next 并将其放入循环条件中。不过你也可以直接break。

如果m 是0，

__builtin_ctz(m) 是未定义的行为。使用_tzcnt_u32(m) 是您想要明确定义的行为。大多数（全部？）AVX2 CPU 都有 tzcnt。在实践中，__bultin_ctz(m) 将对大多数编译器使用 tzcnt，因此无论如何都会产生明确定义的结果。但是，是的，同意@chtz，循环计数器更新看起来很糟糕，我认为在上一个比较和下一个迭代的加载地址之间创建一个依赖链，所以你的瓶颈是 latency 而不是吞吐量。跨度>

【参考方案1】：

我正在检查 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, _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 是否在界限内，则可以在退出循环时缓解这种情况。但是，如上所示，我宁愿测试m 在 __builtin_ctz 之前，并使用它来缩短 __builtin_ctz 并作为中断循环的标志。）

这样做的问题是打包是按 128 位通道完成的，这意味着您必须先在通道之间打乱字节，然后才能使用结果。这和打包本身会增加开销，这可能会在一定程度上抵消这种优化带来的好处。如果您使用 128 位向量，则可以节省混洗，它可能会提高整体性能。我没有对代码进行基准测试，您必须进行测试。

如果没有一个比较是true，则要考虑的另一个可能的优化是缩短打包/改组和_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 的循环更差的性能。

---

这是上面代码的更新版本，它基于 cmets 中 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。换句话说，竞争端口 5 的两条打包指令不会成为瓶颈。

vpmovmskb 指令只有一个 µop，延迟 2-3 个周期，vptest 是两个 µops（3 个周期）。随后的test 将与jz/jnz 融合，因此_mm256_movemask_epi8 上的条件可以执行得稍微快一些。请注意，此时_mm256_movemask_epi8 被应用于虚拟向量ymm5，以后不会使用它来产生正确的结果。

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

【讨论】：

感谢您的回答。只有一件事：我不明白第一个 sn-p 将如何工作？偏移量不应该是+8, +16, ... 而不是+1, +2？如果没有，那将如何更有效？对于__builtin_ctz Peter 建议使用_tzcnt_u32。我想我也可以这样做（不知道它在内部是如何工作的，也许它就像你说的那样）。

@Hrant +1 等被添加到 __m256i* 指针，以便按预期工作。 _tzcnt_u32 是一个有效的选项，但是在这里，由于您无论如何都想测试比较结果，因此不需要将其应用于零输入，因此 __builtin_ctz 工作正常。

@Hrant 将+1 添加到__m256i* 指针会将其增加32 个字节（或8 个32 位元素）。因此 4 个加载不会重叠并加载 128 个相邻字节的数据。至于延迟，是的，每个加载指令都有延迟，但是现代 CPU 可以并行发出多个加载（例如，如果我没记错的话，Skylake 每个周期最多 2 个加载），因此可以隐藏延迟。另外，CPU 可能能够推测性地并行执行多个循环迭代，但在这种情况下，如果数据是冷的，它可能无论如何都会受到内存限制。

另外，vpmovmskb / test+jnz 在宏融合后只有 2 uops，但 vptest / jnz 是 2+1，所以实际上即使在“错误”的 v 上 if(movemask(v)) 也比当您的 v 来自比较结果时使用if(_mm256_testz_si256(v,v))（因此符号位具有您想要的数据。）

@AndreySemashev：哦，脑子放个屁，是的，_mm256_permute4x64_epi64 是vpermq。嗯，对，经过 2 步打包后，你在 dword 粒度上混淆了，而不仅仅是 qword。所以你需要vpermd 一步修复它，这需要一个矢量控制常数。 （如果您经常搜索，可能值得加载。）