performance - 使用 SIMD/AVX/SSE 进行树遍历

Question

我目前正在研究是否可以加快 van Emde Boas(或任何树)树的遍历速度。给定一个搜索查询作为输入，在缓存行中已经有多个树节点(van emde Boas 布局)，树遍历似乎是指令瓶颈。

作为 SIMD/AVX/SSE 指令的新手，我想从该主题的专家那里知道是否可以一次将多个节点与一个值进行比较，然后找出要进一步遵循的树路径。我的研究导致了以下问题:

在 SIMD/AVX/SSE 寄存器等的构造上浪费了多少 cpu 周期/指令。如果构造比手动遍历整个子树花费更多的时间(2+4+8 个节点在1 个大小为 64 字节的缓存行)。

在找到正确的 SIMD/AVX/SSE 寄存器上浪费了多少 cpu 周期/指令来保存要遵循哪条路径的答案？任何人都可以想出一个聪明的方法，以便可以使用那些“findMinimumInteger” AVX 指令在 1 (??) cpu 周期内决定它？

你猜怎么着？

另一种加速树遍历的更棘手的方法是让多个搜索查询同时运行，当最后一个树级别中的节点很可能紧密相连时。对此有何猜测？ Ofc 它必须将那些不再属于同一子树的查询放在一边，然后在完成树的第一次“并行遍历”后递归找到它们。 树查询具有顺序但不是恒定的访问模式(查询 [i] 总是 < 比查询 [i+1])。

重要:这个东西是关于整数树的，这就是为什么使用 van Emde Boas 树(稍后可能会尝试 x-fast/y-fast)

我很好奇你在这个问题上的 50 美分是多少，因为人们可能对大型树的最高性能感兴趣。提前感谢您在这方面花费的时间:-)

Answer 1

我使用 SSE2/AVX2 来帮助执行 B+树搜索。这是在 AVX2 中对 16 个 DWORD 的完整缓存行执行二进制搜索的代码:

// perf-critical: ensure this is 64-byte aligned. (a full cache line)
union bnode
{
    int32_t i32[16];
    __m256i m256[2];
};

// returns from 0 (if value < i32[0]) to 16 (if value >= i32[15]) 
unsigned bsearch_avx2(bnode const* const node, __m256i const value)
{
    __m256i const perm_mask = _mm256_set_epi32(7, 6, 3, 2, 5, 4, 1, 0);

    // compare the two halves of the cache line.

    __m256i cmp1 = _mm256_load_si256(&node->m256[0]);
    __m256i cmp2 = _mm256_load_si256(&node->m256[1]);

    cmp1 = _mm256_cmpgt_epi32(cmp1, value); // PCMPGTD
    cmp2 = _mm256_cmpgt_epi32(cmp2, value); // PCMPGTD

    // merge the comparisons back together.
    //
    // a permute is required to get the pack results back into order
    // because AVX-256 introduced that unfortunate two-lane interleave.
    //
    // alternately, you could pre-process your data to remove the need
    // for the permute.

    __m256i cmp = _mm256_packs_epi32(cmp1, cmp2); // PACKSSDW
    cmp = _mm256_permutevar8x32_epi32(cmp, perm_mask); // PERMD

    // finally create a move mask and count trailing
    // zeroes to get an index to the next node.

    unsigned mask = _mm256_movemask_epi8(cmp); // PMOVMSKB
    return _tzcnt_u32(mask) / 2; // TZCNT
}

你最终会得到一个高度可预测的分支 bnode , 测试是否已到达树的末端。

这应该可以轻松扩展到 AVX-512。

预处理并摆脱缓慢 PERMD指令，这将用于:

void preprocess_avx2(bnode* const node)
{
    __m256i const perm_mask = _mm256_set_epi32(3, 2, 1, 0, 7, 6, 5, 4);
    __m256i *const middle = (__m256i*)&node->i32[4];

    __m256i x = _mm256_loadu_si256(middle);
    x = _mm256_permutevar8x32_epi32(x, perm_mask);
    _mm256_storeu_si256(middle, x);
}