使用 AVX2 查找元素索引 - 代码优化

Question

我正在摆弄 AVX2 来编写一些代码，该代码能够在具有 14 个条目和 return 找到的条目的索引的数组中搜索 32 位散列。

因为很可能绝大多数命中都在数组的前 8 个条目内，因此可以改进此代码，添加 __builtin_expect 的用法现在不是我的首要任务。

虽然哈希数组（在代码中由变量 hashes 表示）的长度始终为 14 个条目，但它包含在这种结构中

typedef struct chain_ring chain_ring_t;
struct chain_ring {
    uint32_t hashes[14];
    chain_ring_t* next;
    ...other stuff...
} __attribute__((aligned(16)))

代码在这里

int8_t hash32_find_14_avx2(uint32_t hash, volatile uint32_t* hashes) {
    uint32_t compacted_result_mask, leading_zeroes;
    __m256i cmp_vector, ring_vector, result_mask_vector;
    int8_t found_index = -1;

    if (hashes[0] == hash) {
        return 0;
    }

    for(uint8_t base_index = 0; base_index < 14; base_index += 8) {
        cmp_vector = _mm256_set1_epi32(hash);
        ring_vector = _mm256_stream_load_si256((__m256i*) (hashes + base_index));

        result_mask_vector = _mm256_cmpeq_epi32(ring_vector, cmp_vector);
        compacted_result_mask = _mm256_movemask_epi8(result_mask_vector);

        if (compacted_result_mask != 0) {
            leading_zeroes = 32 - __builtin_clz(compacted_result_mask);
            found_index = base_index + (leading_zeroes >> 2u) - 1;
            break;
        }
    }

    return found_index > 13 ? -1 : found_index;
}

简单解释一下逻辑，它搜索前 8 个条目，然后搜索后 8 个条目。如果找到的索引大于 13，则意味着它找到了与某些不属于数组的内容的匹配项，因此必须将其视为不匹配。

备注：

为了加速加载（来自对齐内存）我正在使用 _mm256_stream_load_si256
因为上面提到的，我需要检查一下 returned 值是否大于 13，我不太喜欢这个特定的部分，我应该使用 _mm256_maskload_epi32?
我正在使用 for 循环来避免重复代码，gcc 当然会展开循环
我正在使用 __builtin_clz 但我正在使用 -mlzcnt 编译代码，因为据我所知，AMD cpus 比运行 bsr 指令慢得多， gcc 使用 lzcnt 而不是带有标志
第一个 IF 平均引入了大约 0.30 ns 的延迟，但平均减少了第一次匹配的时间 0.6ns
代码仅适用于 64 位机器
在某些时候我需要为 aarch64 优化此代码

这里有一个很好的 link 对生产的组件的 godbolt https://godbolt.org/z/5bxbN6

我也实现了 SSE 版本（在要点中）但逻辑是一样的，虽然我不确定它的性能是否值得

作为参考，我构建了一个简单的线性搜索函数，并使用 google-benchmark lib

与它进行了性能比较

int8_t hash32_find_14_loop(uint32_t hash, volatile uint32_t* hashes) {
    for(uint8_t index = 0; index <= 14; index++) {
        if (hashes[index] == hash) {
            return index;
        }
    }

    return -1;
}

完整代码可在此处获得 url https://gist.github.com/danielealbano/9fcbc1ff0a42cc9ad61be205366bdb5f

除了 google-benchmark 库的必要标志外，我正在使用 -avx2 -avx -msse4 -O3 -mbmi -mlzcnt

执行每个元素的工作台（我想比较循环与备选方案）

