SIMD (avx) 处理如何工作?例如,如果我想要 10 个 32 位浮点数,如何放入 256 位 avx 向量?

che*_*z93 -1 c simd avx

我正在学习 C avx 内在函数,我想知道它是如何工作的。

我很熟悉我可以做这样的事情:

__m256 evens = _mm256_set_ps(2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0);
Run Code Online (Sandbox Code Playgroud)

这里我存储 8 个 32 位浮点数。所以这是 256 位。

但是假设我正在编写一个线性代数库。然后我如何处理任意数量的向量;例如,如何将 10 个 32 位浮点数放入 avx 向量中?

如果您能提供一些例子,我将非常感激

Hom*_*512 5

嗯,这是一个广泛的问题,有一个简单的答案:使用循环并处理不是所需向量宽度倍数的元素计数。

\n

由于过于宽泛的问题会带来过于宽泛的答案,而且我找不到关于 SO 的良好参考,也许值得写一下如何进行矢量化循环。

\n

标量变换循环

\n

作为一个例子,考虑一个简单的out[i] = in[i] * scale循环。当然,编译器完全有能力自行对其进行向量化,但它可以很好地为我们提供演示。基本实现如下所示:

\n
#include <immintrin.h>\n\n#include <stddef.h>\n/* using ptrdiff_t */\n\nvoid vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    /* broadcast factor to all vector elements */\n    const __m256 vscale = _mm256_set1_ps(scale);\n    ptrdiff_t i;\n    for(i = 0; i + 8 <= n; i += 8) {\n        /* process 8 elements at a time */\n        __m256 cur = _mm256_loadu_ps(in + i);\n        cur = _mm256_mul_ps(cur, vscale);\n        _mm256_storeu_ps(out + i, cur);\n    }\n    /* deal with last 0-7 elements */\n    for(; i < n; ++i)\n        out[i] = in[i] * scale;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

别名问题

\n

与纯 C 版本相比,向量化此循环有一个后果:输入和输出可能不会重叠,以致在一次迭代中存储向量会覆盖后续迭代之一中的输入。它们可能是彼此的直接别名 ( in == out),但您不能拥有in = out + 1或类似的别名。自动矢量化代码将包含对此的检查和后备,但如果您手动矢量化,您通常不会打扰并简单地将其定义为要求。

\n

如果您想阻止编译器进行这些检查和回退并将要求集成到函数声明中,请使用restrict指针

\n
void vector_scale(\n      float* restrict out, const float* in, ptrdiff_t n, float scale);\n
Run Code Online (Sandbox Code Playgroud)\n

或者,您可以在手动矢量化代码中执行相同的检查。

