是否可以对数组中的每 3 个相邻元素求和，并使用向量指令使它们中的每一个都等于总和？

Posted 2023-02-16

【中文标题】是否可以对数组中的每 3 个相邻元素求和，并使用向量指令使它们中的每一个都等于总和？

【英文标题】：Is it possible to sum every 3 neighbouring elements in an array and make each of them equal to the sum using vector instructions?

【发布时间】：2019-01-21 10:05:24

【问题描述】：

在我的程序中，我有一个 32 位整数的大数组。我必须对其进行以下操作：

sum = array[i] + array[i+1] + array[i+2]
array[i] = sum
array[i+1] = sum
array[i+2] = sum
i+=3

或者，正如我在汇编中写的那样：

loop: ;R12 - address of the array, R11 - last element, R10 - iterator

mov eax, [R12 + R10]
add eax, [R12 + R10 + 4]
add eax, [R12 + R10 + 8]

mov [R12 + R10], eax
mov [R12 + R10 + 4], eax
mov [R12 + R10 + 8], eax

mov rax, 0
mov rdx, 0

add R10, 12
cmp R10, R11
jb loop

是否可以使用向量指令来做到这一点？如果有，怎么做？

【问题讨论】：

使用 SSE 或 AVX2 都非常容易。 AVX 虽然只有浮点算术指令。我建议你使用内在函数而不是汇编（更快。更容易，更不容易出错，更便携）。

请注意a decent compiler will vectorize this for you，使用SSE 或AVX2。

@PaulR：对于整数，AVX1 增加了 3 操作数无损操作的效率，而不是例如vpalignr xmm0, xmm1, xmm2, 4 代替 movdqa xmm0, xmm1 / palignr xmm0, xmm2, 4 和未对齐的内存源操作数，而无需在 tmp 寄存器中使用单独的 movdqu。如果您遇到前端吞吐量的瓶颈，这两种方法都会很有帮助。

@matjag：你根本没有使用 RDX，mov rax, 0 完全没有意义，因为每次迭代的第一次加载 (mov eax, [r12 + r10]) 都会破坏对 RAX 旧值的任何依赖。

【参考方案1】：

编译器可以为您进行矢量化，但使用内在函数进行矢量化 可能会导致更有效的代码。函数sum3neighb 下面求和 3 个相邻 具有 12 个整数元素的数组的元素。它没有使用许多 shuffle，而是使用重叠加载来获得 数据在正确的位置。

/*  gcc -O3 -Wall -march=sandybridge -m64 neighb3.c                */
#include <stdio.h>
#include <immintrin.h>

inline __m128i _mm_shufps_epi32(__m128i a, __m128i b,int imm)
    return _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(a),_mm_castsi128_ps(b),imm));


/* For an integer array of 12 elements, sum every 3 neighbouring elements */
void sum3neighb(int * a)
    __m128i a_3210 = _mm_loadu_si128((__m128i*)&a[0]);
    __m128i a_9876 = _mm_loadu_si128((__m128i*)&a[6]);
    __m128i a_9630 = _mm_shufps_epi32(a_3210, a_9876, 0b11001100);

    __m128i a_4321 = _mm_loadu_si128((__m128i*)&a[1]);
    __m128i a_A987 = _mm_loadu_si128((__m128i*)&a[7]);
    __m128i a_A741 = _mm_shufps_epi32(a_4321, a_A987, 0b11001100);

    __m128i a_5432 = _mm_loadu_si128((__m128i*)&a[2]);
    __m128i a_BA98 = _mm_loadu_si128((__m128i*)&a[8]);
    __m128i a_B852 = _mm_shufps_epi32(a_5432, a_BA98, 0b11001100);

    __m128i sum = _mm_add_epi32(a_9630, a_A741);
            sum = _mm_add_epi32(sum, a_B852);    /* B+A+9, 8+7+6, 5+4+3, 2+1+0 */

    __m128i sum_3210 = _mm_shuffle_epi32(sum, 0b01000000);
    __m128i sum_7654 = _mm_shuffle_epi32(sum, 0b10100101);
    __m128i sum_BA98 = _mm_shuffle_epi32(sum, 0b11111110);

            _mm_storeu_si128((__m128i*)&a[0], sum_3210);
            _mm_storeu_si128((__m128i*)&a[4], sum_7654);
            _mm_storeu_si128((__m128i*)&a[8], sum_BA98);