----------------------------------------------------------------------------------------------------
Benchmark                                                          Time             CPU   Iterations
----------------------------------------------------------------------------------------------------
bench_template_hash32_find_14_loop/0/iterations:100000000       0.610 ns        0.610 ns    100000000
bench_template_hash32_find_14_loop/1/iterations:100000000        1.16 ns         1.16 ns    100000000
bench_template_hash32_find_14_loop/2/iterations:100000000        1.18 ns         1.18 ns    100000000
bench_template_hash32_find_14_loop/3/iterations:100000000        1.19 ns         1.19 ns    100000000
bench_template_hash32_find_14_loop/4/iterations:100000000        1.28 ns         1.28 ns    100000000
bench_template_hash32_find_14_loop/5/iterations:100000000        1.26 ns         1.26 ns    100000000
bench_template_hash32_find_14_loop/6/iterations:100000000        1.52 ns         1.52 ns    100000000
bench_template_hash32_find_14_loop/7/iterations:100000000        2.15 ns         2.15 ns    100000000
bench_template_hash32_find_14_loop/8/iterations:100000000        1.66 ns         1.66 ns    100000000
bench_template_hash32_find_14_loop/9/iterations:100000000        1.67 ns         1.67 ns    100000000
bench_template_hash32_find_14_loop/10/iterations:100000000       1.90 ns         1.90 ns    100000000
bench_template_hash32_find_14_loop/11/iterations:100000000       1.89 ns         1.89 ns    100000000
bench_template_hash32_find_14_loop/12/iterations:100000000       2.13 ns         2.13 ns    100000000
bench_template_hash32_find_14_loop/13/iterations:100000000       2.20 ns         2.20 ns    100000000
bench_template_hash32_find_14_loop/14/iterations:100000000       2.32 ns         2.32 ns    100000000
bench_template_hash32_find_14_loop/15/iterations:100000000       2.53 ns         2.53 ns    100000000
bench_template_hash32_find_14_sse/0/iterations:100000000        0.531 ns        0.531 ns    100000000
bench_template_hash32_find_14_sse/1/iterations:100000000         1.42 ns         1.42 ns    100000000
bench_template_hash32_find_14_sse/2/iterations:100000000         2.53 ns         2.53 ns    100000000
bench_template_hash32_find_14_sse/3/iterations:100000000         1.45 ns         1.45 ns    100000000
bench_template_hash32_find_14_sse/4/iterations:100000000         2.26 ns         2.26 ns    100000000
bench_template_hash32_find_14_sse/5/iterations:100000000         1.90 ns         1.90 ns    100000000
bench_template_hash32_find_14_sse/6/iterations:100000000         1.90 ns         1.90 ns    100000000
bench_template_hash32_find_14_sse/7/iterations:100000000         1.93 ns         1.93 ns    100000000
bench_template_hash32_find_14_sse/8/iterations:100000000         2.07 ns         2.07 ns    100000000
bench_template_hash32_find_14_sse/9/iterations:100000000         2.05 ns         2.05 ns    100000000
bench_template_hash32_find_14_sse/10/iterations:100000000        2.08 ns         2.08 ns    100000000
bench_template_hash32_find_14_sse/11/iterations:100000000        2.08 ns         2.08 ns    100000000
bench_template_hash32_find_14_sse/12/iterations:100000000        2.55 ns         2.55 ns    100000000
bench_template_hash32_find_14_sse/13/iterations:100000000        2.53 ns         2.53 ns    100000000
bench_template_hash32_find_14_sse/14/iterations:100000000        2.37 ns         2.37 ns    100000000
bench_template_hash32_find_14_sse/15/iterations:100000000        2.59 ns         2.59 ns    100000000
bench_template_hash32_find_14_avx2/0/iterations:100000000       0.537 ns        0.537 ns    100000000
bench_template_hash32_find_14_avx2/1/iterations:100000000        1.37 ns         1.37 ns    100000000
bench_template_hash32_find_14_avx2/2/iterations:100000000        1.38 ns         1.38 ns    100000000
bench_template_hash32_find_14_avx2/3/iterations:100000000        1.36 ns         1.36 ns    100000000
bench_template_hash32_find_14_avx2/4/iterations:100000000        1.37 ns         1.37 ns    100000000
bench_template_hash32_find_14_avx2/5/iterations:100000000        1.38 ns         1.38 ns    100000000
bench_template_hash32_find_14_avx2/6/iterations:100000000        1.40 ns         1.40 ns    100000000
bench_template_hash32_find_14_avx2/7/iterations:100000000        1.39 ns         1.39 ns    100000000
bench_template_hash32_find_14_avx2/8/iterations:100000000        1.99 ns         1.99 ns    100000000
bench_template_hash32_find_14_avx2/9/iterations:100000000        2.02 ns         2.02 ns    100000000
bench_template_hash32_find_14_avx2/10/iterations:100000000       1.98 ns         1.98 ns    100000000
bench_template_hash32_find_14_avx2/11/iterations:100000000       1.98 ns         1.98 ns    100000000
bench_template_hash32_find_14_avx2/12/iterations:100000000       2.03 ns         2.03 ns    100000000
bench_template_hash32_find_14_avx2/13/iterations:100000000       1.98 ns         1.98 ns    100000000
bench_template_hash32_find_14_avx2/14/iterations:100000000       1.96 ns         1.96 ns    100000000
bench_template_hash32_find_14_avx2/15/iterations:100000000       1.97 ns         1.97 ns    100000000

