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

Question

在我的程序中,我有一个很大的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

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

Answer 1

编译器可以为您进行矢量化,但使用内在函数进行矢量化可能会导致更高效的代码.sum3neighb下面的函数对具有12个整数元素的数组的3个相邻元素求和.它使用重叠载荷来获取正确位置的数据,而不是使用许多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;
}

这将编译为以下程序集(使用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链接)

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));
}

或宏:

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

在较新的英特尔架构(自Haswell)以来,整数向量加法指令比随机指令更快,请参阅 Agner Fog的指令表.在这种情况下,以下代码可能稍微更有效.它需要额外增加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版本,见下面的代码,使用了车道交叉shuffle,因此不太适合AMD处理器,请参阅chtz的答案.

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);
}