利用ARM NEON intrinsic优化常用数学运算

ARM NEON是arm平台下的SIMD指令集，利用好这些指令可以使程序获得很大的速度提升。不过对很多人来说，直接利用汇编指令优化代码难度较大，这时就可以利用ARM NEON intrinsic指令，它是底层汇编指令的封装，不需要用户考虑底层寄存器的分配，但同时又可以达到原始汇编指令的性能。
所有的intrinsic指令可以参考博客《ARM Neon Intrinsics各函数介绍》，本文不再赘述。我们着重介绍两个利用ARM NEON intrinsic指令加速数学运算的例子。

1. 向量的点积

给定两个向量$\vec A=(a_1,a_2,a_3,...,a_n)$，$\vec B=(b_1,b_2,b_3,...,b_n)$，向量的点积为

$$\vec A \cdot\vec B=a_1b_1+a_2b_2+a_3b_3+...a_nb_n$$
可以看出上面的运算就是向量的每一维相乘然后相加，相乘之间具有良好的并行性，所以可以通过ARM NEON intrinsic指令进行加速。下面是代码实现：

float dot(float* A,float* B,int K)
{
    float sum=0;
    float32x4_t sum_vec=vdupq_n_f32(0),left_vec,right_vec;
    for(int k=0;k<K;k+=4)
    {
        left_vec=vld1q_f32(A+ k);
        right_vec=vld1q_f32(B+ k);
        sum_vec=vmlaq_f32(sum_vec,left_vec,right_vec);
    }

    float32x2_t r=vadd_f32(vget_high_f32(sum_vec),vget_low_f32(sum_vec));
    sum+=vget_lane_f32(vpadd_f32(r,r),0);

    return sum;
}

代码比较简单，核心代码就是先将两个数组每次4个存入ARM NEON intrinsic下的128位变量中，然后利用一个乘加指令计算4个乘积的累加和。最后将4个sum再相加就得到最终的结果。相比于串行代码，上面的代码有接近4倍的加速比。当数据类型是short或者char时，可以取得更高的加速比，下面以char举例：

int dot(char* A,char* B,int K)
{
    int sum=0;
    int16x8_t sum_vec=vdupq_n_s16(0);
    int8x8_t left_vec, right_vec;
    int32x4_t part_sum4;
    int32x2_t part_sum2;

    //有溢出的风险
    for(k=0;k<K;k+=8)
    {
        left_vec=vld1_s8(A+A_pos+k);
        right_vec=vld1_s8(B+B_pos+k);
        sum_vec=vmlal_s8(sum_vec,left_vec,right_vec);
    }

    part_sum4=vaddl_s16(vget_high_s16(sum_vec),vget_low_s16(sum_vec));   
    part_sum2=vadd_s32(vget_high_s32(part_sum4),vget_low_s32(part_sum4));
    sum+=vget_lane_s32(vpadd_s32(part_sum2,part_sum2),0);

    return sum;
}

基于char类型的点积代码和float类型的类似，不过由于char乘法存在溢出的可能性，所以相乘之后我们需要将数据类型升级成short。上面的代码也特别注释了一句：可能存在溢出，这是因为单个乘法不会溢出，但是乘法的结果相加可能会存在溢出。如果合理设计两个向量的值溢出的概率就会很低，更重要的一点是上面代码的加速比是float类型的2倍还要多，所以在速度要求非常严格的程序中，上面代码会带来非常明显的速度提升。

2. exp加速

现在的深度神经网络中，经常会使用到sigmoid函数或者softmax函数，而这些函数中都使用了浮点数的幂指数函数($e^x$)。常规的数学函数库$e^x$的精度高，但是速度慢，在博客《快速浮点数exp算法》中，作者给出一种快速算法，只需要两次乘法，速度相当快，缺点就是误差略大，普遍在0~5%之间。不过对神经网络来说，这样的精度已经足够了。
先看一下原始exp算法：

inline float fast_exp(float x)
{
    union {uint32_t i;float f;} v;
    v.i=(1<<23)*(1.4426950409*x+126.93490512f);

    return v.f;
}

算法的基本原理是考虑了float数据类型在内存中的布局而精巧设计的，想了解更多细节可以参考原博客，本文只介绍如何将其用ARM NEON intrinsic指令进行加速（相比原始博客，代码中第二个常量有点变化，该新常量是我试验出来的，误差更小）。
ARM NEON intrinsic指令的优势是并行计算，所以我们对一个数组的每一个元素进行exp并相加，然后将其加速：

inline float expf_sum(float* score,int len)
{
    float sum=0;
    float32x4_t sum_vec=vdupq_n_f32(0);
    float32x4_t ai=vdupq_n_f32(1064807160.56887296), bi;
    int32x4_t   int_vec;
    int value;
    for(int i=0;i<len;i+=4)
    {
        bi=vld1q_f32(score+4*i);
        sum_vec=vmlaq_n_f32(ai,bi,12102203.1616540672);
        int_vec=vcvtq_s32_f32(sum_vec);

        value=vgetq_lane_s32(int_vec,0);
        sum+=(*(float*)(&value));
        value=vgetq_lane_s32(int_vec,1);
        sum+=(*(float*)(&value));
        value=vgetq_lane_s32(int_vec,2);
        sum+=(*(float*)(&value));
        value=vgetq_lane_s32(int_vec,3);
        sum+=(*(float*)(&value));
    }

    return sum;
}

在原始算法中是先计算(1<<23)，然后将其和另外一部分相乘，我们将其简化成一个乘加操作：12102203.1616540672*x+1064807160.56887296。算法和点积很相似，先加载4个变量，然后执行乘加操作。之后的操作首先是将float类型的变量转成int型变量，之后再通过地址强转获取float值并累加。相比原始的exp累加，速度能有5、6倍左右的提升。
上述两个运算对搞深度学习的人来意味着什么不用我再做介绍了吧~