int main()
    int i;
    int a[24];  
    for (i = 0; i < 24; i++) a[i] = i + 4;  /* example input */
    for (i = 0; i < 24; i++) printf("%3i  ",a[i]);
    printf("\n");
    for (i = 0; i < 24; i = i + 12)
       sum3neighb(&a[i]);
    
    for (i = 0; i < 24; i++) printf("%3i  ",a[i]);
    printf("\n");
    return 0;

这将编译为以下程序集(with gcc 8.2)：

sum3neighb:
  vmovups xmm4, XMMWORD PTR [rdi+4]
  vshufps xmm2, xmm4, XMMWORD PTR [rdi+28], 204
  vmovups xmm3, XMMWORD PTR [rdi]
  vshufps xmm0, xmm3, XMMWORD PTR [rdi+24], 204
  vpaddd xmm0, xmm0, xmm2
  vmovups xmm5, XMMWORD PTR [rdi+8]
  vshufps xmm1, xmm5, XMMWORD PTR [rdi+32], 204
  vpaddd xmm0, xmm0, xmm1
  vpshufd xmm2, xmm0, 64
  vpshufd xmm1, xmm0, 165
  vmovups XMMWORD PTR [rdi], xmm2
  vpshufd xmm0, xmm0, 254
  vmovups XMMWORD PTR [rdi+16], xmm1
  vmovups XMMWORD PTR [rdi+32], xmm0
  ret

示例程序的输出为：（第一行是输入，第二行是输出， 行被截断。）

 4    5    6    7    8    9   10   11   12   13   14   15   16   17   18   19   ...   
15   15   15   24   24   24   33   33   33   42   42   42   51   51   51   60   ...

---

clang 不接受 _mm_shufps_epi32 函数，请参阅 Peter 的评论。 有两种选择： 模板函数（参见 chtz 的评论，Godbolt link）

template<int imm>
inline  __m128i _mm_shufps_epi32(__m128i a, __m128i b)
    return _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(a),_mm_castsi128_ps(b),imm));

或者macro：

#define _mm_shufps_epi32(a,b,i) _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(a),_mm_castsi128_ps(b),i)) 

---

在较新的 Intel 架构（自 Haswell 起）上，整数向量加法指令比 shuffle 指令更快，请参阅 Agner Fog's instruction tables。在这种情况下，下面的代码可能会稍微高效一些。它需要2个加法， 但也少了 2 次洗牌：

void sum3neighb_v3(int * a)
    __m128i a_3210 = _mm_loadu_si128((__m128i*)&a[0]);
    __m128i a_4321 = _mm_loadu_si128((__m128i*)&a[1]);
    __m128i a_5432 = _mm_loadu_si128((__m128i*)&a[2]);
    __m128i sum53_20 = _mm_add_epi32(a_3210, a_5432);
    __m128i sum543_210 = _mm_add_epi32(sum53_20, a_4321);

    __m128i a_9876 = _mm_loadu_si128((__m128i*)&a[6]);
    __m128i a_A987 = _mm_loadu_si128((__m128i*)&a[7]);
    __m128i a_BA98 = _mm_loadu_si128((__m128i*)&a[8]);
    __m128i sumB9_86 = _mm_add_epi32(a_9876, a_BA98);
    __m128i sumBA9_876 = _mm_add_epi32(sumB9_86, a_A987
    );        
    __m128i sum = _mm_shufps_epi32(sum543_210, sumBA9_876, 0b11001100);

    __m128i sum_3210 = _mm_shuffle_epi32(sum, 0b01000000);
    __m128i sum_7654 = _mm_shuffle_epi32(sum, 0b10100101);
    __m128i sum_BA98 = _mm_shuffle_epi32(sum, 0b11111110);

            _mm_storeu_si128((__m128i*)&a[0], sum_3210);
            _mm_storeu_si128((__m128i*)&a[4], sum_7654);
            _mm_storeu_si128((__m128i*)&a[8], sum_BA98);

