c++ - 高效移位或大位 vector

Question

我有一个大的内存数组作为指针 uint64_t * arr (加上大小)，表示普通位。我需要非常有效地(最高性能/最快)将这些位从 0 向右移动到 63。

通过移动整个数组，我的意思是不移动每个元素(如 a[i] <<= Shift )，而是将其作为单个大位 vector 进行移动。换句话说，对于每个中间位置 i (第一个和最后一个元素除外)我可以循环执行以下操作:

dst[i] = w | (src[i] << Shift);
w = src[i] >> (64 - Shift);

在哪里 w是一些临时变量，保存前一个数组元素的右移值。

上面的这个解决方案简单明了。但我需要更高效的东西，因为我有千兆字节的数据。

理想情况下会为此使用一些 SIMD 指令，因此我正在寻找专家的 SIMD 建议。我需要为所有四种流行指令集 - SSE-SSE4.2/AVX/AVX-2/AVX-512 实现移位代码。

但据我所知，例如 SSE2 仅存在 _mm_slli_si128()内在/指令，仅移动 8 的倍数(换句话说，字节移动)。而且我需要按任意位大小进行移位，而不仅仅是字节移位。

如果没有 SIMD，我也可以通过使用 shld reg, reg, reg 一次移动 128 位指令，它允许进行 128 位移位。它被实现为内在的__shiftleft128()在 MSVC 中，并生成可以是 seen here 的汇编代码.

顺便说一句，我需要所有 MSVC/GCC/CLang 的解决方案。

同样在单循环迭代中，我可以在顺序操作中移动 4 或 8 个字，这将使用 CPU 流水线来加速多条指令的并行乱序执行。

如果需要，我的位 vector 可以与内存中的任意数量的字节对齐，如果这有助于例如通过对齐读/写来提高 SIMD 速度。源和目标位 vector 内存也不同(不重叠)。

换句话说，我正在寻找有关如何在不同的 Intel CPU 上最有效(最高效地)解决我的任务的所有建议。

注意，澄清一下，我实际上必须做几个类次，而不仅仅是一次类次。我有大位 vector X ，以及数百个类次大小s0, s1, ..., sN ，其中每个移位大小不同并且也可能很大(例如移位 100K 位)，然后我想计算得到的大位 vector Y = (X << s0) | (X << s1) | ... | (X << sN) .我只是将我对 StackOverflow 的问题简化为移动单个 vector 。但可能这个关于原始任务的细节非常重要。

应@Jake'Alquimista'LEE 的要求，我决定实现一个现成的玩具最小可重现示例，以计算输入位 vector 的移位或 src到产生或-ed最终dst位 vector 。这个例子根本没有优化，只是我的任务如何解决的一个简单的变体。为简单起见，这个例子的输入 vector 很小，不像我的例子那样是千兆字节。这是一个玩具示例，我没有检查它是否正确解决了任务，它可能包含一些小错误:

Try it online!

#include <cstdint>
#include <vector>
#include <random>

#define bit_sizeof(x) (sizeof(x) * 8)

using u64 = uint64_t;
using T = u64;

