如何初始化范围从 0 到 N 的 SIMD 向量?

How do initialize an SIMD vector with a range from 0 to N?

我有以下功能,我正在尝试为其编写 AXV 版本:

void
hashids_shuffle(char *str, size_t str_length, char *salt, size_t salt_length)
{
    size_t i, j, v, p;
    char temp;

    if (!salt_length) {
        return;
    }

    for (i = str_length - 1, v = 0, p = 0; i > 0; --i, ++v) {
        v %= salt_length;
        p += salt[v];
        j = (salt[v] + v + p) % i;

        temp = str[i];
        str[i] = str[j];
        str[j] = temp;
    }
}

我正在尝试矢量化 v %= salt_length;
我想初始化一个包含从 0 到 str_length 的数字的向量,以便使用 SVML 的 _mm_rem_epu64 来计算每个循环迭代的 v。
如何正确初始化向量?

AVX512在余数运算中可以使用的最小数字是8位,内在是:_mm512_rem_epu8

但是,要用 0 到 N 的值初始化它的输入,您需要使用 _mm512_set_epi32 并将 8 位整数传递给 32 位整数,因为似乎没有内部取 64 个整数,每个 8 位。代码如下所示:

const __m512i sseConst = _mm512_set_epi32(
    (63<<24) | (62<<16) | (61<<8) | 60,
    (59<<24) | (58<<16) | (57<<8) | 56,
    ... etc ...);

如果您不喜欢这样的打字或害怕打字错误,请考虑为此编写一个代码生成器。

如果您不使用 malloc() 分配 __m512i 类型,则应自动为您处理对齐。否则,请在 "aligned malloc" 中搜索您的编译器。您需要的对齐方式是 64 字节(等于 512 位)。

当你需要向量中的下一个 N 个整数时,你可以这样做:

const __m512i inc = _mm512_set1_epi32((N<<24) | (N<<16) | (N<<8) | N);

然后您可以添加 incsseConst(内在 _mm512_add_epi32)。

[免责声明:考虑 32 位应用程序 - 其中 size_t 是无符号整数]

对齐分配可以通过使用 aligned_malloc 函数来完成,您可以在 windows and linux.

中找到这两个函数

以这种方式分配您的字符串(到 64 字节边界)将允许您通过对所有字节使用对齐加载 _mm256_load_si256 将数据直接加载到 _mm256i 寄存器中。

如果字符串没有正确对齐,您可以使用未对齐加载 _mm256_loadu_si256 来加载第一个字节。

您执行的第一个模运算 (v %= salt_length) 是使用常量操作数完成的,您可以使用 _mm256_set1_epi32 在 avx 寄存器中的循环之前对其进行初始化:

__m256i mod = _mm256_set2_epi32(salt_length);

对于下一个,您可以使用 _mm256_set_epi32,它使用提供的所有值初始化一个寄存器(注意 逆序 )。

请注意,如果您使用的是 AVX2 或 AVX512(而不仅仅是 AVX——您的问题对指令集有点混淆),您还可以使用 gather instructions 加载数据,从在第二个参数中指定的索引处的向量。

询问如何初始化向量基本上是在询问教程。 Google 上传一些 Intel 的内部函数使用指南。我会假装这个问题不是微不足道的,回答有关尝试有效地实现此功能的问题。这绝对不是您作为初学者尝试矢量化的那种函数。

请参阅 tag wiki for links to docs, esp. Intel's intrinsics guide. See the tag wiki for a link to a very nice intro to SIMD programming with SSE intrinsics, and how to use SIMD effectively,以及其他链接。

目录摘要:

  • Unrolling/vectorizing 免费照顾 v % salt_length
  • 如果 v++; v %= loop_invariant; 不是 2 的幂或 compile-time 常量,您如何对其进行矢量化。包括有关使用 _mm_set_epi8 或为此目的初始化向量的其他方法的标题问题的答案。
  • 如何像这样开始矢量化复杂循环:从展开以查找串行依赖项开始。
  • 完整函数的未经测试版本,它向量化除 % i 和交换之外的所有内容。 (即矢量化所有便宜的操作,就像你问的那样)。

    • (v + salt[v] + p)(以及导致它的所有内容)向量化为两个 vpaddw 指令。用于矢量化 p 的循环外的 prefix-sum 设置很棘手,但我最终也对其进行了矢量化。

    • 函数的绝大部分 运行 时间将在 j 元素向量上的标量内循环中,在 div 上成为瓶颈(或其他SVML 可以做到),and/or 缓存未命中非常大的字符串。