AVX2 版本

AVX2 版本，见下面的代码，使用车道交叉洗牌，因此不太适合 AMD 处理器，另见 chtz's answer。

void sum3neighb_avx2(int * a)
    __m256i a_0  = _mm256_loadu_si256((__m256i*)&a[0]);
    __m256i a_1  = _mm256_loadu_si256((__m256i*)&a[1]);
    __m256i a_2  = _mm256_loadu_si256((__m256i*)&a[2]);

    __m256i a_8  = _mm256_loadu_si256((__m256i*)&a[8]);
    __m256i a_9  = _mm256_loadu_si256((__m256i*)&a[9]);
    __m256i a_10 = _mm256_loadu_si256((__m256i*)&a[10]);

    __m256i a_16 = _mm256_loadu_si256((__m256i*)&a[16]);
    __m256i a_17 = _mm256_loadu_si256((__m256i*)&a[17]);
    __m256i a_18 = _mm256_loadu_si256((__m256i*)&a[18]);

    __m256i sum_0  = _mm256_add_epi32(_mm256_add_epi32(a_0,  a_1),  a_2);
    __m256i sum_8  = _mm256_add_epi32(_mm256_add_epi32(a_8,  a_9),  a_10);
    __m256i sum_16 = _mm256_add_epi32(_mm256_add_epi32(a_16, a_17), a_18);

    __m256i sum_8_0 = _mm256_blend_epi32(sum_0,  sum_8, 0b10010010);
    __m256i sum     = _mm256_blend_epi32(sum_8_0, sum_16, 0b00100100);

    __m256i sum_7_0   = _mm256_permutevar8x32_epi32(sum, _mm256_set_epi32(6,6,3,3,3,0,0,0));
    __m256i sum_15_8  = _mm256_permutevar8x32_epi32(sum, _mm256_set_epi32(7,4,4,4,1,1,1,6));
    __m256i sum_23_16 = _mm256_permutevar8x32_epi32(sum, _mm256_set_epi32(5,5,5,2,2,2,7,7));

            _mm256_storeu_si256((__m256i*)&a[0],  sum_7_0  );
            _mm256_storeu_si256((__m256i*)&a[8],  sum_15_8 );
            _mm256_storeu_si256((__m256i*)&a[16], sum_23_16);

【讨论】：

不幸的是，clang 不支持为shufps 编写包装函数，因为它在内联之前检查编译时常量，即使启用了优化也是如此。不过，您可以将其设为 #define 宏。无论如何，clang 的自动矢量化策略是使用 vpinsrd 手动收集以创建输出所需的矢量：/ 有趣的是与 gcc 不同。 godbolt.org/z/Gox1P3 它在您编写时或多或少地编译您的手动矢量化版本，因此它的 shuffle 优化器没有找到任何东西。

我想知道重叠存储是否是一个好主意，以节省向量之间的一些洗牌。 （即存储out0,out1,out2, x 的向量，然后将其重叠并重写x。）2x load + palignr 可以提供 2x paddd + movups。嗯，也许并不比你正在做的更好。如果针对非冗余输出格式（每个结果一次而不是重复 3 次），phaddd 可能很有用，但即便如此我也不确定 3 的倍数分组。

@PeterCordes 而不是宏，我通常更喜欢模板函数godbolt.org/z/pmqlEr。不过，这确实会导致语法不统一（除非您也包装所有其他 shuffle）。