int main() {
    std::mt19937_64 rng{123};

    // Random generate source bit vector
    std::vector<T> src(100'000);
    for (size_t i = 0; i < src.size(); ++i)
        src[i] = rng();

    size_t const src_bitsize = src.size() * bit_sizeof(T);

    // Destination bit vector, for example twice bigger in size
    std::vector<T> dst(src.size() * 2);

    // Random generate shifts
    std::vector<u64> shifts(200);
    for (size_t i = 0; i < shifts.size(); ++i)
        shifts[i] = rng() % src_bitsize;

    // Right-shift that handles overflow
    auto Shr = [](auto x, size_t s) {
        return s >= bit_sizeof(x) ? 0 : (x >> s);
    };

    // Do actual Shift-Ors
    for (auto orig_shift: shifts) {
        size_t const
            word_off = orig_shift / bit_sizeof(T),
            bit_off = orig_shift % bit_sizeof(T);

        if (word_off >= dst.size())
            continue;
        
        size_t const
            lim = std::min(src.size(), dst.size() - word_off);

        T w = 0;
        
        for (size_t i = 0; i < lim; ++i) {
            dst[word_off + i] |= w | (src[i] << bit_off);
            w = Shr(src[i], bit_sizeof(T) - bit_off);
        }

        // Special case of handling for last word
        if (word_off + lim < dst.size())
            dst[word_off + lim] |= w;
    }
}

我真实项目的当前代码与上面的玩具示例不同。该项目已经正确解决了现实世界的任务。我只需要做额外的优化。我已经做了一些优化，比如使用 OpenMP在所有内核上并行化移位或操作。同样如评论中所说，我为每个类次大小创建了专门的模板函数，总共 64 个函数，并选择 64 个函数中的一个来执行实际的类次或。每个 C++ 函数都有移位大小的编译时间值，因此编译器会根据编译时间值进行额外的优化。

Answer 1

您可以，甚至可能不需要明确使用 SIMD 指令。目标编译器 GCC、CLANG 和 MSVC 以及 ICC 等其他编译器都支持自动矢量化。虽然手动优化的汇编可以胜过编译器生成的矢量化指令，但通常更难实现，并且您可能需要针对不同架构的多个版本。生成高效自动向量化指令的通用代码是一种可以跨多个平台移植的解决方案。

例如一个简单的 shiftvec 版本

void shiftvec(uint64_t* dst, uint64_t* src, int size, int shift)
{
    for (int i = 0; i < size; ++i,++src,++dst)
    {
        *dst = ((*src)<<shift) | (*(src+1)>>(64-shift));
    }
}

使用最近的 GCC 编译(或 CLANG 也可以)，-O3 -std=c++11 -mavx2 导致程序集核心循环中的 SIMD 指令

.L5:
  vmovdqu ymm4, YMMWORD PTR [rsi+rax]
  vmovdqu ymm5, YMMWORD PTR [rsi+8+rax]
  vpsllq ymm0, ymm4, xmm2
  vpsrlq ymm1, ymm5, xmm3
  vpor ymm0, ymm0, ymm1
  vmovdqu YMMWORD PTR [rdi+rax], ymm0
  add rax, 32
  cmp rax, rdx
  jne .L5

参见 godbolt.org:https://godbolt.org/z/5TxhqMhnK

这也概括了如果您想在 dst 中组合多个类次:

void shiftvec2(uint64_t* dst, uint64_t* src1, uint64_t* src2, int size1, int size2, int shift1, int shift2)
{
    int size = size1<size2 ? size1 : size2;
    for (int i = 0; i < size; ++i,++src1,++src2,++dst)
    {
        *dst = ((*src1)<<shift1) | (*(src1+1)>>(64-shift1));
        *dst |= ((*src2)<<shift2) | (*(src2+1)>>(64-shift2)); 
    }
    for (int i = size; i < size1; ++i,++src1,++dst)
    {
        *dst = ((*src1)<<shift1) | (*(src1+1)>>(64-shift1));        
    }
    for (int i = size; i < size2; ++i,++src2,++dst)
    {
        *dst = ((*src2)<<shift2) | (*(src2+1)>>(64-shift2));
    }
}

编译成核心循环:

.L38:
  vmovdqu ymm7, YMMWORD PTR [rsi+rcx]
  vpsllq ymm1, ymm7, xmm4
  vmovdqu ymm7, YMMWORD PTR [rsi+8+rcx]
  vpsrlq ymm0, ymm7, xmm6
  vpor ymm1, ymm1, ymm0
  vmovdqu YMMWORD PTR [rax+rcx], ymm1
  vmovdqu ymm7, YMMWORD PTR [rdx+rcx]
  vpsllq ymm0, ymm7, xmm3
  vmovdqu ymm7, YMMWORD PTR [rdx+8+rcx]
  vpsrlq ymm2, ymm7, xmm5
  vpor ymm0, ymm0, ymm2
  vpor ymm0, ymm0, ymm1
  vmovdqu YMMWORD PTR [rax+rcx], ymm0
  add rcx, 32
  cmp r10, rcx
  jne .L38

在一个循环中组合多个源将减少用于加载/写入目标的内存带宽总量。您可以组合的数量当然受到可用寄存器的限制。请注意，shiftvec 的 xmm2 和 xmm3 包含移位值，因此编译时已知移位值的不同版本可能会释放这些寄存器。

另外对每个指针使用 __restrict(受 GCC、CLANG、MSVC 支持)将告诉编译器范围不重叠。

我最初在 MSVC 提供适当的自动矢量化代码时遇到问题，但似乎添加更多类似 SIMD 的结构将使其适用于所有三个所需的编译器 GCC、CLANG 和 MSVC:

void shiftvec(uint64_t* __restrict dst, const uint64_t* __restrict src, int size, int shift)
{
    int i = 0;
    // MSVC: use steps of 2 for SSE, 4 for AVX2, 8 for AVX512
    for (; i+4 < size; i+=4,dst+=4,src+=4)
    {
        for (int j = 0; j < 4; ++j)
            *(dst+j) = (*(src+j))<<shift;
        for (int j = 0; j < 4; ++j)
            *(dst+j) |= (*(src+1)>>(64-shift));
    }
    for (; i < size; ++i,++src,++dst)
    {
        *dst = ((*src)<<shift) | (*(src+1)>>(64-shift));
    }    
}