使用Intel SSE/AVX指令集（SIMD）加速向量内积计算

最近在写遗传感知机。感知机其实就是神经网络的结构模型，某一层中的某一个节点需要对前面一层中每个节点的输出进行加权求和，这其实就是一个向量内积运算。即使撇开什么感知机、什么神经网络不说，单单是向量内积操作本身就是一个很常用但又很耗时的计算任务。因此，我把向量内积运算封装成一个函数，并且通过Intel SSE/AVX指令集进行了性能优化。优化后的性能是本来的4倍！

假设有两个向量：

float a[5]={1,2,3,4,5};
float b[5]={5,4,3,2,1};

首先来看最最原始的求向量内积的函数：

#include <stdint.h>

float vector_dot(uint32_t len,const float* a,const float* b)
{
    double sum=0;
    for(uint32_t i=0;i<len;i++)
        sum+=a[i]*b[i];
    return (float)sum;
}

在这种实现中，各个分量的相乘是串行的。这种串行方式叫做SISD，也就是Single Instruction Single Data，即一条指令处理一条数据。那好，我要介绍啥叫SIMD了！SIMD是Single Instruction Multiply Data的缩写，也就是说，一条指令同时处理多条数据。试想一下，如果能够把5个分量各自的乘法运算同时进行，最后在把各个分量的乘积相加，是不是能快一些？Intel SSE指令集给予我们这样的能力，即可以把四组单精度浮点数同时运算。

我们先来看一个简单点的例子：

test.c

#include <x86intrin.h>
#include <stdio.h>

int main()
{    
    float __attribute__((aligned(16))) a[4]={1,3,5,7};
    float __attribute__((aligned(16))) b[4]={5,10,15,20};
    __m128 A=_mm_load_ps(a);
    __m128 B=_mm_load_ps(b);
    __m128 C=_mm_add_ps(A,B);
    for(int i=0;i<4;i++)
        printf("%f ",C[i]);
    printf("\n");
    return 0;
}

编译命令中需要加入-msse4.1参数（或者-msse4也可以）：

gcc -msse4.1 test.c -o test

运行之后，输出结果：

6.000000 13.000000 20.000000 27.000000

就是四个分量各自相加的结果。代码中

__attribute__((aligned(16)))

是GCC的扩展属性，就是告诉GCC，这个数组要16字节对齐（起始地址要能够被16整除）。这个是Intel的规定，即__m128这个数据结构指向的地址必须是128-bit aligned的，也就是16字节对齐。

计算向量内积（也叫点乘）也有专门的SIMD指令，直接看用SSE指令集改进后的代码：

#include <stdint.h>
#include <x86intrin.h>

float vector_dot(uint32_t len,const float* a,const float* b)
{
    double sum=0;
    uint32_t i=0;
    uint32_t end=len&(~3);
    for(;i<end;i+=4)
    {
        __m128 A=_mm_load_ps(a+i);
        __m128 B=_mm_load_ps(b+i);
        __m128 C=_mm_dp_ps(A,B,0xf1);
        sum+=C[0];
    }
    for(;i<len;i++)
        sum+=a[i]*b[i];
    return (float)sum;
}

是不是发现就是在本来的代码上加了一点东西？！代码中的第一个for循环中，把4个分量作为一个LOOP地处理，通过_mm_dp_ps来计算四对float的点乘，然后把所有的点乘结果累计。由于len可能不能被4整除，所以最后再用本来的办法把剩余的几个分量累加。这里需要说明的是_mm_dp_ps的第三个参数mask。第三个参数mask的高4位用来指明哪些分量参与计算。比如我只想计算第0,2,3这3个分量的内积，那么高4位就是1101。由于我要4个分量都参与计算，所以就是1111。低4位表面计算结果放到哪些位置。由于计算结果也是一个__m128类型，即也是4个float，那么就必须说明结果存放在哪个float中。我想放在第0个float，那么就是0001，所以最终mask就是11110001,即0xf1。

后来，Intel又推出了AVX指令集，功能相似，但是并行度更加高了，支持8个float同时运算。改进后的代码是这样的：

#include <stdint.h>
#include <x86intrin.h>

float vector_dot(uint32_t len,const float* a,const float* b)
{
    double sum=0;
    uint32_t i=0;
    uint32_t end;
    end=len&(~7);
    for(;i<end;i+=8)
    {
        __m256 tmp=_mm256_dp_ps(_mm256_load_ps(a+i),_mm256_load_ps(b+i),0xf1);
        sum+=tmp[0]+tmp[4];
    }
    end=len&(~3);
    for(;i<end;i+=4)
    {
        __m128 C=_mm_dp_ps(_mm_load_ps(a+i),_mm_load_ps(b+i),0xf1);
        sum+=C[0];
    }
    for(;i<len;i++)
        sum+=a[i]*b[i];
    return (float)sum;
}

编译命令是：

gcc -mavx test.c -o test

当AVX被使能时，SSE自动使能。

为了以后使用方便，我加入了条件编译：

#include <stdint.h>
#include <x86intrin.h>

#define PTR_ADDR(p) ((unsigned long)p)

#if defined(__AVX__)
    #define ALIGNMENT 32
#elif defined(__SSE4_1__)
    #define ALIGNMENT 16
#else
    #define ALIGNMENT 1
#endif

float vector_dot(uint32_t len,const float* a,const float* b)
{
    assert(PTR_ADDR(a)%ALIGNMENT==0&&PTR_ADDR(b)%ALIGNMENT==0);
    double sum=0;
    uint32_t i=0;
    uint32_t end;
    #ifdef __AVX__
        end=len&(~7);
        for(;i<end;i+=8)
        {
            __m256 tmp=_mm256_dp_ps(_mm256_load_ps(a+i),_mm256_load_ps(b+i),0xf1);
            sum+=tmp[0]+tmp[4];
        }
    #endif
    #ifdef __SSE4_1__
        end=len&(~3);
        for(;i<end;i+=4)
        {
            __m128 tmp=_mm_dp_ps(_mm_load_ps(a+i),_mm_load_ps(b+i),0xf1);
            sum+=tmp[0];
        }
    #endif
    for(;i<len;i++)
        sum+=a[i]*b[i];
    return (float)sum;
}

编译的时候可以加-mavx参数或者-msse4.1参数。注意一点，-mavx自动包含-msse4.1。另外，传入的向量的起始地址必须是16字节（SSE）或者32字节（AVX）对齐的，我加入了断言。否则运行时会出现Segment Fault。对于静态分配的数组，可以使用GCC扩展属性__attribute__((aligned(...)))指定，而对于动态分配内存，我将在后续博客中讲解。