@PeterCordes：值得注意的是，clang 和 gcc 都使用了 13 次 shuffle 操作（vpalignr、vpshufd 等）。可以在重叠加载后立即进行一些添加。这减少了 shuffle 的数量，但增加了 add 的数量，这对 Haswell 和更新的版本很好。在这种特殊情况下（重叠）广播商店会很好。显然它们不存在。

【参考方案2】：

如果有人正在寻找 AVX2 变体，这里有一个基于 https://***.com/a/45025712 的版本（它本身基于 an article by Intel）：

#include <immintrin.h> 

template<int imm>
inline __m256i _mm256_shufps_epi32(__m256i a, __m256i b)
    return _mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(a),_mm256_castsi256_ps(b),imm));


void sum3neighb24(int * a)
    __m256i a_FEDC_3210 = _mm256_insertf128_si256(_mm256_castsi128_si256(_mm_loadu_si128((__m128i*)&a[0])),_mm_loadu_si128((__m128i*)&a[12]),1) ;
    __m256i a_JIHG_7654 = _mm256_insertf128_si256(_mm256_castsi128_si256(_mm_loadu_si128((__m128i*)&a[4])),_mm_loadu_si128((__m128i*)&a[16]),1) ;
    __m256i a_NMLK_BA98 = _mm256_insertf128_si256(_mm256_castsi128_si256(_mm_loadu_si128((__m128i*)&a[8])),_mm_loadu_si128((__m128i*)&a[20]),1) ;


    __m256i a_MLJI_A976 = _mm256_shufps_epi32<_MM_SHUFFLE( 2,1, 3,2)>(a_JIHG_7654,a_NMLK_BA98);
    __m256i a_HGED_5421 = _mm256_shufps_epi32<_MM_SHUFFLE( 1,0, 2,1)>(a_FEDC_3210,a_JIHG_7654);

    __m256i a_LIFC_9630 = _mm256_shufps_epi32<_MM_SHUFFLE( 2,0, 3,0)>(a_FEDC_3210,a_MLJI_A976);
    __m256i a_MJGD_A741 = _mm256_shufps_epi32<_MM_SHUFFLE( 3,1, 2,0)>(a_HGED_5421,a_MLJI_A976);
    __m256i a_NKHE_B852 = _mm256_shufps_epi32<_MM_SHUFFLE( 3,0, 3,1)>(a_HGED_5421,a_NMLK_BA98);

    __m256i sum = _mm256_add_epi32(a_LIFC_9630, a_MJGD_A741);
            sum = _mm256_add_epi32(sum, a_NKHE_B852);    /* B+A+9, 8+7+6, 5+4+3, 2+1+0 */

    __m256i sum_FEDC_3210 = _mm256_shuffle_epi32(sum, 0b01000000);
    __m256i sum_JIHG_7654 = _mm256_shuffle_epi32(sum, 0b10100101);
    __m256i sum_NMLK_BA98 = _mm256_shuffle_epi32(sum, 0b11111110);

    _mm_storeu_si128((__m128i*)&a[0], _mm256_castsi256_si128(sum_FEDC_3210));
    _mm_storeu_si128((__m128i*)&a[4], _mm256_castsi256_si128(sum_JIHG_7654));
    _mm_storeu_si128((__m128i*)&a[8], _mm256_castsi256_si128(sum_NMLK_BA98));
    _mm_storeu_si128((__m128i*)&a[12], _mm256_extractf128_si256 (sum_FEDC_3210,1));
    _mm_storeu_si128((__m128i*)&a[16], _mm256_extractf128_si256 (sum_JIHG_7654,1));
    _mm_storeu_si128((__m128i*)&a[20], _mm256_extractf128_si256 (sum_NMLK_BA98,1));

反洗牌基于@wim 的回答。实际上，在开始时用更多的负载换取更少的随机播放可能会更好。

【讨论】：

你能提供一个我可以复制到godbolt.org 并编译成汇编的版本吗？我不熟悉内在函数。

只需在顶部添加一个#include <immintrin.h>（并确保使用-O2进行编译）：godbolt.org/z/V60PKI