整个循环无法轻松矢量化,因为与 pseudo-random 索引的交换会创建一个不可预测的 table 序列依赖性。使用 AVX512 gather + shuffle + scatter 和 AVX512CD 来查找冲突位掩码可能是可能的,但这必须是一个单独的问题。我不确定有效地执行此操作有多困难,或者您是否经常最终多次重复向量随机播放,只在一个 non-conflicting 元素中取得进展。


因为 salt_length = sizeof(size_t) 是一个 compile-time 常数并且是比向量长度小的 2 的幂,所以 v++v%=salt_length 不需要循环内的任何代码都,并且作为有效展开循环以并行执行多个v值的side-effect免费发生。

(使用 platform-dependent salt 大小意味着 32 位构建将无法处理使用 64 位 salt 创建的数据。即使 x32 ABI 也有 32 位 size_t ,因此更改为 uint64_t 似乎很有意义,除非您永远不需要在机器之间共享加盐哈希。)

在标量循环中,v 遵循 0 1 2 3 0 1 2 3 ...(或 0..7,对于 64 位)的重复模式。在向量代码中,我们可能一次使用 32 字节向量中的 4B 个元素执行 8 v 个值,或者使用 2B 个元素一次执行 16 次迭代。

所以v变成了loop-invariant常数向量。有趣的是,salt[v] 也是如此,所以我们永远不需要在循环内进行任何 salt table 查找。 事实上,v+salt[v] 可以是 pre-computed 标量和向量。

标量版本也应该 pre-compute v+salt[v] 并展开 4 或 8,删除 LUT 查找,以便所有 memory/cache 吞吐量可用于实际交换。编译器可能不会为您执行此操作,因此您可能需要手动展开并编写额外的代码来处理最后的奇数个字符串字节。 (如果不展开,您仍然可以 pre-compute 查找 v+salt[v] 的 table,类型足够宽,不会环绕)。

即使只是确保 salt_length 在编译时是已知的也将允许更好的代码。 v %= compile_time_constantdiv insn 便宜,而且当它是 2 的幂时非常便宜。(它只是变成 v &= 7)。如果标量版本可以内联,或者如果您使用 salt_length = sizeof(size_t) 而不是将其作为函数参数传递,编译器可能会为您执行此操作。


如果您还不知道 salt_length: 即在您透露有关 salt_length:[=71 的关键信息之前 @harold 的建议=]

因为我们知道 v < salt_length 开始,所以我们最多只需要一个 v -= salt_length 来将它包装回正确的范围并保持不变。这称为 "strength reduction" 运算,因为减法是比除法更弱(且成本更低)的运算。

// The scalar loop would benefit from this transformation, too.
// or better, unroll the scalar loop by 8 so everything becomes constant
v++;
if( v >= salt_length)
    v-= salt_length;

向量化只是:让我们假设我们所知道的是salt_length <= 16,所以我们可以使用32个uint8_t值的向量。 (我们可以使用 pshufb 对 salt[v] LUT 查找进行矢量化)。

// untested  // Vectorizing  v++; v %= unknown_loop_invariant_value;

if (!salt_length) return;
assert(salt_length <= 16);  // so we can use pshufb for the salt[v] step

__m256i vvec = _mm256_setr_epi8(  // setr: lowest element first, unlike set
   0%salt_length,  1%salt_length,  2%salt_length,  3%salt_length, 
   4%salt_length,  5%salt_length,  6%salt_length,  7%salt_length,
   8%salt_length,  9%salt_length, 10%salt_length, 11%salt_length,
  12%salt_length, 13%salt_length, 14%salt_length, 15%salt_length,
  16%salt_length, 17%salt_length, 18%salt_length, 19%salt_length,
  20%salt_length, 21%salt_length, 22%salt_length, 23%salt_length,
  24%salt_length, 25%salt_length, 26%salt_length, 27%salt_length,
  28%salt_length, 29%salt_length, 30%salt_length, 31%salt_length);
__m256i v_increment = _mm256_set1_epi8(32 % salt_length);
__m256i vmodulus    = _mm256_set1_epi8(salt_length);

// salt_lut = _mm256_set1_epi64x(salt_byval);  // for known salt length. (pass it by val in a size_t arg, instead of by char*).