\n
inline _Bool pointers_interfere(\n      const void* out, const void* in, size_t vectorsize)\n{\n    ptrdiff_t distance = (const char*) out - (const char*) in;\n    return distance > 0 && (size_t) distance < vectorsize;\n}\nvoid vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m256 vscale = _mm256_set1_ps(scale);\n    ptrdiff_t i = 0;\n    if(pointers_interfere(out, in, sizeof(__m256))) {\n        for(; i < n; ++i) // unvectorized fallback\n            out[i] = in[i] * scale;\n        return;\n    }\n    ...\n}\n
Run Code Online (Sandbox Code Playgroud)\n

进行此检查时需要注意的一些事项:

\n
    \n
  1. 它允许直接别名out == in
  2. \n
  3. 该代码只是防止与标量代码的差异。它并不能使其正确。在这个例子中,如果out == in + 8 && n > 8,结果仍然是无意义的;这和你从纯 C 版本中得到的废话是一样的。
  4. \n
  5. out - in从技术上讲,如果两个指针都指向不同的数组,则 C 标准没有很好地定义指针算术。然而,我们不是为抽象机编程,而是专门为 x86-64 编程,我们可以假设它有一个平面内存模型。但在其他平台(例如 OpenCL 内核)上重用该代码时,请记住这一点,而该假设可能不成立。
  6. \n
\n

循环条件

\n

我选择了这种形式的循环for(i = 0; i + 8 <= n; i += 8),因为它易于阅读,不易混乱(尤其是无符号类型)并且相当高效。但这不是i+8有效的方法,因为当循环体本身只需要时,条件会强制编译器在循环体之前进行计算i。一个幼稚的循环结构会浪费一个保持寄存器i+8。为了理解编译器如何优化它以及我们如何帮助它,我们必须区分循环条件是有符号类型(例如 )ptrdiff_t和无符号类型(例如 )size_t

\n

对于有符号类型,编译器可以自由地假设i+8永远不会溢出或下溢,并且它也可以使用有符号比较。但是,编译器不会将条件重写为i < n - 7. 如果n = LLONG_MIN(或接近它)wheren-7会环绕,这将会失败。相反,GCC 将条件转换为类似i < ((n - 8) & -8) + 8. 如果循环不从 0 开始,事情会变得更加混乱,就像一些进一步优化的循环中的情况一样。因此我不建议手动进行此优化。

\n

对于无符号类型,可能的优化较少,并且如果 则发生回绕危险n < 8。编译器现在将 转入i循环i+8内部,并使用固定的负偏移量进行每次内存访问,例如[rsi-32+rax*4]代替[rsi+rax*4]; 或者它执行上面概述的简单代码,浪费额外的指令和寄存器。

\n

在签名和未签名的情况下,我们可以帮助编译器进行一些显式检查。像这样的事情应该可以解决问题:

\n
void vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m256 vscale = _mm256_set1_ps(scale);\n    ptrdiff_t i = 0;\n    if(n > 7) {\n        for(; i < n - 7; i += 8) {\n            __m256 cur = _mm256_loadu_ps(in + i);\n            cur = _mm256_mul_ps(cur, vscale);\n            _mm256_storeu_ps(out + i, cur);\n        }\n    }\n    for(; i < n; ++i)\n        out[i] = in[i] * scale;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

如果您像我在这里一样使用有符号循环计数器,那么您可能会得到更好的代码,如果您只是if(n <= 0) return;在函数的早期编写,然后您可以执行一个简单的i < n - 7条件而不必担心环绕。然而,为了方便复制和粘贴,我确保即使循环计数器更改为无符号类型,下面的所有代码仍然有效。

\n

结盟

\n

对于 SSE ,内存访问与 16 字节边界对齐曾经非常重要。AVX 通常可以很好地处理未对齐的访问,但您仍然可能会发现对齐访问的性能有所提高。对于AVX512来说,对齐再次变得更加重要。对齐循环是否值得额外的代码和运行时需要进行测试。我在这里展示如何做到这一点。由于我们有两个数组,因此我们只能真正对齐到一个。如有疑问,通常最好对齐输出,但这也可能取决于具体情况和硬件。

\n

计算对齐位置是微优化的主要候选者。我在下面展示的例程相当高效,并且编写为可读和可重用。C++ 模板或直接内联代码可能可以节省一些指令。想法总是相同的:将内存地址舍入到向量对齐的下一个更高倍数。然后计算出有多少标量条目,并检查下一个对齐是否超出数组末尾。

\n
#include <stdint.h>\n/* using uintptr_t */\n\ninline void* next_aligned(const void* ptr, size_t alignment, const void* end)\n{\n    const uintptr_t pos = (uintptr_t) ptr;\n    const uintptr_t alignedpos = (pos + alignment - 1) & -alignment;\n    void* rtrn = (void*) alignedpos;\n    return rtrn > end ? (void*) end : rtrn;\n}\n\nvoid vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m256 vscale = _mm256_set1_ps(scale);\n    const ptrdiff_t aligned =\n          (float*) next_aligned(out, sizeof(__m256), out + n) - out;\n    ptrdiff_t i;\n    for(i = 0; i < aligned; ++i)\n        out[i] = in[i] * scale;\n    if(n > 7) {\n        for(; i < n - 7; i += 8) {\n            __m256 cur = _mm256_loadu_ps(in + i);\n            cur = _mm256_mul_ps(cur, vscale);\n            _mm256_store_ps(out + i, cur);\n        }\n    }\n    for(; i < n; ++i)\n        out[i] = in[i] * scale;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

向量化头部和尾部

\n

我们可以使用较小的向量大小来处理未对齐的前面和剩余的尾部元素,而不是标量循环。这样做的另一个好处是阻止编译器自动向量化这些短循环。这是否值得需要测试。例如Oracle Java的编译器似乎认为这是不值得的

\n
void vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m256 vscale = _mm256_set1_ps(scale);\n    const ptrdiff_t aligned =\n          (float*) next_aligned(out, sizeof(__m256), out + n) - out;\n    ptrdiff_t i, tail;\n    if(n <= 0) /* otherwise bit-tests fail if negative */\n        return;\n    if((i = aligned & 1) != 0)\n        out[0] = in[0] * scale;\n    if(aligned & 2) {\n        /* process 2 elements */\n        __m128 cur = _mm_loadl_pi(_mm_undefined_ps(), (const __m64*) (in + i));\n        cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n        _mm_storel_pi((__m64*) (out + i), cur);\n        i += 2;\n    }\n    if(aligned & 4) {\n        /* process 4 elements */\n        __m128 cur = _mm_loadu_ps(in + i);\n        cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n        _mm_store_ps(out + i, cur);\n        i += 4;\n    }\n    if(n > 7) {\n        for(; i < n - 7; i += 8) {\n            /* process 8 elements at a time */\n            __m256 cur = _mm256_loadu_ps(in + i);\n            cur = _mm256_mul_ps(cur, vscale);\n            _mm256_store_ps(out + i, cur);\n        }\n    }\n    tail = n - i;\n    if(tail & 4) {\n        __m128 cur = _mm_load_ps(in + i);\n        cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n        _mm_store_ps(out + i, cur);\n        i += 4;\n    }\n    if(tail & 2) {\n        __m128 cur = _mm_loadl_pi(_mm_undefined_ps(), (const __m64*) (in + i));\n        cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n        _mm_storel_pi((__m64*) (out + i), cur);\n        i += 2;\n    }\n    if(tail & 1) {\n        out[i] = in[i] * scale;\n        ++i;\n    }\n}\n
Run Code Online (Sandbox Code Playgroud)\n

