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

左手的ㄟ右手 2022-05-27 08:13 552阅读 0赞

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

## 1. 向量的点积 ##

给定两个向量A⃗ =(a1,a2,a3,...,an)  A → = ( a 1 , a 2 , a 3 , . . . , a n ) ，B⃗ =(b1,b2,b3,...,bn)  B → = ( b 1 , b 2 , b 3 , . . . , b n ) ，向量的点积为

A⃗ ⋅B⃗ =a1b1\+a2b2\+a3b3\+...anbn  A → ⋅ B → = a 1 b 1 + a 2 b 2 + a 3 b 3 + . . . a n b 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函数，而这些函数中都使用了浮点数的幂指数函数(ex  e x )。常规的数学函数库ex  e x 的精度高，但是速度慢，在博客《[快速浮点数exp算法][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型变量，之后再通过[地址强转][Link 1]获取float值并累加。相比原始的exp累加，速度能有5、6倍左右的提升。  
上述两个运算对搞深度学习的人来意味着什么不用我再做介绍了吧~

[ARM Neon Intrinsics]: https://blog.csdn.net/hemmingway/article/details/44828303
[exp]: https://blog.csdn.net/shakingWaves/article/details/78450500?locationNum=6&fps=1
[Link 1]: https://blog.csdn.net/yutianzuijin/article/details/40265995