// duplicate the salt into both lanes of a vector.  Garbage beyond salt_length isn't looked at.
__m256i salt_lut = _mm256_broadcastsi128_si256(_mm_loadu_si128(salt));  // nevermind that this could segfault if salt is short and at the end of a page.

//__m256i v_plus_salt_lut = _mm256_add_epi8(vvec, salt_lut); // not safe with 8-bit elements: could wrap
// We could use 16-bit elements and AVX512 vpermw (or vpermi2w to support longer salts)

for (...) {
    vvec = _mm256_add_epi8(vvec, v_increment);         // ++v;

    // if(!(salt_length > v)) { v-= salt_length; }
    __m256i notlessequal = _mm256_cmpgt_epi8(vmodulus, vvec);  // all-ones where salt_length > v.
    //  all-zero where salt_length <= v, where we need to subtract salt_length
    __m256i conditional_sub = _mm256_and_si256(notlessequal, vmodulus)
    vvec = _mm256_sub_epi8(vvec, conditional_sub);   // subtract 0 or salt_length

    // salt[v] lookup:
    __m256i saltv = _mm256_shuffle_epi8(salt_lut, vvec);  // salt[v]

   // then maybe pmovzx and vextracti128+pmovzx to zero-extend to 16-bit elements?    Maybe vvec should only be a 16-bit vector?
   // or unpack lo/hi with zeros (but that behaves differently from pmovzx at the lane boundary)
   // or  have vvec already holding 16-bit elements with the upper half of each one always zero.  mask after the pshufb to re-zero,
   //   or do something clever with `vvec`, `v_increment` and `vmodulus` so `vvec` can have `0xff` in the odd bytes, so pshufb zeros those elements.
}

当然,如果我们知道 salt_length 是 2 的幂,我们就应该屏蔽掉每个元素中除了相关低位以外的所有位:

vvec = _mm256_add_epi8(vvec, _mm256_set1_epi8(salt_length));       // ++v;
vvec = _mm256_and_si256(vvec, _mm256_set1_epi8(salt_length - 1));  // v &= salt_length - 1; // aka v%=salt_length;

注意到我们从错误的元素大小开始是因为我们意识到一次只向量化一行是一个坏主意,因为现在我们必须去更改我们已经编写的所有代码以使用更宽的元素,有时需要不同的策略来做同样的事情。

当然,你确实需要从一个粗略的大纲开始,在脑海中或写下来,说明你如何分别执行每个步骤。在思考这个问题的过程中,您会看到不同部分如何组合在一起。

对于复杂循环,有用的第一步可能是尝试手动展开标量循环。这将有助于找到串行依赖关系,以及通过展开简化的事情。


(stuff) % i:困难的部分

我们需要足够宽的元素来容纳i的最大值,因为i不是2的幂,而且不是常量,所以该模运算需要工作。任何更宽的宽度都是一种浪费,会降低我们的吞吐量。如果我们可以向量化整个循环的其余部分,甚至可能值得针对 str_length 的不同范围专门使用不同版本的函数。 (或者可能用 64b 元素循环直到 i<= UINT32_MAX,然后循环直到 i<=UINT16_MAX,等等)。如果您知道不需要处理大于 4GiB 的字符串,则可以仅使用 32 位数学来加速常见情况。 (64 位除法比 32 位除法慢,即使高位全为零)。

实际上,我认为 我们需要与最大值 p 一样宽的元素,因为它会永远累积(直到它在原始标量中以 2^64 换行代码)。与常数模数不同,我们不能只使用 p%=i 来控制它,即使模数是分配的。 (123 % 33) % (33-16) != 123 % (33-16)。即使对齐到 16 也无济于事:12345 % 32 != 12345 % 48 % 32

这将很快使 p 对于 i 的重复条件减法(直到条件掩码全部为假)来说太大,即使对于相当大的 i 值也是如此。