头尾重叠

\n

我们可以应用一个简单的技巧,而不是使用较小的向量:我们对前 8 个条目进行一次未对齐的迭代。然后,我们不再向前跳过 8 个元素,而是向前跳到对齐的位置。这意味着第一次对齐迭代将与未对齐迭代部分重叠,但这比我们上面进行的部分处理要便宜得多。然后我们可以对尾部执行相同的操作,从 开始进行未对齐的迭代n - 8。与上述循环相比,这种处理方式多了三个要求。

\n
    \n
  1. 我们至少需要一个充满数据的完整向量。我们必须包括小向量的特殊情况
  2. \n
  3. 向量寄存器内的处理需要与位置无关。这意味着对于每个标量,in[i]向量寄存器内的位置一定无关紧要(例如,这可能发生在随机指令中)
  4. \n
  5. 输入和输出不能互相别名,否则您将对重叠元素应用相同的操作两次
  6. \n
\n
void vector_scale(\n      float* restrict out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m256 vscale = _mm256_set1_ps(scale);\n    const ptrdiff_t aligned =\n          (float*) next_aligned(out + 1, sizeof(__m256), out + n) - out;\n    ptrdiff_t i, tail = n - 8;\n    if(n < 8) {\n        /* special case for small inputs */\n        if(n <= 0)\n            return;\n        if((i = n & 1) != 0)\n            out[0] = in[0] * scale;\n        if(n & 2) {\n            __m128 cur = _mm_loadl_pi(_mm_undefined_ps(), (const __m64*) (in + i));\n            cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n            _mm_storel_pi((__m64*) (out + i), cur);\n            i += 2;\n        }\n        if(n & 4) {\n            __m128 cur = _mm_loadu_ps(in + i);\n            cur = _mm_mul_ps(cur, _mm256_castps256_ps128(vscale));\n            _mm_storeu_ps(out + i, cur);\n            i += 4;\n        }\n        return;\n    }\n    {  /* first iteration, potentially misaligned */\n        __m256 cur = _mm256_loadu_ps(in);\n        cur = _mm256_mul_ps(cur, vscale);\n        _mm256_storeu_ps(out, cur);\n    }\n    for(i = aligned; i < tail; i += 8) {\n        __m256 cur = _mm256_loadu_ps(in + i);\n        cur = _mm256_mul_ps(cur, vscale);\n        _mm256_store_ps(out + i, cur);\n    }\n    {   /* last iteration */\n        __m256 cur = _mm256_loadu_ps(in + tail);\n        cur = _mm256_mul_ps(cur, vscale);\n        _mm256_storeu_ps(out + tail, cur);\n    }\n}\n
Run Code Online (Sandbox Code Playgroud)\n

next_aligned(out, \xe2\x80\xa6)请注意从到 的细微变化next_aligned(out + 1, \xe2\x80\xa6)。否则,正确对齐的输出数组将导致一次完全冗余的迭代。

\n

减少循环

\n

上面的示例涵盖了简单的转换,但是缩减呢?让我们探索一个简单的向量点积:sum(left[i] * right[i])。如上所述,我们从简单开始,然后迭代优化代码。我假设我们有可用于此的 AVX 和 FMA 指令。