感谢任何建议！

---更新

我已经用@chtz 制作的 b运行chless 实现更新了要点，并将 __lzcnt32 替换为 _tzcnt_u32，我不得不稍微改变行为以考虑未找到当 32 被 returned 而不是 -1 但并不重要时。

他们运行所在的 CPU 是英特尔酷睿 i7 8700（6c/12t，3.20GHZ）。

bench 使用 cpu-pinning，使用比物理或逻辑 cpu 核心更多的线程，并执行一些额外的操作，特别是 for 循环，因此存在开销，但两者之间是相同的两个测试，因此它应该以相同的方式影响它们。

如果您想运行测试，您需要调整 CPU_CORE_LOGICAL_COUNT 以手动匹配 cpu 的逻辑 cpu 核心数。

有趣的是，当存在更多争用时（从单线程到 64 线程），性能提升如何从 +17% 跃升至 +41%。我用 128 和 256 个线程进行了更多测试运行使用 AVX2 时速度提高了 +60%，但我没有包括下面的数字。

（bench_template_hash32_find_14_avx2 是 b运行chless 版本的基准，我缩短了名称以使 post 更具可读性）

------------------------------------------------------------------------------------------
Benchmark                                                                 CPU   Iterations
------------------------------------------------------------------------------------------
bench_template_hash32_find_14_loop/iterations:10000000/threads:1      45.2 ns     10000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:2      50.4 ns     20000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:4      52.1 ns     40000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:8      70.9 ns     80000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:16     86.8 ns    160000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:32     87.3 ns    320000000
bench_template_hash32_find_14_loop/iterations:10000000/threads:64     92.9 ns    640000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:1      38.4 ns     10000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:2      42.1 ns     20000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:4      46.5 ns     40000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:8      52.6 ns     80000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:16     60.0 ns    160000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:32     62.1 ns    320000000
bench_template_hash32_find_14_avx2/iterations:10000000/threads:64     65.8 ns    640000000

Answer 1

您完全可以在没有分支的情况下实现这一点，方法是比较数组的两个重叠部分，对它们进行位或运算，然后使用单个 lzcnt 获得最后一位的位置。此外，使用 vmovmskps 而不是 vpmovmskb 可以将结果除以 4（不过我不确定这是否会导致任何跨域延迟）。

int8_t hash32_find_14_avx2(uint32_t hash, volatile uint32_t* hashes) {
    uint32_t compacted_result_mask = 0;
    __m256i cmp_vector = _mm256_set1_epi32(hash);
    for(uint8_t base_index = 0; base_index < 12; base_index += 6) {
        __m256i ring_vector = _mm256_loadu_si256((__m256i*) (hashes + base_index));

        __m256i result_mask_vector = _mm256_cmpeq_epi32(ring_vector, cmp_vector);
        compacted_result_mask |= _mm256_movemask_ps(_mm256_castsi256_ps(result_mask_vector)) << (base_index);
    }
    int32_t leading_zeros = __lzcnt32(compacted_result_mask);
    return (31 - leading_zeros);
}

正如彼得在评论中已经指出的那样，在大多数情况下 _mm256_stream_load_si256 比正常负载更糟糕。另外，请注意，在 gcc 中使用未对齐加载时，您必须使用 -mno-avx256-split-unaligned-load 进行编译（或者实际上只使用 -march=native）- .

Godbolt-Link 带有简单的测试代码（请注意，如果数组中有多个匹配值，则循环版本和 avx2 版本的行为会有所不同）： https://godbolt.org/z/2jNWqK

使用 AVX2 查找元素索引 - 代码优化

Find element index with AVX2 - code optimization

c

linux

performance

simd

avx2