有已知整数常量取模的技巧(参见 http://libdivide.com/),但 AFAIK 计算附近除数的乘法模逆(即使 power-of-two 步幅为 16)不是'这比完全独立的号码更容易。所以我们不能廉价地调整下一个 i 值向量的常量。

小数定律或许可以让worth-while剥掉最后一对 向量迭代,具有 pre-computed 个乘法模逆向量 所以 % i 可以用向量来完成。一旦我们接近字符串的末尾,它可能在 L1 缓存中很热,所以我们在 div 而不是交换 loads/stores 上完全成为瓶颈。为此,我们可能会使用序言来达到 16 的倍数的 i 值,因此当我们接近 i=0 时,最后几个向量始终具有与 i 值相同的对齐方式。否则我们将有一个 i 值范围的常量 LUT,并简单地从中进行未对齐的加载。这意味着我们不必旋转 salt_vp.

可能转换为 FP 会很有用,因为最近的 Intel CPU(尤其是 Skylake)具有非常强大的 FP 分区硬件和显着的流水线(吞吐量:延迟比)。如果我们可以通过正确的舍入选择获得准确的结果,那就太好了。 (floatdouble 可以准确地表示任何整数,最大约为其尾数的大小。)

我猜值得尝试英特尔的 _mm_rem_epu16(使用 i 值的矢量,您使用 set1(16) 的矢量递减)。如果他们使用 FP 来获得准确的结果,那就太好了。如果它只是解压缩为标量并进行整数除法,那么将值返回到向量中会浪费时间。

无论如何,最简单的解决方案当然是使用标量循环迭代向量元素。直到你想出一些非常奇特的使用 AVX512CD 进行交换的东西,这似乎是合理的,但它可能比交换慢一个数量级,如果它们在 L1 缓存中都是热的。


(未测试)partially-vectorized 函数版本:

这是 the code on the Godbolt compiler explorer, with full design-notes comments, including the diagrams I made while figuring out the SIMD prefix-sum algo. I eventually remembered I'd seen a narrower version of this as a building block in @ZBoson's floating point SSE Prefix sum answer,但大部分 re-inventing 之后我自己才开始。

// See the godbolt link for full design-notes comments
// comments about what makes nice ASM or not.
#include <stdint.h>
#include <stdlib.h>
#include <immintrin.h>
#include <assert.h>

static inline
__m256i init_p_and_increment(size_t salt_length, __m256i *p_increment, __m256i saltv_u16, __m128i saltv_u8)
{  // return initial p vector (for first 16 i values).
   // return increment vector by reference.

  if (salt_length == 4) {
    assert(0); // unimplemented
    // should be about the same as length == 8.  Can maybe factor out some common parts, like up to psum2
  } else {
    assert(salt_length == 8);

    // SIMD prefix sum for n elements in a vector in O(log2(n)) steps.
    __m128i sv     = _mm256_castsi256_si128(saltv_u16);
    __m128i pshift1 = _mm_bslli_si128(sv, 2);        // 1 elem (uint16_t)
    __m128i psum1   = _mm_add_epi16(pshift1, sv);
    __m128i pshift2 = _mm_bslli_si128(psum1, 4);     // 2 elem
    __m128i psum2   = _mm_add_epi16(pshift2, psum1);
    __m128i pshift3 = _mm_bslli_si128(psum2, 8);     // 4 elem
    __m128i psum3   = _mm_add_epi16(pshift3, psum2); // p_initial low 128.  2^3 = 8 elements = salt_length
    // psum3 = the repeating pattern of p values.  Later values just add sum(salt[0..7]) to every element
     __m128i p_init_low = psum3;

    __m128i sum8_low = _mm_sad_epu8(saltv_u8, _mm_setzero_si128());  // sum(s0..s7) in each 64-bit half
    // alternative:
    //        sum8_low = _mm_bsrli_si128(p_init_low, 14); // has to wait for psum3 to be ready: lower ILP than doing psadbw separately
    __m256i sum8 = _mm256_broadcastw_epi16(sum8_low);

    *p_increment = _mm256_slli_epi16(sum8, 1);   // set1_epi16(2*sum(salt[0..7]))

    __m128i p_init_high = _mm_add_epi16(p_init_low, _mm256_castsi256_si128(sum8));
    __m256i p_init = _mm256_castsi128_si256(p_init_low);
    p_init = _mm256_inserti128_si256(p_init, p_init_high, 1);
      // not supported by gcc _mm256_set_m128i(p_init_high, psum3);

    return p_init;
  }

}

void
hashids_shuffle_simd(char *restrict str, size_t str_length, size_t salt_byval)
{
    //assert(salt_length <= 16); // so we can use pshufb for the salt[v] step for non-constant salt length.

    // platform-dependent salt size seems weird. Why not uint64_t?
    size_t salt_length = sizeof(size_t);

    assert(str_length-1 < UINT16_MAX);   // we do p + v + salt[v] in 16-bit elements
    // TODO: assert((str_length-1)/salt_length * p_increment < UINT16_MAX);

    __m128i saltv_u8;
    __m256i v, saltv;
    if(salt_length == 4) {
          v = _mm256_set1_epi64x(0x0003000200010000);   // `v%salt_length` is 0 1 2 3 repeating
      saltv_u8 = _mm_set1_epi32( salt_byval );
      saltv = _mm256_cvtepu8_epi16( saltv_u8 );         // salt[v] repeats with the same pattern: expand it to 16b elements with pmovzx
    } else {
        assert(salt_length == 8);
            v = _mm256_cvtepu8_epi16( _mm_set1_epi64x(0x0706050403020100) );
        saltv_u8 = _mm_set1_epi64x( salt_byval );
        saltv = _mm256_cvtepu8_epi16( saltv_u8 );
    }

    __m256i v_saltv = _mm256_add_epi16(v, saltv);

    __m256i p_increment;
    __m256i p = init_p_and_increment(salt_length, &p_increment, saltv, saltv_u8);


    for (unsigned i=str_length-1; i>0 ; /*i-=16 */){
        // 16 uint16_t j values per iteration.  i-- happens inside the scalar shuffle loop.
        p = _mm256_add_epi16(p, p_increment);    // p += salt[v]; with serial dependencies accounted for, prefix-sum style

        __m256i j_unmodded = _mm256_add_epi16(v_saltv, p);

        // size_t j = (v + saltv[v] + p) % i;
        //////// scalar loop over 16 j elements, doing the modulo and swap
        // alignas(32) uint16_t jbuf[16];   // portable C++11 syntax
        uint16_t jbuf[16] __attribute__((aligned(32)));  // GNU C syntax
        _mm256_store_si256((__m256i*)jbuf, j_unmodded);

        const int jcount = sizeof(jbuf)/sizeof(jbuf[0]);
        for (int elem = 0 ; elem < jcount ; elem++) {
          if (--i == 0) break;  // in fact returns from the whole function.

              // 32-bit division is significantly faster than 64-bit division
          unsigned j = jbuf[elem] % (uint32_t)i;
          // doubtful that vectorizing this with Intel SVML _mm_rem_epu16 would be a win
          // since there's no hardware support for it.  Until AVX512CD, we need each element in a gp reg as an array index anyway.

          char temp = str[i];
          str[i] = str[j];
          str[j] = temp;
        }

    }
}

这编译成 asm 看起来是正确的,但我没有 运行 它。

Clang 制作了一个相当合理的内部循环。这是 -fno-unroll-loops 的可读性。将其排除在性能之外,尽管在这里无关紧要,因为循环开销不是瓶颈。

 # The loop part of clang3.8.1's output.  -march=haswell -fno-unroll-loops (only for human readability.  Normally it unrolls by 2).
.LBB0_6:  # outer loop                  #   in Loop: Header=BB0_3 Depth=1
    add     esi, 1
.LBB0_3:  # first iteration entry point # =>This Loop Header: Depth=1
    vpaddw  ymm2, ymm2, ymm1           # p += p_increment
    vpaddw  ymm3, ymm0, ymm2           # v+salt[v] + p
    vmovdqa ymmword ptr [rsp], ymm3    # store jbuf
    add     esi, -1
    lea     r8, [rdi + rsi]
    mov     ecx, 1
.LBB0_4:  # inner loop                  #   Parent Loop BB0_3 Depth=1
    # gcc's version fully unrolls the inner loop, leading to code bloat
    test    esi, esi                            # if(i==0) return
    je      .LBB0_8
    movzx   eax, word ptr [rsp + 2*rcx - 2]     # load jbuf
    xor     edx, edx
    div     esi
    mov     r9b, byte ptr [r8]                  # swap
    mov     al, byte ptr [rdi + rdx]
    mov     byte ptr [r8], al
    mov     byte ptr [rdi + rdx], r9b
    add     esi, -1
    add     r8, -1
    cmp     rcx, 16                     # silly clang, not macro-fusing cmp/jl because it wants to use a weird way to increment.
    lea     rcx, [rcx + 1]
    jl      .LBB0_4                     # inner loop
    jmp     .LBB0_6                     # outer loop