\n
float vector_dot(const float* left, const float* right, ptrdiff_t n)\n{\n    __m256 vsum = _mm256_setzero_ps();\n    __m128 high_half, low_half;\n    float headsum, tailsum = 0.f;\n    ptrdiff_t i = 0;\n    if(n > 7) {\n        for(; i < n - 7; i += 8) {\n            const __m256 lefti = _mm256_loadu_ps(left + i);\n            const __m256 righti = _mm256_loadu_ps(right + i);\n            vsum = _mm256_fmadd_ps(lefti, righti, vsum);\n        }\n    }\n    /* Reduce 8 elements to 4 */\n    high_half = _mm256_extractf128_ps(vsum, 1);\n    low_half = _mm256_castps256_ps128(vsum);\n    low_half = _mm_add_ps(low_half, high_half);\n    /* Reduce 4 to 2 */\n    high_half = _mm_movehl_ps(low_half, low_half);\n    low_half = _mm_add_ps(low_half, high_half);\n    /* Reduce 2 to 1 */\n    high_half = _mm_movehdup_ps(low_half);\n    low_half = _mm_add_ss(low_half, high_half);\n    headsum = _mm_cvtss_f32(low_half);\n    /* Process last 0-7 elements */\n    for(; i < n; ++i)\n        tailsum += left[i] * right[i];\n    return headsum + tailsum;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

循环携带的依赖链

\n

我们的归约主循环有一个关键问题:FMA 指令取决于前一迭代的值。在大多数系统上,FMA 的延迟为 4 个周期(尽管 AMD Zen2 或 Intel Broadwell 等较旧的 CPU 使用 5 个周期)。这意味着,我们最多每 4 个周期迭代一次,而不是每个周期迭代一次。为了避免这种情况,我们必须交错至少 4 个 FMA 依赖链;基本上保留四部分金额。让这个额外痛苦的是你不能依赖内部循环来实现这一点,除非你知道你的编译器会展开这些循环(GCC 不会)。

\n
float vector_dot(const float* left, const float* right, ptrdiff_t n)\n{\n    __m256 vsum1 = _mm256_setzero_ps();\n    __m256 vsum2 = vsum1, vsum3 = vsum1, vsum4 = vsum1;\n    __m128 high_half, low_half;\n    float headsum, tailsum = 0.f;\n    ptrdiff_t i = 0;\n    if(n > 31) {\n        for(; i < n - 31; i += 32) {\n            __m256 lefti = _mm256_loadu_ps(left + i);\n            __m256 righti = _mm256_loadu_ps(right + i);\n            vsum1 = _mm256_fmadd_ps(lefti, righti, vsum1);\n            lefti = _mm256_loadu_ps(left + i + 8);\n            righti = _mm256_loadu_ps(right + i + 8);\n            vsum2 = _mm256_fmadd_ps(lefti, righti, vsum2);\n            lefti = _mm256_loadu_ps(left + i + 16);\n            righti = _mm256_loadu_ps(right + i + 16);\n            vsum3 = _mm256_fmadd_ps(lefti, righti, vsum3);\n            lefti = _mm256_loadu_ps(left + i + 24);\n            righti = _mm256_loadu_ps(right + i + 24);\n            vsum4 = _mm256_fmadd_ps(lefti, righti, vsum4);\n        }\n    }\n    if(n - i >= 16) {\n        __m256 lefti = _mm256_loadu_ps(left + i);\n        __m256 righti = _mm256_loadu_ps(right + i);\n        vsum1 = _mm256_fmadd_ps(lefti, righti, vsum1);\n        lefti = _mm256_loadu_ps(left + i + 8);\n        righti = _mm256_loadu_ps(right + i + 8);\n        vsum2 = _mm256_fmadd_ps(lefti, righti, vsum2);\n        i += 16;\n    }\n    if(n - i >= 8) {\n        __m256 lefti = _mm256_loadu_ps(left + i);\n        __m256 righti = _mm256_loadu_ps(right + i);\n        vsum1 = _mm256_fmadd_ps(lefti, righti, vsum1);\n        i += 8;\n    }\n    /* Reduce 32 to 16 */\n    vsum1 = _mm256_add_ps(vsum1, vsum3);\n    vsum2 = _mm256_add_ps(vsum2, vsum4);\n    /* Reduce 16 to 8 */\n    vsum1 = _mm256_add_ps(vsum1, vsum2);\n    /* Reduce 8 elements to 4 */\n    high_half = _mm256_extractf128_ps(vsum1, 1);\n    low_half = _mm256_castps256_ps128(vsum1);\n    low_half = _mm_add_ps(low_half, high_half);\n    /* Reduce 4 to 2 */\n    high_half = _mm_movehl_ps(low_half, low_half);\n    low_half = _mm_add_ps(low_half, high_half);\n    /* Reduce 2 to 1 */\n    high_half = _mm_movehdup_ps(low_half);\n    low_half = _mm_add_ss(low_half, high_half);\n    headsum = _mm_cvtss_f32(low_half);\n    /* Process last 0-7 elements */\n    for(; i < n; ++i)\n        tailsum += left[i] * right[i];\n    return headsum + tailsum;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

进一步优化

\n

以下优化是可能的,但不会详细介绍:

\n
    \n
  1. 第一次迭代不需要完整的乘加,因为我们只加零。我们可以剥离第一次迭代并将其变成一个简单的乘法。然而,该步骤对输入数据的长度特别敏感。仅当向量很短时它才真正相关。但是,您必须平衡用于初始化部分总和的分支数量与这些分支的成本(以及潜在的错误预测)。可能有必要为短向量创建专门的代码路径。
  2. \n
  3. 最后 0-7 个元素可以向量化。可能是这样的:
  4. \n
\n
    ptrdiff_t tail;\n    if(n <= 0) /* check so bit tests don\'t fail with negative n */\n        return 0.f;\n    ... /* main loop */\n    tail = n - i;\n    /* Reduce 8 elements to 4 */\n    high_half = _mm256_extractf128_ps(vsum1, 1);\n    low_half = _mm256_castps256_ps128(vsum1);\n    low_half = _mm_add_ps(low_half, high_half);\n    if(tail & 4) {\n        __m128 lefti = _mm_loadu_ps(left + i);\n        __m128 righti = _mm_loadu_ps(right + i);\n        low_half = _mm_fmadd_ps(lefti, righti, low_half);\n        i += 4;\n    }\n    if(tail & 2) {\n        __m128 lefti = _mm_loadl_pi(_mm_setzero_ps(), (const __m64*) (left + i));\n        __m128 righti = _mm_loadl_pi(_mm_setzero_ps(), (const __m64*) (right + i));\n        low_half = _mm_fmadd_ps(lefti, righti, low_half);\n        i += 2;\n    }\n    /* Reduce 2 to 1 */\n    high_half = _mm_movehdup_ps(low_half);\n    low_half = _mm_add_ss(low_half, high_half);\n    float sum = _mm_cvtss_f32(low_half);\n    if(tail & 1)\n        sum += left[i] * right[i];\n    return sum;\n
Run Code Online (Sandbox Code Playgroud)\n
    \n
  1. 我们可以对齐输入向量之一。选择哪个向量很难预测。使用 SSE 指令,检查生成的代码以查看适当的对齐方式是否允许使用短movaps指令或将对齐的负载折叠到算术指令中可能是有意义的。
  2. \n
\n

AVX-512

\n

有了 AVX-512 的掩码寄存器,事情就变得容易多了。对齐和尾循环可以通过屏蔽加载和存储简单地实现。

\n
void vector_scale(float* out, const float* in, ptrdiff_t n, float scale)\n{\n    const __m512 vscale = _mm512_set1_ps(scale);\n    const ptrdiff_t aligned =\n          (float*) next_aligned(out + 1, sizeof(__m512), out + n) - out;\n    __mmask16 mask = _cvtu32_mask16((1 << (unsigned) aligned) - 1);\n    ptrdiff_t i = aligned, tail;\n    __m512 cur;\n    if(n <= 0)\n        return;\n    cur = _mm512_maskz_loadu_ps(mask, in);\n    cur = _mm512_mul_ps(cur, vscale);\n    _mm512_mask_storeu_ps(out, mask, cur);\n    if(n > 15) {\n        for(; i < n - 15; i += 16) {\n            cur = _mm512_load_ps(in + i);\n            cur = _mm512_mul_ps(cur, vscale);\n            _mm512_store_ps(out + i, cur);\n        }\n    }\n    tail = n - i;\n    mask = _cvtu32_mask16((1 << (unsigned) tail) - 1);\n    cur = _mm512_maskz_load_ps(mask, in + i);\n    cur = _mm512_mul_ps(cur, vscale);\n    _mm512_mask_store_ps(out + i, mask, cur);\n}\n
Run Code Online (Sandbox Code Playgroud)\n

改编此代码时需要注意的一件事是掩码生成:如果您需要__mmask32,则中间体1 << ones需要是 64 位类型,以避免溢出(ones == 32如果__mmask64您需要 )__int128。特别是在第二种情况下,更好的选择是__mask64 mask = _cvtu64_mask64((uint64_t) -1 >> (64 - ones)).

\n

与 AVX-512 一样,选择 512 位或 256 位向量需要对特定硬件进行测试和了解,因为大向量可能会也可能不会降低时钟速率,从而可能影响代码的其他部分,并且可能有其他限制因素

\n